new proxy type: stcp(secret tcp)

server/manager.go

@@ -16,7 +16,12 @@ package server
 
 import (
 	"fmt"
+	"io"
 	"sync"
+
+	frpIo "github.com/fatedier/frp/utils/io"
+	frpNet "github.com/fatedier/frp/utils/net"
+	"github.com/fatedier/frp/utils/util"
 )
 
 type ControlManager struct {
@@ -87,3 +92,72 @@ func (pm *ProxyManager) GetByName(name string) (pxy Proxy, ok bool) {
 	pxy, ok = pm.pxys[name]
 	return
 }
+
+// Manager for vistor listeners.
+type VistorManager struct {
+	vistorListeners map[string]*frpNet.CustomListener
+	skMap           map[string]string
+
+	mu sync.RWMutex
+}
+
+func NewVistorManager() *VistorManager {
+	return &VistorManager{
+		vistorListeners: make(map[string]*frpNet.CustomListener),
+		skMap:           make(map[string]string),
+	}
+}
+
+func (vm *VistorManager) Listen(name string, sk string) (l *frpNet.CustomListener, err error) {
+	vm.mu.Lock()
+	defer vm.mu.Unlock()
+
+	if _, ok := vm.vistorListeners[name]; ok {
+		err = fmt.Errorf("custom listener for [%s] is repeated", name)
+		return
+	}
+
+	l = frpNet.NewCustomListener()
+	vm.vistorListeners[name] = l
+	vm.skMap[name] = sk
+	return
+}
+
+func (vm *VistorManager) NewConn(name string, conn frpNet.Conn, timestamp int64, signKey string,
+	useEncryption bool, useCompression bool) (err error) {
+
+	vm.mu.RLock()
+	defer vm.mu.RUnlock()
+
+	if l, ok := vm.vistorListeners[name]; ok {
+		var sk string
+		if sk = vm.skMap[name]; util.GetAuthKey(sk, timestamp) != signKey {
+			err = fmt.Errorf("vistor connection of [%s] auth failed", name)
+			return
+		}
+
+		var rwc io.ReadWriteCloser = conn
+		if useEncryption {
+			if rwc, err = frpIo.WithEncryption(rwc, []byte(sk)); err != nil {
+				err = fmt.Errorf("create encryption connection failed: %v", err)
+				return
+			}
+		}
+		if useCompression {
+			rwc = frpIo.WithCompression(rwc)
+		}
+		err = l.PutConn(frpNet.WrapReadWriteCloserToConn(rwc))
+	} else {
+		err = fmt.Errorf("custom listener for [%s] doesn't exist", name)
+		return
+	}
+	return
+}
+
+func (vm *VistorManager) CloseListener(name string) {
+	vm.mu.Lock()
+	defer vm.mu.Unlock()
+
+	delete(vm.vistorListeners, name)
+	delete(vm.skMap, name)
+}

server/proxy.go

@@ -143,6 +143,11 @@ func NewProxy(ctl *Control, pxyConf config.ProxyConf) (pxy Proxy, err error) {
 			BaseProxy: basePxy,
 			cfg:       cfg,
 		}
+	case *config.StcpProxyConf:
+		pxy = &StcpProxy{
+			BaseProxy: basePxy,
+			cfg:       cfg,
+		}
 	default:
 		return pxy, fmt.Errorf("proxy type not support")
 	}
@@ -274,6 +279,33 @@ func (pxy *HttpsProxy) Close() {
 	pxy.BaseProxy.Close()
 }
 
+type StcpProxy struct {
+	BaseProxy
+	cfg *config.StcpProxyConf
+}
+
+func (pxy *StcpProxy) Run() error {
+	listener, err := pxy.ctl.svr.vistorManager.Listen(pxy.GetName(), pxy.cfg.Sk)
+	if err != nil {
+		return err
+	}
+	listener.AddLogPrefix(pxy.name)
+	pxy.listeners = append(pxy.listeners, listener)
+	pxy.Info("stcp proxy custom listen success")
+
+	pxy.startListenHandler(pxy, HandleUserTcpConnection)
+	return nil
+}
+
+func (pxy *StcpProxy) GetConf() config.ProxyConf {
+	return pxy.cfg
+}
+
+func (pxy *StcpProxy) Close() {
+	pxy.BaseProxy.Close()
+	pxy.ctl.svr.vistorManager.CloseListener(pxy.GetName())
+}
+
 type UdpProxy struct {
 	BaseProxy
 	cfg *config.UdpProxyConf

server/service.go

@@ -55,12 +55,16 @@ type Service struct {
 
 	// Manage all proxies.
 	pxyManager *ProxyManager
+
+	// Manage all vistor listeners.
+	vistorManager *VistorManager
 }
 
 func NewService() (svr *Service, err error) {
 	svr = &Service{
-		ctlManager: NewControlManager(),
-		pxyManager: NewProxyManager(),
+		ctlManager:    NewControlManager(),
+		pxyManager:    NewProxyManager(),
+		vistorManager: NewVistorManager(),
 	}
 
 	// Init assets.
@@ -176,6 +180,20 @@ func (svr *Service) HandleListener(l frpNet.Listener) {
 					}
 				case *msg.NewWorkConn:
 					svr.RegisterWorkConn(conn, m)
+				case *msg.NewVistorConn:
+					if err = svr.RegisterVistorConn(conn, m); err != nil {
+						conn.Warn("%v", err)
+						msg.WriteMsg(conn, &msg.NewVistorConnResp{
+							ProxyName: m.ProxyName,
+							Error:     err.Error(),
+						})
+						conn.Close()
+					} else {
+						msg.WriteMsg(conn, &msg.NewVistorConnResp{
+							ProxyName: m.ProxyName,
+							Error:     "",
+						})
+					}
 				default:
 					log.Warn("Error message type for the new connection [%s]", conn.RemoteAddr().String())
 					conn.Close()
@@ -262,9 +280,13 @@ func (svr *Service) RegisterWorkConn(workConn frpNet.Conn, newMsg *msg.NewWorkCo
 	return
 }
 
+func (svr *Service) RegisterVistorConn(vistorConn frpNet.Conn, newMsg *msg.NewVistorConn) error {
+	return svr.vistorManager.NewConn(newMsg.ProxyName, vistorConn, newMsg.Timestamp, newMsg.SignKey,
+		newMsg.UseEncryption, newMsg.UseCompression)
+}
+
 func (svr *Service) RegisterProxy(name string, pxy Proxy) error {
-	err := svr.pxyManager.Add(name, pxy)
-	return err
+	return svr.pxyManager.Add(name, pxy)
 }
 
 func (svr *Service) DelProxy(name string) {