feat: add negotiated binary UDP packets (#5456)

2026-07-31 15:52:54 +08:00 · 2026-07-30 13:19:55 +08:00
parent 5c6d761c12
commit effa496859
23 changed files with 1232 additions and 71 deletions
@@ -79,6 +79,70 @@ func TestWriteWithDeadlineTimesOutAndClearsDeadline(t *testing.T) {
 	}
 }

+func TestServiceAcceptConnectionTracksClientHelloPresence(t *testing.T) {
+	for _, tc := range []struct {
+		name               string
+		clientHelloPresent bool
+		offeredCodecs      []string
+		expectedCodec      string
+	}{
+		{
+			name: "absent Hello",
+		},
+		{
+			name:               "present Hello with JSON fallback",
+			clientHelloPresent: true,
+		},
+		{
+			name:               "present Hello with binary codec",
+			clientHelloPresent: true,
+			offeredCodecs:      []string{wire.UDPPacketCodecBinary},
+			expectedCodec:      wire.UDPPacketCodecBinary,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			serverConn, clientConn := net.Pipe()
+			defer serverConn.Close()
+			defer clientConn.Close()
+
+			clientErrCh := make(chan error, 1)
+			go func() {
+				if err := wire.WriteMagic(clientConn); err != nil {
+					clientErrCh <- err
+					return
+				}
+				wireConn := wire.NewConn(clientConn)
+				if tc.clientHelloPresent {
+					hello, err := wire.NewClientHello(wire.BootstrapInfo{})
+					if err != nil {
+						clientErrCh <- err
+						return
+					}
+					hello.Capabilities.Message.UDPPacketCodecs = tc.offeredCodecs
+					if err := wireConn.WriteJSONFrame(wire.FrameTypeClientHello, hello); err != nil {
+						clientErrCh <- err
+						return
+					}
+					var serverHello wire.ServerHello
+					if err := wireConn.ReadJSONFrame(wire.FrameTypeServerHello, &serverHello); err != nil {
+						clientErrCh <- err
+						return
+					}
+				}
+				clientErrCh <- msg.NewV2ReadWriterWithConn(wireConn).WriteMsg(&msg.NewWorkConn{RunID: "shared-run"})
+			}()
+
+			acceptedConn, err := (&Service{}).acceptConnection(t.Context(), serverConn)
+			require.NoError(t, err)
+			require.NoError(t, <-clientErrCh)
+			require.Equal(t, tc.clientHelloPresent, acceptedConn.clientHelloPresent)
+			require.Equal(t, tc.expectedCodec, acceptedConn.udpPacketCodec)
+			require.IsType(t, &msg.NewWorkConn{}, acceptedConn.firstMsg)
+			require.NoError(t, acceptedConn.conn.Close())
+		})
+	}
+}
+
 func TestSharedPortHTTPListenerProtocols(t *testing.T) {
 	listener, err := net.Listen("tcp", "127.0.0.1:0")
 	require.NoError(t, err)
@@ -428,7 +492,7 @@ func TestServiceWorkConnRoutingRequiresCurrentRunningControl(t *testing.T) {

 	pendingConn := newCountingCloseConn()
 	pendingMsgConn := msg.NewConn(pendingConn, msg.NewV1ReadWriter(pendingConn))
-	err = registerWorkConnAsCaller(svr, pendingMsgConn, &msg.NewWorkConn{RunID: "shared-run"})
+	err = registerWorkConnAsCaller(svr, pendingMsgConn, &msg.NewWorkConn{RunID: "shared-run"}, wire.ProtocolV1, false)
 	require.Error(t, err)
 	require.Equal(t, int64(1), pendingConn.closeCount.Load())
 	require.Len(t, ctl.workConnCh, 0)
@@ -442,7 +506,7 @@ func TestServiceWorkConnRoutingRequiresCurrentRunningControl(t *testing.T) {

 	runningConn := newCountingCloseConn()
 	runningMsgConn := msg.NewConn(runningConn, msg.NewV1ReadWriter(runningConn))
-	require.NoError(t, svr.RegisterWorkConn(runningMsgConn, &msg.NewWorkConn{RunID: "shared-run"}))
+	require.NoError(t, svr.RegisterWorkConn(runningMsgConn, &msg.NewWorkConn{RunID: "shared-run"}, wire.ProtocolV1, false))
 	require.Len(t, ctl.workConnCh, 1)

 	require.NoError(t, ctl.Close())
@@ -450,6 +514,123 @@ func TestServiceWorkConnRoutingRequiresCurrentRunningControl(t *testing.T) {
 	require.Equal(t, int64(1), runningConn.closeCount.Load())
 }

+func TestServiceWorkConnRoutingRejectsWireProtocolMismatch(t *testing.T) {
+	svr := newControlTestService(t)
+	ctl, controlConn, err := registerLifecycleTestControl(svr)
+	require.NoError(t, err)
+	require.NoError(t, svr.completeControlLogin(ctl, func() error { return nil }))
+	waitForSignal(t, controlConn.readStarted, "control reader to start")
+
+	workConn := newCountingCloseConn()
+	workMsgConn := msg.NewConn(workConn, msg.NewV2ReadWriter(workConn))
+	err = svr.RegisterWorkConn(workMsgConn, &msg.NewWorkConn{RunID: "shared-run"}, wire.ProtocolV2, false)
+	require.ErrorContains(t, err, "wire protocol mismatch")
+	require.Len(t, ctl.workConnCh, 0)
+	_ = workMsgConn.Close()
+	require.NoError(t, ctl.Close())
+}
+
+func TestServiceWorkConnRoutingClientHelloPolicy(t *testing.T) {
+	for _, tc := range []struct {
+		name                   string
+		controlUDPPacketCodec  string
+		workClientHelloPresent bool
+		errorSubstring         string
+	}{
+		{
+			name: "JSON control allows work connection without Hello",
+		},
+		{
+			name:                  "binary control allows work connection without Hello",
+			controlUDPPacketCodec: wire.UDPPacketCodecBinary,
+		},
+		{
+			name:                   "JSON control rejects work connection with Hello",
+			workClientHelloPresent: true,
+			errorSubstring:         "ClientHello is not allowed",
+		},
+		{
+			name:                   "binary control rejects work connection with Hello",
+			controlUDPPacketCodec:  wire.UDPPacketCodecBinary,
+			workClientHelloPresent: true,
+			errorSubstring:         "ClientHello is not allowed",
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			svr := newControlTestService(t)
+			controlConn := newDeadlineReadConn()
+			controlMsgConn := msg.NewConn(controlConn, msg.NewV2ReadWriter(controlConn))
+			ctl, err := svr.RegisterControl(controlMsgConn, &msg.Login{
+				RunID:    "shared-run",
+				ClientID: "client",
+				ClientSpec: msg.ClientSpec{
+					AlwaysAuthPass: true,
+				},
+			}, true, wire.ProtocolV2, tc.controlUDPPacketCodec)
+			require.NoError(t, err)
+			require.NoError(t, svr.completeControlLogin(ctl, func() error { return nil }))
+			waitForSignal(t, controlConn.readStarted, "control reader to start")
+
+			workConn := newCountingCloseConn()
+			workMsgConn := msg.NewConn(workConn, msg.NewV2ReadWriter(workConn))
+			err = svr.RegisterWorkConn(
+				workMsgConn,
+				&msg.NewWorkConn{RunID: "shared-run"},
+				wire.ProtocolV2,
+				tc.workClientHelloPresent,
+			)
+			if tc.errorSubstring != "" {
+				require.ErrorContains(t, err, tc.errorSubstring)
+				require.Len(t, ctl.workConnCh, 0)
+				require.NoError(t, workMsgConn.Close())
+			} else {
+				require.NoError(t, err)
+				require.Len(t, ctl.workConnCh, 1)
+			}
+
+			require.NoError(t, ctl.Close())
+			waitForControlDone(t, ctl)
+			require.Equal(t, int64(1), workConn.closeCount.Load())
+		})
+	}
+}
+
+func TestServiceRegisterControlRejectsInvalidCodecSelection(t *testing.T) {
+	for _, tc := range []struct {
+		name           string
+		wireProtocol   string
+		udpPacketCodec string
+		errorSubstring string
+	}{
+		{
+			name:           "binary codec over v1",
+			wireProtocol:   wire.ProtocolV1,
+			udpPacketCodec: wire.UDPPacketCodecBinary,
+			errorSubstring: "requires wire protocol v2",
+		},
+		{
+			name:           "unknown v2 codec",
+			wireProtocol:   wire.ProtocolV2,
+			udpPacketCodec: "unknown",
+			errorSubstring: "unsupported UDP packet codec",
+		},
+		{
+			name:           "unknown wire protocol",
+			wireProtocol:   "unknown",
+			errorSubstring: "unsupported wire protocol",
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			svr := newControlTestService(t)
+			conn := newDeadlineReadConn()
+			msgConn := msg.NewConn(conn, msg.NewV1ReadWriter(conn))
+			ctl, err := svr.RegisterControl(msgConn, &msg.Login{}, true, tc.wireProtocol, tc.udpPacketCodec)
+			require.Nil(t, ctl)
+			require.ErrorContains(t, err, tc.errorSubstring)
+		})
+	}
+}
+
 func TestServiceWorkConnRoutingRejectsLostGeneration(t *testing.T) {
 	for _, action := range []string{"replace", "close"} {
 		t.Run(action, func(t *testing.T) {
@@ -465,7 +646,7 @@ func TestServiceWorkConnRoutingRejectsLostGeneration(t *testing.T) {
 			workMsgConn := msg.NewConn(workConn, msg.NewV1ReadWriter(workConn))
 			routeDone := make(chan error, 1)
 			go func() {
-				routeDone <- registerWorkConnAsCaller(svr, workMsgConn, &msg.NewWorkConn{RunID: "shared-run"})
+				routeDone <- registerWorkConnAsCaller(svr, workMsgConn, &msg.NewWorkConn{RunID: "shared-run"}, wire.ProtocolV1, false)
 			}()
 			waitForSignal(t, barrier.entered, "work connection plugin barrier")

@@ -509,7 +690,7 @@ func TestServiceVisitorRoutingExcludesPendingUser(t *testing.T) {
 		ClientSpec: msg.ClientSpec{
 			AlwaysAuthPass: true,
 		},
-	}, true, wire.ProtocolV1)
+	}, true, wire.ProtocolV1, "")
 	require.NoError(t, err)

 	timestamp := time.Now().Unix()
@@ -567,7 +748,7 @@ func registerLifecycleTestControl(svr *Service) (*Control, *deadlineReadConn, er
 		ClientSpec: msg.ClientSpec{
 			AlwaysAuthPass: true,
 		},
-	}, true, wire.ProtocolV1)
+	}, true, wire.ProtocolV1, "")
 	return ctl, conn, err
 }

@@ -584,8 +765,14 @@ func waitForDifferentCurrentControl(t *testing.T, manager *ControlManager, runID
 	return nil
 }

-func registerWorkConnAsCaller(svr *Service, workConn *msg.Conn, newMsg *msg.NewWorkConn) error {
-	err := svr.RegisterWorkConn(workConn, newMsg)
+func registerWorkConnAsCaller(
+	svr *Service,
+	workConn *msg.Conn,
+	newMsg *msg.NewWorkConn,
+	wireProtocol string,
+	clientHelloPresent bool,
+) error {
+	err := svr.RegisterWorkConn(workConn, newMsg, wireProtocol, clientHelloPresent)
 	if err != nil {
 		_ = workConn.Close()
 	}