add auth to in-memory coordinator

coder · coadler · Apr 10, 2024 · Mar 27, 2024 · Mar 29, 2024 · Mar 29, 2024
commit 5660e03d02c73bedada4082d627208a078b5c92a
diff --git a/tailnet/coordinator.go b/tailnet/coordinator.go
@@ -711,6 +711,10 @@ func (c *core) handleReadyForHandshakeLocked(src *peer, rfhs []*proto.Coordinate
 			return xerrors.Errorf("unable to convert bytes to UUID: %w", err)
 		}
 
+		if !c.tunnels.tunnelExists(src.id, dstID) {
+			return xerrors.Errorf("tunnel does not exist between %s and %s", src.id.String(), dstID.String())
+		}
+
 		dst, ok := c.peers[dstID]
 		if ok {
 			select {

diff --git a/tailnet/coordinator_test.go b/tailnet/coordinator_test.go
@@ -422,8 +422,28 @@ func TestCoordinator(t *testing.T) {
 		clientID := uuid.New()
 		agentID := uuid.New()
 
-		aReq, _ := coordinator.Coordinate(ctx, agentID, agentID.String(), tailnet.AgentCoordinateeAuth{ID: agentID})
-		_, cRes := coordinator.Coordinate(ctx, clientID, clientID.String(), tailnet.ClientCoordinateeAuth{AgentID: agentID})
+		aReq, aRes := coordinator.Coordinate(ctx, agentID, agentID.String(), tailnet.AgentCoordinateeAuth{ID: agentID})
+		cReq, cRes := coordinator.Coordinate(ctx, clientID, clientID.String(), tailnet.ClientCoordinateeAuth{AgentID: agentID})
+
+		{
+			nk, err := key.NewNode().Public().MarshalBinary()
+			require.NoError(t, err)
+			dk, err := key.NewDisco().Public().MarshalText()
+			require.NoError(t, err)
+			cReq <- &proto.CoordinateRequest{UpdateSelf: &proto.CoordinateRequest_UpdateSelf{
+				Node: &proto.Node{
+					Id:    3,
+					Key:   nk,
+					Disco: string(dk),
+				},
+			}}
+		}
+
+		cReq <- &proto.CoordinateRequest{AddTunnel: &proto.CoordinateRequest_Tunnel{
+			Id: agentID[:],
+		}}
+
+		testutil.RequireRecvCtx(ctx, t, aRes)
 
 		aReq <- &proto.CoordinateRequest{ReadyForHandshake: []*proto.CoordinateRequest_ReadyForHandshake{{
 			Id: clientID[:],
@@ -434,6 +454,44 @@ func TestCoordinator(t *testing.T) {
 		require.Equal(t, proto.CoordinateResponse_PeerUpdate_READY_FOR_HANDSHAKE, ack.PeerUpdates[0].Kind)
 		require.Equal(t, agentID[:], ack.PeerUpdates[0].Id)
 	})
+
+	t.Run("AgentAck_NoPermission", func(t *testing.T) {
+		t.Parallel()
+		logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
+		coordinator := tailnet.NewCoordinator(logger)
+		ctx := testutil.Context(t, testutil.WaitShort)
+
+		clientID := uuid.New()
+		agentID := uuid.New()
+
+		aReq, _ := coordinator.Coordinate(ctx, agentID, agentID.String(), tailnet.AgentCoordinateeAuth{ID: agentID})
+		_, _ = coordinator.Coordinate(ctx, clientID, clientID.String(), tailnet.ClientCoordinateeAuth{AgentID: agentID})
+
+		nk, err := key.NewNode().Public().MarshalBinary()
+		require.NoError(t, err)
+		dk, err := key.NewDisco().Public().MarshalText()
+		require.NoError(t, err)
+		aReq <- &proto.CoordinateRequest{UpdateSelf: &proto.CoordinateRequest_UpdateSelf{
+			Node: &proto.Node{
+				Id:    3,
+				Key:   nk,
+				Disco: string(dk),
+			},
+		}}
+
+		require.Eventually(t, func() bool {
+			return coordinator.Node(agentID) != nil
+		}, testutil.WaitShort, testutil.IntervalFast)
+
+		aReq <- &proto.CoordinateRequest{ReadyForHandshake: []*proto.CoordinateRequest_ReadyForHandshake{{
+			Id: clientID[:],
+		}}}
+
+		// The agent node should disappear, indicating it was booted off.
+		require.Eventually(t, func() bool {
+			return coordinator.Node(agentID) == nil
+		}, testutil.WaitShort, testutil.IntervalFast)
+	})
 }
 
 // TestCoordinator_AgentUpdateWhileClientConnects tests for regression on

diff --git a/tailnet/tunnel.go b/tailnet/tunnel.go
@@ -151,6 +151,12 @@ func (s *tunnelStore) findTunnelPeers(id uuid.UUID) []uuid.UUID {
 	return out
 }
 
+func (s *tunnelStore) tunnelExists(src, dst uuid.UUID) bool {
+	_, srcOK := s.bySrc[src][dst]
+	_, dstOK := s.byDst[src][dst]
+	return srcOK || dstOK
+}
+
 func (s *tunnelStore) htmlDebug() []HTMLTunnel {
 	out := make([]HTMLTunnel, 0)
 	for src, dsts := range s.bySrc {

diff --git a/tailnet/tunnel_internal_test.go b/tailnet/tunnel_internal_test.go
@@ -43,3 +43,18 @@ func TestTunnelStore_RemoveAll(t *testing.T) {
 	require.Len(t, uut.findTunnelPeers(p2), 0)
 	require.Len(t, uut.findTunnelPeers(p3), 0)
 }
+
+func TestTunnelStore_TunnelExists(t *testing.T) {
+	t.Parallel()
+	p1 := uuid.UUID{1}
+	p2 := uuid.UUID{2}
+	uut := newTunnelStore()
+	require.False(t, uut.tunnelExists(p1, p2))
+	require.False(t, uut.tunnelExists(p2, p1))
+	uut.add(p1, p2)
+	require.True(t, uut.tunnelExists(p1, p2))
+	require.True(t, uut.tunnelExists(p2, p1))
+	uut.remove(p1, p2)
+	require.False(t, uut.tunnelExists(p1, p2))
+	require.False(t, uut.tunnelExists(p2, p1))
+}