tests for derp map changing

deansheather · deansheather · commit 0e6d39a7abbe · 2023-06-13T19:18:20.000Z
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -1564,6 +1564,109 @@ func TestAgent_Dial(t *testing.T) {
 	}
 }
 
+// TestAgent_UpdatedDERP checks that agents can handle their DERP map being
+// updated, and that clients can also handle it.
+func TestAgent_UpdatedDERP(t *testing.T) {
+	t.Parallel()
+
+	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
+
+	derpMap := *tailnettest.RunDERPAndSTUN(t)
+	metadata := agentsdk.Manifest{
+		DERPMap: &derpMap,
+	}
+	coordinator := tailnet.NewCoordinator(logger, func() *tailcfg.DERPMap {
+		return &derpMap
+	})
+	defer func() {
+		_ = coordinator.Close()
+	}()
+	agentID := uuid.New()
+	statsCh := make(chan *agentsdk.Stats, 50)
+	fs := afero.NewMemMapFs()
+	c := &client{
+		t:           t,
+		agentID:     agentID,
+		manifest:    metadata,
+		statsChan:   statsCh,
+		coordinator: coordinator,
+	}
+	closer := agent.New(agent.Options{
+		Client:                 c,
+		Filesystem:             fs,
+		Logger:                 logger.Named("agent"),
+		ReconnectingPTYTimeout: time.Minute,
+	})
+	defer func() {
+		_ = closer.Close()
+	}()
+
+	// Setup a client connection.
+	newClientConn := func() *codersdk.WorkspaceAgentConn {
+		conn, err := tailnet.NewConn(&tailnet.Options{
+			Addresses: []netip.Prefix{netip.PrefixFrom(tailnet.IP(), 128)},
+			DERPMap:   metadata.DERPMap,
+			Logger:    logger.Named("client"),
+		})
+		require.NoError(t, err)
+		clientConn, serverConn := net.Pipe()
+		serveClientDone := make(chan struct{})
+		t.Cleanup(func() {
+			_ = clientConn.Close()
+			_ = serverConn.Close()
+			_ = conn.Close()
+			<-serveClientDone
+		})
+		go func() {
+			defer close(serveClientDone)
+			err := coordinator.ServeClient(serverConn, uuid.New(), agentID)
+			assert.NoError(t, err)
+		}()
+		sendNode, _ := tailnet.ServeCoordinator(clientConn, func(update tailnet.CoordinatorNodeUpdate) error {
+			if tailnet.CompareDERPMaps(conn.DERPMap(), update.DERPMap) {
+				conn.SetDERPMap(update.DERPMap)
+			}
+			return conn.UpdateNodes(update.Nodes, false)
+		})
+		conn.SetNodeCallback(sendNode)
+
+		sdkConn := &codersdk.WorkspaceAgentConn{
+			Conn: conn,
+		}
+		t.Cleanup(func() {
+			_ = sdkConn.Close()
+		})
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+		if !sdkConn.AwaitReachable(ctx) {
+			t.Fatal("agent not reachable")
+		}
+
+		return sdkConn
+	}
+	conn1 := newClientConn()
+
+	// Change the DERP map.
+	derpMap = *tailnettest.RunDERPAndSTUN(t)
+	// Change the region ID.
+	derpMap.Regions[2] = derpMap.Regions[1]
+	delete(derpMap.Regions, 1)
+	derpMap.Regions[2].RegionID = 2
+	for _, node := range derpMap.Regions[2].Nodes {
+		node.RegionID = 2
+	}
+
+	// Connect from a second client and make sure it uses the new DERP map.
+	conn2 := newClientConn()
+	require.Equal(t, []int{2}, conn2.DERPMap().RegionIDs())
+
+	// Make sure the first client is still reachable (it received a a DERP map
+	// update when the agent's nodes changed).
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+	defer cancel()
+	require.True(t, conn1.AwaitReachable(ctx))
+}
+
 func TestAgent_Speedtest(t *testing.T) {
 	t.Parallel()
 	t.Skip("This test is relatively flakey because of Tailscale's speedtest code...")
diff --git a/tailnet/conn_test.go b/tailnet/conn_test.go
@@ -195,3 +195,124 @@ func TestConn_PreferredDERP(t *testing.T) {
 		t.Fatal("timed out waiting for node")
 	}
 }
