Fix client publish race

coder · kylecarbs · Oct 17, 2022 · Sep 22, 2022 · Sep 23, 2022 · Sep 23, 2022
commit 9878fc51d4b6a464e87bcb5f59f6891e1e39d5d9
diff --git a/enterprise/tailnet/coordinator.go b/enterprise/tailnet/coordinator.go
@@ -56,8 +56,8 @@ type haCoordinator struct {
 
 // Node returns an in-memory node by ID.
 func (c *haCoordinator) Node(id uuid.UUID) *agpl.Node {
-	c.mutex.RLock()
-	defer c.mutex.RUnlock()
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
 	node := c.nodes[id]
 	return node
 }
@@ -79,6 +79,11 @@ func (c *haCoordinator) ServeClient(conn net.Conn, id uuid.UUID, agent uuid.UUID
 		if err != nil {
 			return xerrors.Errorf("write nodes: %w", err)
 		}
+	} else {
+		err := c.publishClientHello(agent)
+		if err != nil {
+			return xerrors.Errorf("publish client hello: %w", err)
+		}
 	}
 
 	c.mutex.Lock()
@@ -205,7 +210,7 @@ func (c *haCoordinator) ServeAgent(conn net.Conn, id uuid.UUID) error {
 
 	decoder := json.NewDecoder(conn)
 	for {
-		node, err := c.hangleAgentUpdate(id, decoder)
+		node, err := c.handleAgentUpdate(id, decoder)
 		if err != nil {
 			if errors.Is(err, io.EOF) {
 				return nil
@@ -240,7 +245,17 @@ func (c *haCoordinator) nodesSubscribedToAgent(agentID uuid.UUID) []*agpl.Node {
 	return nodes
 }
 
-func (c *haCoordinator) hangleAgentUpdate(id uuid.UUID, decoder *json.Decoder) (*agpl.Node, error) {
+func (c *haCoordinator) handleClientHello(id uuid.UUID) error {
+	c.mutex.Lock()
+	node, ok := c.nodes[id]
+	c.mutex.Unlock()
+	if !ok {
+		return nil
+	}
+	return c.publishAgentToNodes(id, node)
+}
+
+func (c *haCoordinator) handleAgentUpdate(id uuid.UUID, decoder *json.Decoder) (*agpl.Node, error) {
 	var node agpl.Node
 	err := decoder.Decode(&node)
 	if err != nil {
@@ -343,6 +358,18 @@ func (c *haCoordinator) publishAgentHello(id uuid.UUID) error {
 	return nil
 }
 
+func (c *haCoordinator) publishClientHello(id uuid.UUID) error {
+	msg, err := c.formatClientHello(id)
+	if err != nil {
+		return xerrors.Errorf("format client hello: %w", err)
+	}
+	err = c.pubsub.Publish("wireguard_peers", msg)
+	if err != nil {
+		return xerrors.Errorf("publish client hello: %w", err)
+	}
+	return nil
+}
+
 func (c *haCoordinator) publishAgentToNodes(id uuid.UUID, node *agpl.Node) error {
 	msg, err := c.formatAgentUpdate(id, node)
 	if err != nil {
@@ -408,6 +435,18 @@ func (c *haCoordinator) runPubsub() error {
 				c.log.Error(ctx, "send callmemaybe to agent", slog.Error(err))
 				return
 			}
+		case "clienthello":
+			agentUUID, err := uuid.ParseBytes(agentID)
+			if err != nil {
+				c.log.Error(ctx, "invalid agent id", slog.F("id", string(agentID)))
+				return
+			}
+
+			err = c.handleClientHello(agentUUID)
+			if err != nil {
+				c.log.Error(ctx, "handle agent request node", slog.Error(err))
+				return
+			}
 		case "agenthello":
 			agentUUID, err := uuid.ParseBytes(agentID)
 			if err != nil {
@@ -431,7 +470,7 @@ func (c *haCoordinator) runPubsub() error {
 			}
 
 			decoder := json.NewDecoder(bytes.NewReader(nodeJSON))
-			_, err = c.hangleAgentUpdate(agentUUID, decoder)
+			_, err = c.handleAgentUpdate(agentUUID, decoder)
 			if err != nil {
 				c.log.Error(ctx, "handle agent update", slog.Error(err))
 				return
@@ -478,6 +517,17 @@ func (c *haCoordinator) formatAgentHello(id uuid.UUID) ([]byte, error) {
 	return buf.Bytes(), nil
 }
 
+// format: <coordinator id>|clienthello|<agent id>|
+func (c *haCoordinator) formatClientHello(id uuid.UUID) ([]byte, error) {
+	buf := bytes.Buffer{}
+
+	buf.WriteString(c.id.String() + "|")
+	buf.WriteString("clienthello|")
+	buf.WriteString(id.String() + "|")
+
+	return buf.Bytes(), nil
+}
+
 // format: <coordinator id>|agentupdate|<node id>|<node json>
 func (c *haCoordinator) formatAgentUpdate(id uuid.UUID, node *agpl.Node) ([]byte, error) {
 	buf := bytes.Buffer{}

diff --git a/enterprise/tailnet/coordinator_test.go b/enterprise/tailnet/coordinator_test.go
@@ -11,6 +11,7 @@ import (
 	"cdr.dev/slog/sloggers/slogtest"
 
 	"github.com/coder/coder/coderd/database"
+	"github.com/coder/coder/coderd/database/dbtestutil"
 	"github.com/coder/coder/enterprise/tailnet"
 	agpl "github.com/coder/coder/tailnet"
 	"github.com/coder/coder/testutil"
@@ -167,16 +168,12 @@ func TestCoordinatorHA(t *testing.T) {
 	t.Run("AgentWithClient", func(t *testing.T) {
 		t.Parallel()
 
-		pubsub := database.NewPubsubInMemory()
+		_, pubsub := dbtestutil.NewDB(t)
 
 		coordinator1, err := tailnet.NewCoordinator(slogtest.Make(t, nil), pubsub)
 		require.NoError(t, err)
 		defer coordinator1.Close()
 
-		coordinator2, err := tailnet.NewCoordinator(slogtest.Make(t, nil), pubsub)
-		require.NoError(t, err)
-		defer coordinator2.Close()
-
 		agentWS, agentServerWS := net.Pipe()
 		defer agentWS.Close()
 		agentNodeChan := make(chan []*agpl.Node)
@@ -196,6 +193,10 @@ func TestCoordinatorHA(t *testing.T) {
 			return coordinator1.Node(agentID) != nil
 		}, testutil.WaitShort, testutil.IntervalFast)
 
+		coordinator2, err := tailnet.NewCoordinator(slogtest.Make(t, nil), pubsub)
+		require.NoError(t, err)
+		defer coordinator2.Close()
+
 		clientWS, clientServerWS := net.Pipe()
 		defer clientWS.Close()
 		defer clientServerWS.Close()