fix: Close tailnet if agent is closed during creation (#5375)

mafredri · web-flow · commit 88bb90128305 · 2022-12-12T11:26:49.000Z
diff --git a/agent/agent.go b/agent/agent.go
@@ -215,8 +215,16 @@ func (a *agent) run(ctx context.Context) error {
 			return xerrors.Errorf("create tailnet: %w", err)
 		}
 		a.closeMutex.Lock()
-		a.network = network
+		// Re-check if agent was closed while initializing the network.
+		closed := a.isClosed()
+		if !closed {
+			a.network = network
+		}
 		a.closeMutex.Unlock()
+		if closed {
+			_ = network.Close()
+			return xerrors.New("agent is closed")
+		}
 	} else {
 		// Update the DERP map!
 		network.SetDERPMap(metadata.DERPMap)
@@ -246,27 +254,20 @@ func (a *agent) trackConnGoroutine(fn func()) error {
 }
 
 func (a *agent) createTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) (_ *tailnet.Conn, err error) {
-	a.closeMutex.Lock()
-	if a.isClosed() {
-		a.closeMutex.Unlock()
-		return nil, xerrors.New("closed")
-	}
 	network, err := tailnet.NewConn(&tailnet.Options{
 		Addresses:          []netip.Prefix{netip.PrefixFrom(codersdk.TailnetIP, 128)},
 		DERPMap:            derpMap,
 		Logger:             a.logger.Named("tailnet"),
 		EnableTrafficStats: true,
 	})
 	if err != nil {
-		a.closeMutex.Unlock()
 		return nil, xerrors.Errorf("create tailnet: %w", err)
 	}
 	defer func() {
 		if err != nil {
 			network.Close()
 		}
 	}()
-	a.closeMutex.Unlock()
 
 	sshListener, err := network.Listen("tcp", ":"+strconv.Itoa(codersdk.TailnetSSHPort))
 	if err != nil {
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -806,12 +806,17 @@ func setupAgent(t *testing.T, metadata codersdk.WorkspaceAgentMetadata, ptyTimeo
 	})
 	require.NoError(t, err)
 	clientConn, serverConn := net.Pipe()
+	serveClientDone := make(chan struct{})
 	t.Cleanup(func() {
 		_ = clientConn.Close()
 		_ = serverConn.Close()
 		_ = conn.Close()
+		<-serveClientDone
 	})
-	go coordinator.ServeClient(serverConn, uuid.New(), agentID)
+	go func() {
+		defer close(serveClientDone)
+		coordinator.ServeClient(serverConn, uuid.New(), agentID)
+	}()
 	sendNode, _ := tailnet.ServeCoordinator(clientConn, func(node []*tailnet.Node) error {
 		return conn.UpdateNodes(node)
 	})
