From 5eab21476889f02ccb6da9cedb02383906faf3a7 Mon Sep 17 00:00:00 2001
From: fatedier <fatedier@gmail.com>
Date: Fri, 6 Mar 2026 23:39:12 +0800
Subject: [PATCH] client/proxy: extract wrapWorkConn to deduplicate UDP/SUDP
 connection wrapping

Move the repeated rate-limiting, encryption, and compression wrapping
logic from UDPProxy and SUDPProxy into a shared BaseProxy.wrapWorkConn
method, reducing ~18 lines of duplication in each proxy type.
---
 client/proxy/proxy.go | 24 ++++++++++++++++++++++++
 client/proxy/sudp.go  | 24 +++---------------------
 client/proxy/udp.go   | 24 +++---------------------
 3 files changed, 30 insertions(+), 42 deletions(-)

diff --git a/client/proxy/proxy.go b/client/proxy/proxy.go
index ca7d905b..148e0e90 100644
--- a/client/proxy/proxy.go
+++ b/client/proxy/proxy.go
@@ -16,6 +16,7 @@ package proxy
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"net"
 	"reflect"
@@ -122,6 +123,29 @@ func (pxy *BaseProxy) Close() {
 	}
 }
 
+// wrapWorkConn applies rate limiting, encryption, and compression
+// to a work connection based on the proxy's transport configuration.
+func (pxy *BaseProxy) wrapWorkConn(conn net.Conn) (net.Conn, error) {
+	var rwc io.ReadWriteCloser = conn
+	if pxy.limiter != nil {
+		rwc = libio.WrapReadWriteCloser(limit.NewReader(conn, pxy.limiter), limit.NewWriter(conn, pxy.limiter), func() error {
+			return conn.Close()
+		})
+	}
+	if pxy.baseCfg.Transport.UseEncryption {
+		var err error
+		rwc, err = libio.WithEncryption(rwc, pxy.encryptionKey)
+		if err != nil {
+			conn.Close()
+			return nil, fmt.Errorf("create encryption stream error: %w", err)
+		}
+	}
+	if pxy.baseCfg.Transport.UseCompression {
+		rwc = libio.WithCompression(rwc)
+	}
+	return netpkg.WrapReadWriteCloserToConn(rwc, conn), nil
+}
+
 func (pxy *BaseProxy) SetInWorkConnCallback(cb func(*v1.ProxyBaseConfig, net.Conn, *msg.StartWorkConn) bool) {
 	pxy.inWorkConnCallback = cb
 }
diff --git a/client/proxy/sudp.go b/client/proxy/sudp.go
index 3a7af19c..8b8c4985 100644
--- a/client/proxy/sudp.go
+++ b/client/proxy/sudp.go
@@ -17,7 +17,6 @@
 package proxy
 
 import (
-	"io"
 	"net"
 	"reflect"
 	"strconv"
@@ -25,13 +24,10 @@ import (
 	"time"
 
 	"github.com/fatedier/golib/errors"
-	libio "github.com/fatedier/golib/io"
 
 	v1 "github.com/fatedier/frp/pkg/config/v1"
 	"github.com/fatedier/frp/pkg/msg"
 	"github.com/fatedier/frp/pkg/proto/udp"
-	"github.com/fatedier/frp/pkg/util/limit"
-	netpkg "github.com/fatedier/frp/pkg/util/net"
 )
 
 func init() {
@@ -83,25 +79,11 @@ func (pxy *SUDPProxy) InWorkConn(conn net.Conn, _ *msg.StartWorkConn) {
 	xl := pxy.xl
 	xl.Infof("incoming a new work connection for sudp proxy, %s", conn.RemoteAddr().String())
 
-	var rwc io.ReadWriteCloser = conn
 	var err error
-	if pxy.limiter != nil {
-		rwc = libio.WrapReadWriteCloser(limit.NewReader(conn, pxy.limiter), limit.NewWriter(conn, pxy.limiter), func() error {
-			return conn.Close()
-		})
+	if conn, err = pxy.wrapWorkConn(conn); err != nil {
+		xl.Errorf("wrap work conn error: %v", err)
+		return
 	}
-	if pxy.cfg.Transport.UseEncryption {
-		rwc, err = libio.WithEncryption(rwc, pxy.encryptionKey)
-		if err != nil {
-			conn.Close()
-			xl.Errorf("create encryption stream error: %v", err)
-			return
-		}
-	}
-	if pxy.cfg.Transport.UseCompression {
-		rwc = libio.WithCompression(rwc)
-	}
-	conn = netpkg.WrapReadWriteCloserToConn(rwc, conn)
 
 	workConn := conn
 	readCh := make(chan *msg.UDPPacket, 1024)
diff --git a/client/proxy/udp.go b/client/proxy/udp.go
index 68426dc6..8848abc8 100644
--- a/client/proxy/udp.go
+++ b/client/proxy/udp.go
@@ -17,20 +17,16 @@
 package proxy
 
 import (
-	"io"
 	"net"
 	"reflect"
 	"strconv"
 	"time"
 
 	"github.com/fatedier/golib/errors"
-	libio "github.com/fatedier/golib/io"
 
 	v1 "github.com/fatedier/frp/pkg/config/v1"
 	"github.com/fatedier/frp/pkg/msg"
 	"github.com/fatedier/frp/pkg/proto/udp"
-	"github.com/fatedier/frp/pkg/util/limit"
-	netpkg "github.com/fatedier/frp/pkg/util/net"
 )
 
 func init() {
@@ -94,25 +90,11 @@ func (pxy *UDPProxy) InWorkConn(conn net.Conn, _ *msg.StartWorkConn) {
 	// close resources related with old workConn
 	pxy.Close()
 
-	var rwc io.ReadWriteCloser = conn
 	var err error
-	if pxy.limiter != nil {
-		rwc = libio.WrapReadWriteCloser(limit.NewReader(conn, pxy.limiter), limit.NewWriter(conn, pxy.limiter), func() error {
-			return conn.Close()
-		})
+	if conn, err = pxy.wrapWorkConn(conn); err != nil {
+		xl.Errorf("wrap work conn error: %v", err)
+		return
 	}
-	if pxy.cfg.Transport.UseEncryption {
-		rwc, err = libio.WithEncryption(rwc, pxy.encryptionKey)
-		if err != nil {
-			conn.Close()
-			xl.Errorf("create encryption stream error: %v", err)
-			return
-		}
-	}
-	if pxy.cfg.Transport.UseCompression {
-		rwc = libio.WithCompression(rwc)
-	}
-	conn = netpkg.WrapReadWriteCloserToConn(rwc, conn)
 
 	pxy.mu.Lock()
 	pxy.workConn = conn