fix(wsconncache): only allow one peer per connection

coadler · coadler · commit 269bc1e15049 · 2023-01-26T16:07:54.000-06:00
If an agent went away and reconnected, the wsconncache connection would be polluted for about 10m because there would be two peers with the same IP. The old peer always had priority, which caused the dashboard to try and always dial the old peer until it was removed. Fixes: #5292
diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -278,18 +278,21 @@ func (api *API) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
 	_, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageBinary)
 	defer wsNetConn.Close() // Also closes conn.
 
+	fmt.Println("getting conn from cache")
 	agentConn, release, err := api.workspaceAgentCache.Acquire(r, workspaceAgent.ID)
 	if err != nil {
 		_ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("dial workspace agent: %s", err))
 		return
 	}
 	defer release()
+	fmt.Println("got conn from cache")
 	ptNetConn, err := agentConn.ReconnectingPTY(ctx, reconnect, uint16(height), uint16(width), r.URL.Query().Get("command"))
 	if err != nil {
 		_ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("dial: %s", err))
 		return
 	}
 	defer ptNetConn.Close()
+	fmt.Println("dialed reconnecting pty, doing bicopy")
 	agent.Bicopy(ctx, wsNetConn, ptNetConn)
 }
 
@@ -402,11 +405,11 @@ func (api *API) workspaceAgentListeningPorts(rw http.ResponseWriter, r *http.Req
 
 func (api *API) dialWorkspaceAgentTailnet(r *http.Request, agentID uuid.UUID) (*codersdk.AgentConn, error) {
 	clientConn, serverConn := net.Pipe()
-	go func() {
-		<-r.Context().Done()
-		_ = clientConn.Close()
-		_ = serverConn.Close()
-	}()
+	// go func() {
+	// 	<-r.Context().Done()
+	// 	_ = clientConn.Close()
+	// 	_ = serverConn.Close()
+	// }()
 
 	derpMap := api.DERPMap.Clone()
 	for _, region := range derpMap.Regions {
@@ -453,7 +456,16 @@ func (api *API) dialWorkspaceAgentTailnet(r *http.Request, agentID uuid.UUID) (*
 	}
 
 	sendNodes, _ := tailnet.ServeCoordinator(clientConn, func(node []*tailnet.Node) error {
-		return conn.UpdateNodes(node)
+		err := conn.RemoveAllPeers()
+		if err != nil {
+			return xerrors.Errorf("remove all peers: %w", err)
+		}
+
+		err = conn.UpdateNodes(node)
+		if err != nil {
+			return xerrors.Errorf("update nodes: %w", err)
+		}
+		return nil
 	})
 	conn.SetNodeCallback(sendNodes)
 	go func() {
@@ -465,6 +477,10 @@ func (api *API) dialWorkspaceAgentTailnet(r *http.Request, agentID uuid.UUID) (*
 	}()
 	return &codersdk.AgentConn{
 		Conn: conn,
+		CloseFunc: func() {
+			_ = clientConn.Close()
+			_ = serverConn.Close()
+		},
 	}, nil
 }
 
diff --git a/tailnet/conn.go b/tailnet/conn.go
@@ -325,6 +325,31 @@ func (c *Conn) SetDERPMap(derpMap *tailcfg.DERPMap) {
 	c.wireguardEngine.SetDERPMap(derpMap)
 }
 
+func (c *Conn) RemoveAllPeers() error {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
+	c.netMap.Peers = []*tailcfg.Node{}
+	c.peerMap = map[tailcfg.NodeID]*tailcfg.Node{}
+	netMapCopy := *c.netMap
+	c.wireguardEngine.SetNetworkMap(&netMapCopy)
+	cfg, err := nmcfg.WGCfg(c.netMap, Logger(c.logger.Named("wgconfig")), netmap.AllowSingleHosts, "")
+	if err != nil {
+		return xerrors.Errorf("update wireguard config: %w", err)
+	}
+	err = c.wireguardEngine.Reconfig(cfg, c.wireguardRouter, &dns.Config{}, &tailcfg.Debug{})
+	if err != nil {
+		if c.isClosed() {
+			return nil
+		}
+		if errors.Is(err, wgengine.ErrNoChanges) {
+			return nil
+		}
+		return xerrors.Errorf("reconfig: %w", err)
+	}
+	return nil
+}
+
 // UpdateNodes connects with a set of peers. This can be constantly updated,
 // and peers will continually be reconnected as necessary.
 func (c *Conn) UpdateNodes(nodes []*Node) error {
