stcp, xtcp, sudp: support allow_users and specified server user (#3472)

server/control.go

@@ -524,7 +524,7 @@ func (ctl *Control) manager() {
 }
 
 func (ctl *Control) HandleNatHoleVisitor(m *msg.NatHoleVisitor) {
-	ctl.rc.NatHoleController.HandleVisitor(m, ctl.msgTransporter)
+	ctl.rc.NatHoleController.HandleVisitor(m, ctl.msgTransporter, ctl.loginMsg.User)
 }
 
 func (ctl *Control) HandleNatHoleClient(m *msg.NatHoleClient) {
@@ -537,7 +537,7 @@ func (ctl *Control) HandleNatHoleReport(m *msg.NatHoleReport) {
 
 func (ctl *Control) RegisterProxy(pxyMsg *msg.NewProxy) (remoteAddr string, err error) {
 	var pxyConf config.ProxyConf
-	// Load configures from NewProxy message and check.
+	// Load configures from NewProxy message and validate.
 	pxyConf, err = config.NewProxyConfFromMsg(pxyMsg, ctl.serverCfg)
 	if err != nil {
 		return
@@ -550,8 +550,8 @@ func (ctl *Control) RegisterProxy(pxyMsg *msg.NewProxy) (remoteAddr string, err
 		RunID: ctl.runID,
 	}
 
-	// NewProxy will return a interface Proxy.
-	// In fact it create different proxies by different proxy type, we just call run() here.
+	// NewProxy will return an interface Proxy.
+	// In fact, it creates different proxies based on the proxy type. We just call run() here.
 	pxy, err := proxy.NewProxy(ctl.ctx, userInfo, ctl.rc, ctl.poolCount, ctl.GetWorkConn, pxyConf, ctl.serverCfg, ctl.loginMsg)
 	if err != nil {
 		return remoteAddr, err

server/proxy/stcp.go

@@ -27,7 +27,12 @@ type STCPProxy struct {
 
 func (pxy *STCPProxy) Run() (remoteAddr string, err error) {
 	xl := pxy.xl
-	listener, errRet := pxy.rc.VisitorManager.Listen(pxy.GetName(), pxy.cfg.Sk)
+	allowUsers := pxy.cfg.AllowUsers
+	// if allowUsers is empty, only allow same user from proxy
+	if len(allowUsers) == 0 {
+		allowUsers = []string{pxy.GetUserInfo().User}
+	}
+	listener, errRet := pxy.rc.VisitorManager.Listen(pxy.GetName(), pxy.cfg.Sk, allowUsers)
 	if errRet != nil {
 		err = errRet
 		return

server/proxy/sudp.go

@@ -27,8 +27,12 @@ type SUDPProxy struct {
 
 func (pxy *SUDPProxy) Run() (remoteAddr string, err error) {
 	xl := pxy.xl
-
-	listener, errRet := pxy.rc.VisitorManager.Listen(pxy.GetName(), pxy.cfg.Sk)
+	allowUsers := pxy.cfg.AllowUsers
+	// if allowUsers is empty, only allow same user from proxy
+	if len(allowUsers) == 0 {
+		allowUsers = []string{pxy.GetUserInfo().User}
+	}
+	listener, errRet := pxy.rc.VisitorManager.Listen(pxy.GetName(), pxy.cfg.Sk, allowUsers)
 	if errRet != nil {
 		err = errRet
 		return

server/proxy/xtcp.go

@@ -35,11 +35,15 @@ func (pxy *XTCPProxy) Run() (remoteAddr string, err error) {
 	xl := pxy.xl
 
 	if pxy.rc.NatHoleController == nil {
-		xl.Error("udp port for xtcp is not specified.")
 		err = fmt.Errorf("xtcp is not supported in frps")
 		return
 	}
-	sidCh := pxy.rc.NatHoleController.ListenClient(pxy.GetName(), pxy.cfg.Sk)
+	allowUsers := pxy.cfg.AllowUsers
+	// if allowUsers is empty, only allow same user from proxy
+	if len(allowUsers) == 0 {
+		allowUsers = []string{pxy.GetUserInfo().User}
+	}
+	sidCh := pxy.rc.NatHoleController.ListenClient(pxy.GetName(), pxy.cfg.Sk, allowUsers)
 	go func() {
 		for {
 			select {

server/service.go

@@ -587,6 +587,15 @@ func (svr *Service) RegisterWorkConn(workConn net.Conn, newMsg *msg.NewWorkConn)
 }
 
 func (svr *Service) RegisterVisitorConn(visitorConn net.Conn, newMsg *msg.NewVisitorConn) error {
+	visitorUser := ""
+	// TODO: Compatible with old versions, can be without runID, user is empty. In later versions, it will be mandatory to include runID.
+	if newMsg.RunID != "" {
+		ctl, exist := svr.ctlManager.GetByID(newMsg.RunID)
+		if !exist {
+			return fmt.Errorf("no client control found for run id [%s]", newMsg.RunID)
+		}
+		visitorUser = ctl.loginMsg.User
+	}
 	return svr.rc.VisitorManager.NewConn(newMsg.ProxyName, visitorConn, newMsg.Timestamp, newMsg.SignKey,
-		newMsg.UseEncryption, newMsg.UseCompression)
+		newMsg.UseEncryption, newMsg.UseCompression, visitorUser)
 }

server/visitor/visitor.go

@@ -21,57 +21,69 @@ import (
 	"sync"
 
 	libio "github.com/fatedier/golib/io"
+	"github.com/samber/lo"
 
 	utilnet "github.com/fatedier/frp/pkg/util/net"
 	"github.com/fatedier/frp/pkg/util/util"
 )
 
+type listenerBundle struct {
+	l          *utilnet.InternalListener
+	sk         string
+	allowUsers []string
+}
+
 // Manager for visitor listeners.
 type Manager struct {
-	visitorListeners map[string]*utilnet.InternalListener
-	skMap            map[string]string
+	listeners map[string]*listenerBundle
 
 	mu sync.RWMutex
 }
 
 func NewManager() *Manager {
 	return &Manager{
-		visitorListeners: make(map[string]*utilnet.InternalListener),
-		skMap:            make(map[string]string),
+		listeners: make(map[string]*listenerBundle),
 	}
 }
 
-func (vm *Manager) Listen(name string, sk string) (l *utilnet.InternalListener, err error) {
+func (vm *Manager) Listen(name string, sk string, allowUsers []string) (l *utilnet.InternalListener, err error) {
 	vm.mu.Lock()
 	defer vm.mu.Unlock()
 
-	if _, ok := vm.visitorListeners[name]; ok {
+	if _, ok := vm.listeners[name]; ok {
 		err = fmt.Errorf("custom listener for [%s] is repeated", name)
 		return
 	}
 
 	l = utilnet.NewInternalListener()
-	vm.visitorListeners[name] = l
-	vm.skMap[name] = sk
+	vm.listeners[name] = &listenerBundle{
+		l:          l,
+		sk:         sk,
+		allowUsers: allowUsers,
+	}
 	return
 }
 
 func (vm *Manager) NewConn(name string, conn net.Conn, timestamp int64, signKey string,
-	useEncryption bool, useCompression bool,
+	useEncryption bool, useCompression bool, visitorUser string,
 ) (err error) {
 	vm.mu.RLock()
 	defer vm.mu.RUnlock()
 
-	if l, ok := vm.visitorListeners[name]; ok {
-		var sk string
-		if sk = vm.skMap[name]; util.GetAuthKey(sk, timestamp) != signKey {
+	if l, ok := vm.listeners[name]; ok {
+		if util.GetAuthKey(l.sk, timestamp) != signKey {
 			err = fmt.Errorf("visitor connection of [%s] auth failed", name)
 			return
 		}
 
+		if !lo.Contains(l.allowUsers, visitorUser) && !lo.Contains(l.allowUsers, "*") {
+			err = fmt.Errorf("visitor connection of [%s] user [%s] not allowed", name, visitorUser)
+			return
+		}
+
 		var rwc io.ReadWriteCloser = conn
 		if useEncryption {
-			if rwc, err = libio.WithEncryption(rwc, []byte(sk)); err != nil {
+			if rwc, err = libio.WithEncryption(rwc, []byte(l.sk)); err != nil {
 				err = fmt.Errorf("create encryption connection failed: %v", err)
 				return
 			}
@@ -79,7 +91,7 @@ func (vm *Manager) NewConn(name string, conn net.Conn, timestamp int64, signKey
 		if useCompression {
 			rwc = libio.WithCompression(rwc)
 		}
-		err = l.PutConn(utilnet.WrapReadWriteCloserToConn(rwc, conn))
+		err = l.l.PutConn(utilnet.WrapReadWriteCloserToConn(rwc, conn))
 	} else {
 		err = fmt.Errorf("custom listener for [%s] doesn't exist", name)
 		return
@@ -91,6 +103,5 @@ func (vm *Manager) CloseListener(name string) {
 	vm.mu.Lock()
 	defer vm.mu.Unlock()
 
-	delete(vm.visitorListeners, name)
-	delete(vm.skMap, name)
+	delete(vm.listeners, name)
 }