feat: add negotiated binary UDP packets (#5456)

This commit is contained in:
fatedier
2026-07-30 13:19:55 +08:00
committed by GitHub
parent 5c6d761c12
commit effa496859
23 changed files with 1232 additions and 71 deletions
+3 -1
View File
@@ -368,7 +368,8 @@ type SessionContext struct {
// server configuration
ServerCfg *v1.ServerConfig
// negotiated wire protocol for this client session
WireProtocol string
WireProtocol string
UDPPacketCodec string
}
type controlState uint8
@@ -821,6 +822,7 @@ func (ctl *Control) RegisterProxy(pxyMsg *msg.NewProxy) (remoteAddr string, err
ServerCfg: ctl.sessionCtx.ServerCfg,
EncryptionKey: ctl.sessionCtx.EncryptionKey,
WireProtocol: ctl.sessionCtx.WireProtocol,
UDPPacketCodec: ctl.sessionCtx.UDPPacketCodec,
})
if err != nil {
return remoteAddr, err
+30 -27
View File
@@ -82,19 +82,20 @@ type Proxy interface {
}
type BaseProxy struct {
name string
rc *controller.ResourceController
listeners []net.Listener
usedPortsNum int
poolCount int
getWorkConnFn GetWorkConnFn
serverCfg *v1.ServerConfig
encryptionKey []byte
limiter *rate.Limiter
userInfo plugin.UserInfo
loginMsg *msg.Login
configurer v1.ProxyConfigurer
wireProtocol string
name string
rc *controller.ResourceController
listeners []net.Listener
usedPortsNum int
poolCount int
getWorkConnFn GetWorkConnFn
serverCfg *v1.ServerConfig
encryptionKey []byte
limiter *rate.Limiter
userInfo plugin.UserInfo
loginMsg *msg.Login
configurer v1.ProxyConfigurer
wireProtocol string
udpPacketCodec string
mu sync.RWMutex
xl *xlog.Logger
@@ -469,6 +470,7 @@ type Options struct {
ServerCfg *v1.ServerConfig
EncryptionKey []byte
WireProtocol string
UDPPacketCodec string
}
func NewProxy(ctx context.Context, options *Options) (pxy Proxy, err error) {
@@ -482,20 +484,21 @@ func NewProxy(ctx context.Context, options *Options) (pxy Proxy, err error) {
}
basePxy := BaseProxy{
name: configurer.GetBaseConfig().Name,
rc: options.ResourceController,
listeners: make([]net.Listener, 0),
poolCount: options.PoolCount,
getWorkConnFn: options.GetWorkConnFn,
serverCfg: options.ServerCfg,
encryptionKey: options.EncryptionKey,
limiter: limiter,
xl: xl,
ctx: xlog.NewContext(ctx, xl),
userInfo: options.UserInfo,
loginMsg: options.LoginMsg,
configurer: configurer,
wireProtocol: options.WireProtocol,
name: configurer.GetBaseConfig().Name,
rc: options.ResourceController,
listeners: make([]net.Listener, 0),
poolCount: options.PoolCount,
getWorkConnFn: options.GetWorkConnFn,
serverCfg: options.ServerCfg,
encryptionKey: options.EncryptionKey,
limiter: limiter,
xl: xl,
ctx: xlog.NewContext(ctx, xl),
userInfo: options.UserInfo,
loginMsg: options.LoginMsg,
configurer: configurer,
wireProtocol: options.WireProtocol,
udpPacketCodec: options.UDPPacketCodec,
}
factory := proxyFactoryRegistry[reflect.TypeOf(configurer)]
+7 -1
View File
@@ -224,7 +224,13 @@ func (pxy *UDPProxy) Run() (remoteAddr string, err error) {
pxy.workConn = netpkg.WrapReadWriteCloserToConn(rwc, workConn)
// Plain UDP payload follows the negotiated wire protocol for message framing.
payloadConn := msg.NewConn(pxy.workConn, msg.NewReadWriter(pxy.workConn, pxy.wireProtocol))
payloadRW, err := msg.NewUDPPacketReadWriter(pxy.workConn, pxy.wireProtocol, pxy.udpPacketCodec)
if err != nil {
xl.Errorf("create UDP packet read writer: %v", err)
pxy.workConn.Close()
continue
}
payloadConn := msg.NewConn(pxy.workConn, payloadRW)
ctx, cancel := context.WithCancel(context.Background())
go workConnReaderFn(payloadConn)
go workConnSenderFn(payloadConn, ctx)
+50 -16
View File
@@ -470,7 +470,7 @@ func (svr *Service) handleConnection(ctx context.Context, conn net.Conn, interna
}
}
if err == nil {
ctl, err = svr.RegisterControl(controlConn, m, internal, acceptedConn.wireProtocol)
ctl, err = svr.RegisterControl(controlConn, m, internal, acceptedConn.wireProtocol, acceptedConn.udpPacketCodec)
}
}
@@ -509,7 +509,12 @@ func (svr *Service) handleConnection(ctx context.Context, conn net.Conn, interna
return
}
case *msg.NewWorkConn:
if err := svr.RegisterWorkConn(acceptedConn.conn, m); err != nil {
if err := svr.RegisterWorkConn(
acceptedConn.conn,
m,
acceptedConn.wireProtocol,
acceptedConn.clientHelloPresent,
); err != nil {
_ = acceptedConn.conn.WriteMsg(&msg.StartWorkConn{
Error: util.GenerateResponseErrorString("invalid NewWorkConn", err, lo.FromPtr(svr.cfg.DetailedErrorsToClient)),
})
@@ -547,10 +552,12 @@ func (svr *Service) completeControlLogin(ctl *Control, writeSuccess func() error
}
type acceptedConnection struct {
conn *msg.Conn
wireProtocol string
cryptoContext *wire.CryptoContext
firstMsg msg.Message
conn *msg.Conn
wireProtocol string
clientHelloPresent bool
udpPacketCodec string
cryptoContext *wire.CryptoContext
firstMsg msg.Message
}
func (svr *Service) acceptConnection(ctx context.Context, conn net.Conn) (*acceptedConnection, error) {
@@ -618,6 +625,7 @@ func (ac *acceptedConnection) readFirstV2Msg(conn net.Conn, wireConn *wire.Conn)
return nil, fmt.Errorf("read v2 frame: %w", err)
}
if frame.Type == wire.FrameTypeClientHello {
ac.clientHelloPresent = true
if err := ac.handleClientHello(conn, wireConn, frame); err != nil {
return nil, err
}
@@ -666,6 +674,7 @@ func (ac *acceptedConnection) handleClientHello(conn net.Conn, wireConn *wire.Co
return fmt.Errorf("write ServerHello: %w", err)
}
ac.cryptoContext = cryptoContext
ac.udpPacketCodec = serverHello.Selected.Message.UDPPacketCodec
return nil
}
@@ -759,7 +768,20 @@ func (svr *Service) RegisterControl(
loginMsg *msg.Login,
internal bool,
wireProtocol string,
udpPacketCodec string,
) (*Control, error) {
switch wireProtocol {
case wire.ProtocolV1:
if udpPacketCodec != "" {
return nil, fmt.Errorf("UDP packet codec %q requires wire protocol v2", udpPacketCodec)
}
case wire.ProtocolV2:
if udpPacketCodec != "" && udpPacketCodec != wire.UDPPacketCodecBinary {
return nil, fmt.Errorf("unsupported UDP packet codec selection: %s", udpPacketCodec)
}
default:
return nil, fmt.Errorf("unsupported wire protocol: %s", wireProtocol)
}
// If client's RunID is empty, it's a new client, we just create a new controller.
// Otherwise, we check if there is one controller has the same run id. If so, we release previous controller and start new one.
var err error
@@ -787,15 +809,16 @@ func (svr *Service) RegisterControl(
}
ctl, err := NewControl(ctx, &SessionContext{
RC: svr.rc,
PxyManager: svr.pxyManager,
PluginManager: svr.pluginManager,
AuthVerifier: authVerifier,
EncryptionKey: svr.auth.EncryptionKey(),
Conn: ctlConn,
LoginMsg: loginMsg,
ServerCfg: svr.cfg,
WireProtocol: wireProtocol,
RC: svr.rc,
PxyManager: svr.pxyManager,
PluginManager: svr.pluginManager,
AuthVerifier: authVerifier,
EncryptionKey: svr.auth.EncryptionKey(),
Conn: ctlConn,
LoginMsg: loginMsg,
ServerCfg: svr.cfg,
WireProtocol: wireProtocol,
UDPPacketCodec: udpPacketCodec,
})
if err != nil {
xl.Warnf("create new controller error: %v", err)
@@ -820,13 +843,24 @@ func (svr *Service) RegisterControl(
}
// RegisterWorkConn register a new work connection to control and proxies need it.
func (svr *Service) RegisterWorkConn(workConn *msg.Conn, newMsg *msg.NewWorkConn) error {
func (svr *Service) RegisterWorkConn(
workConn *msg.Conn,
newMsg *msg.NewWorkConn,
workWireProtocol string,
workClientHelloPresent bool,
) error {
if workClientHelloPresent {
return fmt.Errorf("ClientHello is not allowed on work connections")
}
xl := netpkg.NewLogFromConn(workConn)
ctl, exist := svr.ctlManager.GetByID(newMsg.RunID)
if !exist {
xl.Warnf("no client control found for run id [%s]", newMsg.RunID)
return fmt.Errorf("no client control found for run id [%s]", newMsg.RunID)
}
if workWireProtocol != ctl.sessionCtx.WireProtocol {
return fmt.Errorf("work connection wire protocol mismatch: got %s want %s", workWireProtocol, ctl.sessionCtx.WireProtocol)
}
// server plugin hook
content := &plugin.NewWorkConnContent{
+194 -7
View File
@@ -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()
}