+
+// TestConn_UpdateDERP tests that when update the DERP map we pick a new
+// preferred DERP server and new connections can be made from clients.
+func TestConn_UpdateDERP(t *testing.T) {
+	t.Parallel()
+	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
+
+	derpMap1 := tailnettest.RunDERPAndSTUN(t)
+	ip := tailnet.IP()
+	conn, err := tailnet.NewConn(&tailnet.Options{
+		Addresses:      []netip.Prefix{netip.PrefixFrom(ip, 128)},
+		Logger:         logger.Named("w1"),
+		DERPMap:        derpMap1,
+		BlockEndpoints: true,
+	})
+	require.NoError(t, err)
+	defer func() {
+		err := conn.Close()
+		assert.NoError(t, err)
+	}()
+
+	// Buffer channel so callback doesn't block
+	nodes := make(chan *tailnet.Node, 50)
+	conn.SetNodeCallback(func(node *tailnet.Node) {
+		nodes <- node
+	})
+
+	ctx1, cancel1 := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer cancel1()
+	select {
+	case node := <-nodes:
+		require.Equal(t, 1, node.PreferredDERP)
+	case <-ctx1.Done():
+		t.Fatal("timed out waiting for node")
+	}
+
+	// Connect from a different client.
+	client1, err := tailnet.NewConn(&tailnet.Options{
+		Addresses:      []netip.Prefix{netip.PrefixFrom(tailnet.IP(), 128)},
+		Logger:         logger.Named("client1"),
+		DERPMap:        derpMap1,
+		BlockEndpoints: true,
+	})
+	require.NoError(t, err)
+	defer func() {
+		err := client1.Close()
+		assert.NoError(t, err)
+	}()
+	client1.SetNodeCallback(func(node *tailnet.Node) {
+		err := conn.UpdateNodes([]*tailnet.Node{node}, false)
+		assert.NoError(t, err)
+	})
+	client1.UpdateNodes([]*tailnet.Node{conn.Node()}, false)
+
+	awaitReachableCtx1, awaitReachableCancel1 := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer awaitReachableCancel1()
+	require.True(t, client1.AwaitReachable(awaitReachableCtx1, ip))
+
+	// Update the DERP map and wait for the preferred DERP server to change.
+	derpMap2 := tailnettest.RunDERPAndSTUN(t)
+	// Change the region ID.
+	derpMap2.Regions[2] = derpMap2.Regions[1]
+	delete(derpMap2.Regions, 1)
+	derpMap2.Regions[2].RegionID = 2
+	for _, node := range derpMap2.Regions[2].Nodes {
+		node.RegionID = 2
+	}
+	conn.SetDERPMap(derpMap2)
+
+	ctx2, cancel2 := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer cancel2()
+parentLoop:
+	for {
+		select {
+		case node := <-nodes:
+			if node.PreferredDERP != 2 {
+				t.Logf("waiting for preferred DERP server to change, got %v", node.PreferredDERP)
+				continue
+			}
+			t.Log("preferred DERP server changed!")
+			break parentLoop
+		case <-ctx2.Done():
+			t.Fatal("timed out waiting for preferred DERP server to change")
+		}
+	}
+
+	// Client1 should be dropped...
+	awaitReachableCtx2, awaitReachableCancel2 := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer awaitReachableCancel2()
+	require.False(t, client1.AwaitReachable(awaitReachableCtx2, ip))
+
+	// ... unless the client updates it's derp map and nodes.
+	client1.SetDERPMap(derpMap2)
+	client1.UpdateNodes([]*tailnet.Node{conn.Node()}, false)
+	awaitReachableCtx3, awaitReachableCancel3 := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer awaitReachableCancel3()
+	require.True(t, client1.AwaitReachable(awaitReachableCtx3, ip))
+
+	// Connect from a different different client with up-to-date derp map and
+	// nodes.
+	client2, err := tailnet.NewConn(&tailnet.Options{
+		Addresses:      []netip.Prefix{netip.PrefixFrom(tailnet.IP(), 128)},
+		Logger:         logger.Named("client2"),
+		DERPMap:        derpMap2,
+		BlockEndpoints: true,
+	})
+	require.NoError(t, err)
+	defer func() {
+		err := client2.Close()
+		assert.NoError(t, err)
+	}()
+	client2.SetNodeCallback(func(node *tailnet.Node) {
+		err := conn.UpdateNodes([]*tailnet.Node{node}, false)
+		assert.NoError(t, err)
+	})
+	client2.UpdateNodes([]*tailnet.Node{conn.Node()}, false)
+
+	awaitReachableCtx4, awaitReachableCancel4 := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer awaitReachableCancel4()
+	require.True(t, client2.AwaitReachable(awaitReachableCtx4, ip))
+}
