frps: new params max_ports_per_client

server/control.go

@@ -55,6 +55,9 @@ type Control struct {
 	// pool count
 	poolCount int
 
+	// ports used, for limitations
+	portsUsedNum int
+
 	// last time got the Ping message
 	lastPing time.Time
 
@@ -84,6 +87,7 @@ func NewControl(svr *Service, ctlConn net.Conn, loginMsg *msg.Login) *Control {
 		workConnCh:      make(chan net.Conn, loginMsg.PoolCount+10),
 		proxies:         make(map[string]Proxy),
 		poolCount:       loginMsg.PoolCount,
+		portsUsedNum:    0,
 		lastPing:        time.Now(),
 		runId:           loginMsg.RunId,
 		status:          consts.Working,
@@ -348,6 +352,26 @@ func (ctl *Control) RegisterProxy(pxyMsg *msg.NewProxy) (remoteAddr string, err
 		return remoteAddr, err
 	}
 
+	// Check ports used number in each client
+	if config.ServerCommonCfg.MaxPortsPerClient > 0 {
+		ctl.mu.Lock()
+		if ctl.portsUsedNum+pxy.GetUsedPortsNum() > int(config.ServerCommonCfg.MaxPortsPerClient) {
+			ctl.mu.Unlock()
+			err = fmt.Errorf("exceed the max_ports_per_client")
+			return
+		}
+		ctl.portsUsedNum = ctl.portsUsedNum + pxy.GetUsedPortsNum()
+		ctl.mu.Unlock()
+
+		defer func() {
+			if err != nil {
+				ctl.mu.Lock()
+				ctl.portsUsedNum = ctl.portsUsedNum - pxy.GetUsedPortsNum()
+				ctl.mu.Unlock()
+			}
+		}()
+	}
+
 	remoteAddr, err = pxy.Run()
 	if err != nil {
 		return
@@ -371,16 +395,21 @@ func (ctl *Control) RegisterProxy(pxyMsg *msg.NewProxy) (remoteAddr string, err
 
 func (ctl *Control) CloseProxy(closeMsg *msg.CloseProxy) (err error) {
 	ctl.mu.Lock()
-	defer ctl.mu.Unlock()
 
 	pxy, ok := ctl.proxies[closeMsg.ProxyName]
 	if !ok {
+		ctl.mu.Unlock()
 		return
 	}
 
+	if config.ServerCommonCfg.MaxPortsPerClient > 0 {
+		ctl.portsUsedNum = ctl.portsUsedNum - pxy.GetUsedPortsNum()
+	}
 	pxy.Close()
 	ctl.svr.DelProxy(pxy.GetName())
 	delete(ctl.proxies, closeMsg.ProxyName)
+	ctl.mu.Unlock()
+
 	StatsCloseProxy(pxy.GetName(), pxy.GetConf().GetBaseInfo().ProxyType)
 	return
 }

server/proxy.go

@@ -40,15 +40,18 @@ type Proxy interface {
 	GetName() string
 	GetConf() config.ProxyConf
 	GetWorkConnFromPool() (workConn frpNet.Conn, err error)
+	GetUsedPortsNum() int
 	Close()
 	log.Logger
 }
 
 type BaseProxy struct {
-	name      string
-	ctl       *Control
-	listeners []frpNet.Listener
-	mu        sync.RWMutex
+	name         string
+	ctl          *Control
+	listeners    []frpNet.Listener
+	usedPortsNum int
+
+	mu sync.RWMutex
 	log.Logger
 }
 
@@ -60,6 +63,10 @@ func (pxy *BaseProxy) GetControl() *Control {
 	return pxy.ctl
 }
 
+func (pxy *BaseProxy) GetUsedPortsNum() int {
+	return pxy.usedPortsNum
+}
+
 func (pxy *BaseProxy) Close() {
 	pxy.Info("proxy closing")
 	for _, l := range pxy.listeners {
@@ -126,6 +133,7 @@ func NewProxy(ctl *Control, pxyConf config.ProxyConf) (pxy Proxy, err error) {
 	}
 	switch cfg := pxyConf.(type) {
 	case *config.TcpProxyConf:
+		basePxy.usedPortsNum = 1
 		pxy = &TcpProxy{
 			BaseProxy: basePxy,
 			cfg:       cfg,
@@ -141,6 +149,7 @@ func NewProxy(ctl *Control, pxyConf config.ProxyConf) (pxy Proxy, err error) {
 			cfg:       cfg,
 		}
 	case *config.UdpProxyConf:
+		basePxy.usedPortsNum = 1
 		pxy = &UdpProxy{
 			BaseProxy: basePxy,
 			cfg:       cfg,