fix: Ensure tailnet coordinations are sent orderly (coder#4198)

kylecarbs · web-flow · commit b8ec5c786ddb · 2022-09-26T10:16:04.000-05:00
diff --git a/codersdk/agentconn.go b/codersdk/agentconn.go
@@ -46,7 +46,7 @@ type AgentConn struct {
 func (c *AgentConn) Ping() (time.Duration, error) {
 	errCh := make(chan error, 1)
 	durCh := make(chan time.Duration, 1)
-	c.Conn.Ping(TailnetIP, tailcfg.PingICMP, func(pr *ipnstate.PingResult) {
+	c.Conn.Ping(TailnetIP, tailcfg.PingDisco, func(pr *ipnstate.PingResult) {
 		if pr.Err != "" {
 			errCh <- xerrors.New(pr.Err)
 			return
diff --git a/tailnet/conn.go b/tailnet/conn.go
@@ -182,6 +182,7 @@ func NewConn(options *Options) (*Conn, error) {
 		magicConn:        magicConn,
 		dialer:           dialer,
 		listeners:        map[listenKey]*listener{},
+		peerMap:          map[tailcfg.NodeID]*tailcfg.Node{},
 		tunDevice:        tunDevice,
 		netMap:           netMap,
 		netStack:         netStack,
@@ -192,10 +193,17 @@ func NewConn(options *Options) (*Conn, error) {
 		wireguardEngine: wireguardEngine,
 	}
 	wireguardEngine.SetStatusCallback(func(s *wgengine.Status, err error) {
+		server.logger.Info(context.Background(), "wireguard status", slog.F("status", s), slog.F("err", err))
 		if err != nil {
 			return
 		}
 		server.lastMutex.Lock()
+		if s.AsOf.Before(server.lastStatus) {
+			// Don't process outdated status!
+			server.lastMutex.Unlock()
+			return
+		}
+		server.lastStatus = s.AsOf
 		server.lastEndpoints = make([]string, 0, len(s.LocalAddrs))
 		for _, addr := range s.LocalAddrs {
 			server.lastEndpoints = append(server.lastEndpoints, addr.Addr.String())
@@ -240,6 +248,7 @@ type Conn struct {
 
 	dialer             *tsdial.Dialer
 	tunDevice          *tstun.Wrapper
+	peerMap            map[tailcfg.NodeID]*tailcfg.Node
 	netMap             *netmap.NetworkMap
 	netStack           *netstack.Impl
 	magicConn          *magicsock.Conn
@@ -254,6 +263,7 @@ type Conn struct {
 	nodeChanged bool
 	// It's only possible to store these values via status functions,
 	// so the values must be stored for retrieval later on.
+	lastStatus        time.Time
 	lastEndpoints     []string
 	lastPreferredDERP int
 	lastDERPLatency   map[string]float64
@@ -282,8 +292,9 @@ func (c *Conn) SetNodeCallback(callback func(node *Node)) {
 func (c *Conn) SetDERPMap(derpMap *tailcfg.DERPMap) {
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
-	c.netMap.DERPMap = derpMap
 	c.logger.Debug(context.Background(), "updating derp map", slog.F("derp_map", derpMap))
+	c.netMap.DERPMap = derpMap
+	c.wireguardEngine.SetNetworkMap(c.netMap)
 	c.wireguardEngine.SetDERPMap(derpMap)
 }
 
@@ -292,18 +303,24 @@ func (c *Conn) SetDERPMap(derpMap *tailcfg.DERPMap) {
 func (c *Conn) UpdateNodes(nodes []*Node) error {
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
-	peerMap := map[tailcfg.NodeID]*tailcfg.Node{}
 	status := c.Status()
 	for _, peer := range c.netMap.Peers {
-		if peerStatus, ok := status.Peer[peer.Key]; ok {
-			// Clear out inactive connections!
-			// If a connection hasn't been active for a minute post creation, we assume it's dead.
-			if !peerStatus.Active && peer.Created.Before(time.Now().Add(-time.Minute)) {
-				c.logger.Debug(context.Background(), "clearing peer", slog.F("peerStatus", peerStatus))
-				continue
-			}
+		peerStatus, ok := status.Peer[peer.Key]
+		if !ok {
+			continue
 		}
-		peerMap[peer.ID] = peer
+		// If this peer was added in the last 5 minutes, assume it
+		// could still be active.
+		if time.Since(peer.Created) < 5*time.Minute {
+			continue
+		}
+		// We double-check that it's safe to remove by ensuring no
+		// handshake has been sent in the past 5 minutes as well. Connections that
+		// are actively exchanging IP traffic will handshake every 2 minutes.
+		if time.Since(peerStatus.LastHandshake) < 5*time.Minute {
+			continue
+		}
+		delete(c.peerMap, peer.ID)
 	}
 	for _, node := range nodes {
 		peerStatus, ok := status.Peer[node.Key]
@@ -322,18 +339,11 @@ func (c *Conn) UpdateNodes(nodes []*Node) error {
 			// reason. TODO: @kylecarbs debug this!
 			KeepAlive: ok && peerStatus.Active,
 		}
-		existingNode, ok := peerMap[node.ID]
-		if ok {
-			peerNode.Created = existingNode.Created
-			c.logger.Debug(context.Background(), "updating peer", slog.F("peer", peerNode))
-		} else {
-			c.logger.Debug(context.Background(), "adding peer", slog.F("peer", peerNode))
-		}
-		peerMap[node.ID] = peerNode
+		c.peerMap[node.ID] = peerNode
 	}
-	c.netMap.Peers = make([]*tailcfg.Node, 0, len(peerMap))
-	for _, peer := range peerMap {
-		c.netMap.Peers = append(c.netMap.Peers, peer)
+	c.netMap.Peers = make([]*tailcfg.Node, 0, len(c.peerMap))
+	for _, peer := range c.peerMap {
+		c.netMap.Peers = append(c.netMap.Peers, peer.Clone())
 	}
 	netMapCopy := *c.netMap
 	c.wireguardEngine.SetNetworkMap(&netMapCopy)
@@ -425,6 +435,7 @@ func (c *Conn) sendNode() {
 	}
 	c.nodeSending = true
 	go func() {
+		c.logger.Info(context.Background(), "sending node", slog.F("node", node))
 		nodeCallback(node)
 		c.lastMutex.Lock()
 		c.nodeSending = false
diff --git a/tailnet/coordinator.go b/tailnet/coordinator.go
@@ -164,6 +164,7 @@ func (c *Coordinator) ServeClient(conn net.Conn, id uuid.UUID, agent uuid.UUID)
 			c.mutex.Unlock()
 			continue
 		}
+		c.mutex.Unlock()
 		// Write the new node from this client to the actively
 		// connected agent.
 		data, err := json.Marshal([]*Node{&node})
@@ -173,14 +174,11 @@ func (c *Coordinator) ServeClient(conn net.Conn, id uuid.UUID, agent uuid.UUID)
 		}
 		_, err = agentSocket.Write(data)
 		if errors.Is(err, io.EOF) {
-			c.mutex.Unlock()
 			return nil
 		}
 		if err != nil {
-			c.mutex.Unlock()
 			return xerrors.Errorf("write json: %w", err)
 		}
-		c.mutex.Unlock()
 	}
 }
 
@@ -259,7 +257,7 @@ func (c *Coordinator) ServeAgent(conn net.Conn, id uuid.UUID) error {
 				wg.Done()
 			}()
 		}
-		wg.Wait()
 		c.mutex.Unlock()
+		wg.Wait()
 	}
 }

Original file line number	Diff line number	Diff line change
`@@ -164,6 +164,7 @@ func (c *Coordinator) ServeClient(conn net.Conn, id uuid.UUID, agent uuid.UUID)`
`164`	`164`	`c.mutex.Unlock()`
`165`	`165`	`continue`
`166`	`166`	`}`
	`167`	`+ c.mutex.Unlock()`
`167`	`168`	`// Write the new node from this client to the actively`
`168`	`169`	`// connected agent.`
`169`	`170`	`data, err := json.Marshal([]*Node{&node})`
`@@ -173,14 +174,11 @@ func (c *Coordinator) ServeClient(conn net.Conn, id uuid.UUID, agent uuid.UUID)`
`173`	`174`	`}`
`174`	`175`	`_, err = agentSocket.Write(data)`
`175`	`176`	`if errors.Is(err, io.EOF) {`
`176`		`- c.mutex.Unlock()`
`177`	`177`	`return nil`
`178`	`178`	`}`
`179`	`179`	`if err != nil {`
`180`		`- c.mutex.Unlock()`
`181`	`180`	`return xerrors.Errorf("write json: %w", err)`
`182`	`181`	`}`
`183`		`- c.mutex.Unlock()`
`184`	`182`	`}`
`185`	`183`	`}`
`186`	`184`
`@@ -259,7 +257,7 @@ func (c *Coordinator) ServeAgent(conn net.Conn, id uuid.UUID) error {`
`259`	`257`	`wg.Done()`
`260`	`258`	`}()`
`261`	`259`	`}`
`262`		`- wg.Wait()`
`263`	`260`	`c.mutex.Unlock()`
	`261`	`+ wg.Wait()`
`264`	`262`	`}`
`265`	`263`	`}`