From cbe23d2ed197a02f247f62725c62a26200acbb37 Mon Sep 17 00:00:00 2001
From: yusing <bookz2i2@duck.com>
Date: Sun, 31 Mar 2024 07:04:08 +0000
Subject: [PATCH] tcp/udp fix

---
 src/go-proxy/stream_route.go | 13 +++++++++++--
 src/go-proxy/tcp_route.go    | 14 +++-----------
 src/go-proxy/udp_route.go    | 16 +++-------------
 3 files changed, 17 insertions(+), 26 deletions(-)

diff --git a/src/go-proxy/stream_route.go b/src/go-proxy/stream_route.go
index 0ffa05f..fbe2990 100755
--- a/src/go-proxy/stream_route.go
+++ b/src/go-proxy/stream_route.go
@@ -83,6 +83,10 @@ func newStreamRouteBase(config *ProxyConfig) (*StreamRouteBase, error) {
 		dstScheme = config.Scheme
 	}
 
+	if srcScheme != dstScheme {
+		return nil, NewNestedError("unsupported").Subjectf("%v -> %v", srcScheme, dstScheme)
+	}
+
 	return &StreamRouteBase{
 		Alias:           config.Alias,
 		Type:            streamType,
@@ -106,14 +110,19 @@ func newStreamRouteBase(config *ProxyConfig) (*StreamRouteBase, error) {
 }
 
 func NewStreamRoute(config *ProxyConfig) (StreamRoute, error) {
+	base, err := newStreamRouteBase(config)
+	if err != nil {
+		return nil, err
+	}
 	switch config.Scheme {
 	case StreamType_TCP:
-		return NewTCPRoute(config)
+		base.StreamImpl = NewTCPRoute(base)
 	case StreamType_UDP:
-		return NewUDPRoute(config)
+		base.StreamImpl = NewUDPRoute(base)
 	default:
 		return nil, NewNestedError("invalid stream type").Subject(config.Scheme)
 	}
+	return base, nil
 }
 
 func (route *StreamRouteBase) ListeningUrl() string {
diff --git a/src/go-proxy/tcp_route.go b/src/go-proxy/tcp_route.go
index 61967b0..a6d6762 100755
--- a/src/go-proxy/tcp_route.go
+++ b/src/go-proxy/tcp_route.go
@@ -16,19 +16,11 @@ type TCPRoute struct {
 	listener net.Listener
 }
 
-func NewTCPRoute(config *ProxyConfig) (StreamRoute, error) {
-	base, err := newStreamRouteBase(config)
-	if err != nil {
-		return nil, NewNestedErrorFrom(err).Subject(config.Alias)
-	}
-	if base.TargetScheme != StreamType_TCP {
-		return nil, NewNestedError("unsupported").Subjectf("tcp -> %s", base.TargetScheme)
-	}
-	base.StreamImpl = &TCPRoute{
+func NewTCPRoute(base *StreamRouteBase) StreamImpl {
+	return &TCPRoute{
 		StreamRouteBase: base,
 		listener:        nil,
 	}
-	return base, nil
 }
 
 func (route *TCPRoute) Setup() error {
@@ -44,7 +36,7 @@ func (route *TCPRoute) Accept() (interface{}, error) {
 	return route.listener.Accept()
 }
 
-func (route *TCPRoute) HandleConnection(c interface{}) error {
+func (route *TCPRoute) Handle(c interface{}) error {
 	clientConn := c.(net.Conn)
 
 	defer clientConn.Close()
diff --git a/src/go-proxy/udp_route.go b/src/go-proxy/udp_route.go
index 1f79a7c..fc274ad 100755
--- a/src/go-proxy/udp_route.go
+++ b/src/go-proxy/udp_route.go
@@ -26,21 +26,11 @@ type UDPConn struct {
 	nReceived     int
 }
 
-func NewUDPRoute(config *ProxyConfig) (StreamRoute, error) {
-	base, err := newStreamRouteBase(config)
-	if err != nil {
-		return nil, err
-	}
-
-	if base.TargetScheme != StreamType_UDP {
-		return nil, NewNestedError("unsupported").Subjectf("udp->%s", base.TargetScheme)
-	}
-
-	base.StreamImpl = &UDPRoute{
+func NewUDPRoute(base *StreamRouteBase) StreamImpl {
+	return &UDPRoute{
 		StreamRouteBase: base,
 		connMap:         make(map[net.Addr]net.Conn),
 	}
-	return base, nil
 }
 
 func (route *UDPRoute) Setup() error {
@@ -83,7 +73,7 @@ func (route *UDPRoute) Accept() (interface{}, error) {
 	return conn, nil
 }
 
-func (route *UDPRoute) HandleConnection(c interface{}) error {
+func (route *UDPRoute) Handle(c interface{}) error {
 	var err error
 
 	conn := c.(*UDPConn)