From 08ea4ef3044c87bf06f22cce1b5f582ba044d410 Mon Sep 17 00:00:00 2001
From: unknown <jimmyhuang454@gmail.com>
Date: Thu, 9 May 2024 17:13:02 +0800
Subject: [PATCH] fix hysteria2OnTCP and use ListenSystemPacket

---
 proxy/hysteria2/client.go              |  3 ++-
 testing/scenarios/hy2_test.go          |  8 ++++----
 transport/internet/hysteria2/dialer.go | 22 ++++++++++++++++++++++
 3 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/proxy/hysteria2/client.go b/proxy/hysteria2/client.go
index 35fe4d767d..6f897152d9 100644
--- a/proxy/hysteria2/client.go
+++ b/proxy/hysteria2/client.go
@@ -81,7 +81,8 @@ func (c *Client) Process(ctx context.Context, link *transport.Link, dialer inter
 	}
 	hyConn, IsHy2Transport := iConn.(*hy2_transport.HyConn)
 
-	if !IsHy2Transport && !hyConn.IsUDPExtension { // is not hysteria2 and proxing UDP
+	if !IsHy2Transport && network == net.Network_UDP {
+		// hysteria2 need to use udp extension to proxy UDP.
 		return newError(hy2_transport.CanNotUseUdpExtension)
 	}
 
diff --git a/testing/scenarios/hy2_test.go b/testing/scenarios/hy2_test.go
index 324937dd72..1e7f5244ca 100644
--- a/testing/scenarios/hy2_test.go
+++ b/testing/scenarios/hy2_test.go
@@ -276,9 +276,9 @@ func testHysteria2Offical(t *testing.T, isUDP bool) {
 									Security: &protocol.SecurityConfig{
 										Type: protocol.SecurityType_NONE,
 									},
-									Congestion: &hy2_transport.Congestion{Type: "brutal", UpMbps: 100, DownMbps: 100},
-									UseUdpExtension:        true,
-									Password:   "password",
+									Congestion:      &hy2_transport.Congestion{Type: "brutal", UpMbps: 100, DownMbps: 100},
+									UseUdpExtension: true,
+									Password:        "password",
 								}),
 							},
 						},
@@ -321,7 +321,7 @@ func testHysteria2Offical(t *testing.T, isUDP bool) {
 	}
 }
 
-func TestHysteria2TCP(t *testing.T) {
+func TestHysteria2OnTCP(t *testing.T) {
 	tcpServer := tcp.Server{
 		MsgProcessor: xor,
 	}
diff --git a/transport/internet/hysteria2/dialer.go b/transport/internet/hysteria2/dialer.go
index 2cd7cfd330..10bd67858b 100644
--- a/transport/internet/hysteria2/dialer.go
+++ b/transport/internet/hysteria2/dialer.go
@@ -52,6 +52,16 @@ func InitAddress(dest net.Destination) (net.Addr, error) {
 	return destAddr, nil
 }
 
+type connFactory struct {
+	hy_client.ConnFactory
+
+	NewFunc func(addr net.Addr) (net.PacketConn, error)
+}
+
+func (f *connFactory) New(addr net.Addr) (net.PacketConn, error) {
+	return f.NewFunc(addr)
+}
+
 func NewHyClient(dest net.Destination, streamSettings *internet.MemoryStreamConfig) (hy_client.Client, error) {
 	tlsConfig, err := InitTLSConifg(streamSettings)
 	if err != nil {
@@ -68,6 +78,18 @@ func NewHyClient(dest net.Destination, streamSettings *internet.MemoryStreamConf
 		TLSConfig:  *tlsConfig,
 		Auth:       config.GetPassword(),
 		ServerAddr: serverAddr,
+		ConnFactory: &connFactory{
+			NewFunc: func(addr net.Addr) (net.PacketConn, error) {
+				rawConn, err := internet.ListenSystemPacket(context.Background(), &net.UDPAddr{
+					IP:   []byte{0, 0, 0, 0},
+					Port: 0,
+				}, streamSettings.SocketSettings)
+				if err != nil {
+					return nil, err
+				}
+				return rawConn.(*net.UDPConn), nil
+			},
+		},
 	})
 	if err != nil {
 		return nil, err