fix: Ensure tailnet coordinator can never be blocked

kylecarbs · kylecarbs · commit 35bf9af78bd5 · 2022-09-25T17:56:58.000Z
diff --git a/tailnet/coordinator.go b/tailnet/coordinator.go
@@ -7,6 +7,7 @@ import (
 	"net"
 	"net/netip"
 	"sync"
+	"time"
 
 	"github.com/google/uuid"
 	"golang.org/x/xerrors"
@@ -164,23 +165,25 @@ 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})
 		if err != nil {
 			c.mutex.Unlock()
 			return xerrors.Errorf("marshal nodes: %w", err)
 		}
+		err = agentSocket.SetWriteDeadline(time.Now().Add(5 * time.Second))
+		if err != nil {
+			return xerrors.Errorf("set write deadline: %w", err)
+		}
 		_, 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()
 	}
 }
 
@@ -255,11 +258,12 @@ func (c *Coordinator) ServeAgent(conn net.Conn, id uuid.UUID) error {
 		for _, connectionSocket := range connectionSockets {
 			connectionSocket := connectionSocket
 			go func() {
+				_ = connectionSocket.SetWriteDeadline(time.Now().Add(5 * time.Second))
 				_, _ = connectionSocket.Write(data)
 				wg.Done()
 			}()
 		}
-		wg.Wait()
 		c.mutex.Unlock()
+		wg.Wait()
 	}
 }

Original file line number	Diff line number	Diff line change
`@@ -7,6 +7,7 @@ import (`
`7`	`7`	`"net"`
`8`	`8`	`"net/netip"`
`9`	`9`	`"sync"`
	`10`	`+ "time"`
`10`	`11`
`11`	`12`	`"github.com/google/uuid"`
`12`	`13`	`"golang.org/x/xerrors"`
`@@ -164,23 +165,25 @@ func (c *Coordinator) ServeClient(conn net.Conn, id uuid.UUID, agent uuid.UUID)`
`164`	`165`	`c.mutex.Unlock()`
`165`	`166`	`continue`
`166`	`167`	`}`
	`168`	`+ c.mutex.Unlock()`
`167`	`169`	`// Write the new node from this client to the actively`
`168`	`170`	`// connected agent.`
`169`	`171`	`data, err := json.Marshal([]*Node{&node})`
`170`	`172`	`if err != nil {`
`171`	`173`	`c.mutex.Unlock()`
`172`	`174`	`return xerrors.Errorf("marshal nodes: %w", err)`
`173`	`175`	`}`
	`176`	`+ err = agentSocket.SetWriteDeadline(time.Now().Add(5 * time.Second))`
	`177`	`+ if err != nil {`
	`178`	`+ return xerrors.Errorf("set write deadline: %w", err)`
	`179`	`+ }`
`174`	`180`	`_, err = agentSocket.Write(data)`
`175`	`181`	`if errors.Is(err, io.EOF) {`
`176`		`- c.mutex.Unlock()`
`177`	`182`	`return nil`
`178`	`183`	`}`
`179`	`184`	`if err != nil {`
`180`		`- c.mutex.Unlock()`
`181`	`185`	`return xerrors.Errorf("write json: %w", err)`
`182`	`186`	`}`
`183`		`- c.mutex.Unlock()`
`184`	`187`	`}`
`185`	`188`	`}`
`186`	`189`
`@@ -255,11 +258,12 @@ func (c *Coordinator) ServeAgent(conn net.Conn, id uuid.UUID) error {`
`255`	`258`	`for _, connectionSocket := range connectionSockets {`
`256`	`259`	`connectionSocket := connectionSocket`
`257`	`260`	`go func() {`
	`261`	`+ _ = connectionSocket.SetWriteDeadline(time.Now().Add(5 * time.Second))`
`258`	`262`	`_, _ = connectionSocket.Write(data)`
`259`	`263`	`wg.Done()`
`260`	`264`	`}()`
`261`	`265`	`}`
`262`		`- wg.Wait()`
`263`	`266`	`c.mutex.Unlock()`
	`267`	`+ wg.Wait()`
`264`	`268`	`}`
`265`	`269`	`}`