Merge pull request #5340 from fatedier/design/wire-v2-workconn-message-framing

feat: use wire v2 framing for UDP workConn payload

server/proxy/proxy_test.go

@@ -15,12 +15,15 @@
 package proxy
 
 import (
+	"context"
 	"net"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 
+	v1 "github.com/fatedier/frp/pkg/config/v1"
 	"github.com/fatedier/frp/pkg/msg"
+	"github.com/fatedier/frp/pkg/proto/wire"
 )
 
 func TestWorkConnStartWritesStartWorkConn(t *testing.T) {
@@ -51,3 +54,56 @@ func TestWorkConnStartWritesStartWorkConn(t *testing.T) {
 	require.NoError(t, result.err)
 	require.Same(t, serverMsgConn, result.conn)
 }
+
+func TestGetWorkConnFromPoolStartWorkConnUnchangedForUDPWireV2(t *testing.T) {
+	startMsg := getStartWorkConnFromPool(t, &v1.UDPProxyConfig{
+		ProxyBaseConfig: v1.ProxyBaseConfig{Name: "udp", Type: string(v1.ProxyTypeUDP)},
+	}, wire.ProtocolV2)
+
+	require.Equal(t, msg.StartWorkConn{ProxyName: "udp"}, startMsg)
+}
+
+func TestGetWorkConnFromPoolLeavesRawTCPPayloadUnframed(t *testing.T) {
+	startMsg := getStartWorkConnFromPool(t, &v1.TCPProxyConfig{
+		ProxyBaseConfig: v1.ProxyBaseConfig{Name: "tcp", Type: string(v1.ProxyTypeTCP)},
+	}, wire.ProtocolV2)
+
+	require.Equal(t, msg.StartWorkConn{ProxyName: "tcp"}, startMsg)
+}
+
+func getStartWorkConnFromPool(t *testing.T, cfg v1.ProxyConfigurer, wireProtocol string) msg.StartWorkConn {
+	t.Helper()
+
+	client, server := net.Pipe()
+	t.Cleanup(func() {
+		client.Close()
+		server.Close()
+	})
+
+	serverMsgConn := msg.NewConn(server, msg.NewV2ReadWriter(server))
+	clientMsgConn := msg.NewConn(client, msg.NewV2ReadWriter(client))
+	pxy := &BaseProxy{
+		name:         cfg.GetBaseConfig().Name,
+		configurer:   cfg,
+		poolCount:    0,
+		ctx:          context.Background(),
+		wireProtocol: wireProtocol,
+		getWorkConnFn: func() (*WorkConn, error) {
+			return NewWorkConn(serverMsgConn), nil
+		},
+	}
+
+	errCh := make(chan error, 1)
+	go func() {
+		conn, err := pxy.GetWorkConnFromPool(nil, nil)
+		if conn != nil {
+			conn.Close()
+		}
+		errCh <- err
+	}()
+
+	var startMsg msg.StartWorkConn
+	require.NoError(t, clientMsgConn.ReadMsgInto(&startMsg))
+	require.NoError(t, <-errCh)
+	return startMsg
+}