sshTunnelGateway refactor (#3784)

client/proxy/proxy.go

@@ -47,10 +47,9 @@ func RegisterProxyFactory(proxyConfType reflect.Type, factory func(*BaseProxy, v
 // Proxy defines how to handle work connections for different proxy type.
 type Proxy interface {
 	Run() error
-
 	// InWorkConn accept work connections registered to server.
 	InWorkConn(net.Conn, *msg.StartWorkConn)
-
+	SetInWorkConnCallback(func(*v1.ProxyBaseConfig, net.Conn, *msg.StartWorkConn) /* continue */ bool)
 	Close()
 }
 
@@ -89,7 +88,8 @@ type BaseProxy struct {
 	limiter        *rate.Limiter
 	// proxyPlugin is used to handle connections instead of dialing to local service.
 	// It's only validate for TCP protocol now.
-	proxyPlugin plugin.Plugin
+	proxyPlugin        plugin.Plugin
+	inWorkConnCallback func(*v1.ProxyBaseConfig, net.Conn, *msg.StartWorkConn) /* continue */ bool
 
 	mu  sync.RWMutex
 	xl  *xlog.Logger
@@ -113,7 +113,16 @@ func (pxy *BaseProxy) Close() {
 	}
 }
 
+func (pxy *BaseProxy) SetInWorkConnCallback(cb func(*v1.ProxyBaseConfig, net.Conn, *msg.StartWorkConn) bool) {
+	pxy.inWorkConnCallback = cb
+}
+
 func (pxy *BaseProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
+	if pxy.inWorkConnCallback != nil {
+		if !pxy.inWorkConnCallback(pxy.baseCfg, conn, m) {
+			return
+		}
+	}
 	pxy.HandleTCPWorkConnection(conn, m, []byte(pxy.clientCfg.Auth.Token))
 }
 
@@ -132,7 +141,7 @@ func (pxy *BaseProxy) HandleTCPWorkConnection(workConn net.Conn, m *msg.StartWor
 		})
 	}
 
-	xl.Trace("handle tcp work connection, use_encryption: %t, use_compression: %t",
+	xl.Trace("handle tcp work connection, useEncryption: %t, useCompression: %t",
 		baseCfg.Transport.UseEncryption, baseCfg.Transport.UseCompression)
 	if baseCfg.Transport.UseEncryption {
 		remote, err = libio.WithEncryption(remote, encKey)

client/proxy/proxy_manager.go

@@ -31,8 +31,9 @@ import (
 )
 
 type Manager struct {
-	proxies        map[string]*Wrapper
-	msgTransporter transport.MessageTransporter
+	proxies            map[string]*Wrapper
+	msgTransporter     transport.MessageTransporter
+	inWorkConnCallback func(*v1.ProxyBaseConfig, net.Conn, *msg.StartWorkConn) bool
 
 	closed bool
 	mu     sync.RWMutex
@@ -71,6 +72,10 @@ func (pm *Manager) StartProxy(name string, remoteAddr string, serverRespErr stri
 	return nil
 }
 
+func (pm *Manager) SetInWorkConnCallback(cb func(*v1.ProxyBaseConfig, net.Conn, *msg.StartWorkConn) bool) {
+	pm.inWorkConnCallback = cb
+}
+
 func (pm *Manager) Close() {
 	pm.mu.Lock()
 	defer pm.mu.Unlock()
@@ -146,6 +151,9 @@ func (pm *Manager) Reload(pxyCfgs []v1.ProxyConfigurer) {
 		name := cfg.GetBaseConfig().Name
 		if _, ok := pm.proxies[name]; !ok {
 			pxy := NewWrapper(pm.ctx, cfg, pm.clientCfg, pm.HandleEvent, pm.msgTransporter)
+			if pm.inWorkConnCallback != nil {
+				pxy.SetInWorkConnCallback(pm.inWorkConnCallback)
+			}
 			pm.proxies[name] = pxy
 			addPxyNames = append(addPxyNames, name)
 

client/proxy/proxy_wrapper.go

@@ -121,6 +121,10 @@ func NewWrapper(
 	return pw
 }
 
+func (pw *Wrapper) SetInWorkConnCallback(cb func(*v1.ProxyBaseConfig, net.Conn, *msg.StartWorkConn) bool) {
+	pw.pxy.SetInWorkConnCallback(cb)
+}
+
 func (pw *Wrapper) SetRunningStatus(remoteAddr string, respErr string) error {
 	pw.mu.Lock()
 	defer pw.mu.Unlock()

client/proxy/sudp.go

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+//go:build !frps
+
 package proxy
 
 import (

client/proxy/udp.go

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+//go:build !frps
+
 package proxy
 
 import (

client/proxy/xtcp.go

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+//go:build !frps
+
 package proxy
 
 import (