support bandwidth_limit set by server plugin (#3271)

* support bandwidth_limit set by server plugin

* limiter at proxy level

* bandwidth_limit_mode

* updates tests for bandwidth_limit_mode default

* bandwidth_limit_mode as string

* add checkForSrv for bandwidth_limit_mode

* bandwidth_limit flags for sub cmds

* gci write

pkg/config/client_test.go

@@ -74,6 +74,7 @@ var testClientBytesWithFull = []byte(`
 		local_ip = 127.0.0.9
 		local_port = 29
 		bandwidth_limit = 19MB
+		bandwidth_limit_mode = server
 		use_encryption
 		use_compression
 		remote_port = 6009
@@ -309,13 +310,14 @@ func Test_LoadClientBasicConf(t *testing.T) {
 	proxyExpected := map[string]ProxyConf{
 		testUser + ".ssh": &TCPProxyConf{
 			BaseProxyConf: BaseProxyConf{
-				ProxyName:      testUser + ".ssh",
-				ProxyType:      consts.TCPProxy,
-				UseCompression: true,
-				UseEncryption:  true,
-				Group:          "test_group",
-				GroupKey:       "123456",
-				BandwidthLimit: MustBandwidthQuantity("19MB"),
+				ProxyName:          testUser + ".ssh",
+				ProxyType:          consts.TCPProxy,
+				UseCompression:     true,
+				UseEncryption:      true,
+				Group:              "test_group",
+				GroupKey:           "123456",
+				BandwidthLimit:     MustBandwidthQuantity("19MB"),
+				BandwidthLimitMode: BandwidthLimitModeServer,
 				Metas: map[string]string{
 					"var1": "123",
 					"var2": "234",
@@ -342,6 +344,7 @@ func Test_LoadClientBasicConf(t *testing.T) {
 					LocalIP:   "127.0.0.9",
 					LocalPort: 29,
 				},
+				BandwidthLimitMode: BandwidthLimitModeClient,
 			},
 			RemotePort: 9,
 		},
@@ -353,6 +356,7 @@ func Test_LoadClientBasicConf(t *testing.T) {
 					LocalIP:   "127.0.0.9",
 					LocalPort: 6010,
 				},
+				BandwidthLimitMode: BandwidthLimitModeClient,
 			},
 			RemotePort: 6010,
 		},
@@ -364,6 +368,7 @@ func Test_LoadClientBasicConf(t *testing.T) {
 					LocalIP:   "127.0.0.9",
 					LocalPort: 6011,
 				},
+				BandwidthLimitMode: BandwidthLimitModeClient,
 			},
 			RemotePort: 6011,
 		},
@@ -375,6 +380,7 @@ func Test_LoadClientBasicConf(t *testing.T) {
 					LocalIP:   "127.0.0.9",
 					LocalPort: 6019,
 				},
+				BandwidthLimitMode: BandwidthLimitModeClient,
 			},
 			RemotePort: 6019,
 		},
@@ -388,6 +394,7 @@ func Test_LoadClientBasicConf(t *testing.T) {
 					LocalIP:   "114.114.114.114",
 					LocalPort: 59,
 				},
+				BandwidthLimitMode: BandwidthLimitModeClient,
 			},
 			RemotePort: 6009,
 		},
@@ -401,6 +408,7 @@ func Test_LoadClientBasicConf(t *testing.T) {
 					LocalIP:   "114.114.114.114",
 					LocalPort: 6000,
 				},
+				BandwidthLimitMode: BandwidthLimitModeClient,
 			},
 			RemotePort: 6000,
 		},
@@ -414,6 +422,7 @@ func Test_LoadClientBasicConf(t *testing.T) {
 					LocalIP:   "114.114.114.114",
 					LocalPort: 6010,
 				},
+				BandwidthLimitMode: BandwidthLimitModeClient,
 			},
 			RemotePort: 6010,
 		},
@@ -427,6 +436,7 @@ func Test_LoadClientBasicConf(t *testing.T) {
 					LocalIP:   "114.114.114.114",
 					LocalPort: 6011,
 				},
+				BandwidthLimitMode: BandwidthLimitModeClient,
 			},
 			RemotePort: 6011,
 		},
@@ -447,6 +457,7 @@ func Test_LoadClientBasicConf(t *testing.T) {
 					HealthCheckIntervalS: 19,
 					HealthCheckURL:       "http://127.0.0.9:89/status",
 				},
+				BandwidthLimitMode: BandwidthLimitModeClient,
 			},
 			DomainConf: DomainConf{
 				CustomDomains: []string{"web02.yourdomain.com"},
@@ -471,6 +482,7 @@ func Test_LoadClientBasicConf(t *testing.T) {
 					LocalPort: 8009,
 				},
 				ProxyProtocolVersion: "v2",
+				BandwidthLimitMode:   BandwidthLimitModeClient,
 			},
 			DomainConf: DomainConf{
 				CustomDomains: []string{"web02.yourdomain.com"},
@@ -485,6 +497,7 @@ func Test_LoadClientBasicConf(t *testing.T) {
 					LocalIP:   "127.0.0.1",
 					LocalPort: 22,
 				},
+				BandwidthLimitMode: BandwidthLimitModeClient,
 			},
 			Role: "server",
 			Sk:   "abcdefg",
@@ -497,6 +510,7 @@ func Test_LoadClientBasicConf(t *testing.T) {
 					LocalIP:   "127.0.0.1",
 					LocalPort: 22,
 				},
+				BandwidthLimitMode: BandwidthLimitModeClient,
 			},
 			Role: "server",
 			Sk:   "abcdefg",
@@ -509,6 +523,7 @@ func Test_LoadClientBasicConf(t *testing.T) {
 					LocalIP:   "127.0.0.1",
 					LocalPort: 10701,
 				},
+				BandwidthLimitMode: BandwidthLimitModeClient,
 			},
 			DomainConf: DomainConf{
 				CustomDomains: []string{"tunnel1"},
@@ -527,6 +542,7 @@ func Test_LoadClientBasicConf(t *testing.T) {
 						"plugin_unix_path": "/var/run/docker.sock",
 					},
 				},
+				BandwidthLimitMode: BandwidthLimitModeClient,
 			},
 			RemotePort: 6003,
 		},
@@ -542,6 +558,7 @@ func Test_LoadClientBasicConf(t *testing.T) {
 						"plugin_http_passwd": "abc",
 					},
 				},
+				BandwidthLimitMode: BandwidthLimitModeClient,
 			},
 			RemotePort: 6004,
 		},
@@ -557,6 +574,7 @@ func Test_LoadClientBasicConf(t *testing.T) {
 						"plugin_passwd": "abc",
 					},
 				},
+				BandwidthLimitMode: BandwidthLimitModeClient,
 			},
 			RemotePort: 6005,
 		},
@@ -574,6 +592,7 @@ func Test_LoadClientBasicConf(t *testing.T) {
 						"plugin_http_passwd":  "abc",
 					},
 				},
+				BandwidthLimitMode: BandwidthLimitModeClient,
 			},
 			RemotePort: 6006,
 		},
@@ -592,6 +611,7 @@ func Test_LoadClientBasicConf(t *testing.T) {
 						"plugin_header_X-From-Where": "frp",
 					},
 				},
+				BandwidthLimitMode: BandwidthLimitModeClient,
 			},
 			DomainConf: DomainConf{
 				CustomDomains: []string{"test.yourdomain.com"},
@@ -610,6 +630,7 @@ func Test_LoadClientBasicConf(t *testing.T) {
 						"plugin_header_X-From-Where": "frp",
 					},
 				},
+				BandwidthLimitMode: BandwidthLimitModeClient,
 			},
 			DomainConf: DomainConf{
 				CustomDomains: []string{"test.yourdomain.com"},

pkg/config/proxy.go

@@ -141,6 +141,10 @@ type BaseProxyConf struct {
 	// BandwidthLimit limit the bandwidth
 	// 0 means no limit
 	BandwidthLimit BandwidthQuantity `ini:"bandwidth_limit" json:"bandwidth_limit"`
+	// BandwidthLimitMode specifies whether to limit the bandwidth on the
+	// client or server side. Valid values include "client" and "server".
+	// By default, this value is "client".
+	BandwidthLimitMode string `ini:"bandwidth_limit_mode" json:"bandwidth_limit_mode"`
 
 	// meta info for each proxy
 	Metas map[string]string `ini:"-" json:"metas"`
@@ -319,6 +323,7 @@ func defaultBaseProxyConf(proxyType string) BaseProxyConf {
 		LocalSvrConf: LocalSvrConf{
 			LocalIP: "127.0.0.1",
 		},
+		BandwidthLimitMode: BandwidthLimitModeClient,
 	}
 }
 
@@ -335,6 +340,7 @@ func (cfg *BaseProxyConf) compare(cmp *BaseProxyConf) bool {
 		cfg.GroupKey != cmp.GroupKey ||
 		cfg.ProxyProtocolVersion != cmp.ProxyProtocolVersion ||
 		!cfg.BandwidthLimit.Equal(&cmp.BandwidthLimit) ||
+		cfg.BandwidthLimitMode != cmp.BandwidthLimitMode ||
 		!reflect.DeepEqual(cfg.Metas, cmp.Metas) {
 		return false
 	}
@@ -389,6 +395,8 @@ func (cfg *BaseProxyConf) marshalToMsg(pMsg *msg.NewProxy) {
 	pMsg.ProxyType = cfg.ProxyType
 	pMsg.UseEncryption = cfg.UseEncryption
 	pMsg.UseCompression = cfg.UseCompression
+	pMsg.BandwidthLimit = cfg.BandwidthLimit.String()
+	pMsg.BandwidthLimitMode = cfg.BandwidthLimitMode
 	pMsg.Group = cfg.Group
 	pMsg.GroupKey = cfg.GroupKey
 	pMsg.Metas = cfg.Metas
@@ -399,6 +407,8 @@ func (cfg *BaseProxyConf) unmarshalFromMsg(pMsg *msg.NewProxy) {
 	cfg.ProxyType = pMsg.ProxyType
 	cfg.UseEncryption = pMsg.UseEncryption
 	cfg.UseCompression = pMsg.UseCompression
+	cfg.BandwidthLimit, _ = NewBandwidthQuantity(pMsg.BandwidthLimit)
+	cfg.BandwidthLimitMode = pMsg.BandwidthLimitMode
 	cfg.Group = pMsg.Group
 	cfg.GroupKey = pMsg.GroupKey
 	cfg.Metas = pMsg.Metas
@@ -411,6 +421,10 @@ func (cfg *BaseProxyConf) checkForCli() (err error) {
 		}
 	}
 
+	if cfg.BandwidthLimitMode != "client" && cfg.BandwidthLimitMode != "server" {
+		return fmt.Errorf("bandwidth_limit_mode should be client or server")
+	}
+
 	if err = cfg.LocalSvrConf.checkForCli(); err != nil {
 		return
 	}
@@ -420,6 +434,13 @@ func (cfg *BaseProxyConf) checkForCli() (err error) {
 	return nil
 }
 
+func (cfg *BaseProxyConf) checkForSvr() (err error) {
+	if cfg.BandwidthLimitMode != "client" && cfg.BandwidthLimitMode != "server" {
+		return fmt.Errorf("bandwidth_limit_mode should be client or server")
+	}
+	return nil
+}
+
 // DomainConf
 func (cfg *DomainConf) check() (err error) {
 	if len(cfg.CustomDomains) == 0 && cfg.SubDomain == "" {
@@ -557,6 +578,9 @@ func (cfg *TCPProxyConf) CheckForCli() (err error) {
 }
 
 func (cfg *TCPProxyConf) CheckForSvr(serverCfg ServerCommonConf) error {
+	if err := cfg.BaseProxyConf.checkForSvr(); err != nil {
+		return err
+	}
 	return nil
 }
 
@@ -632,6 +656,10 @@ func (cfg *TCPMuxProxyConf) CheckForCli() (err error) {
 }
 
 func (cfg *TCPMuxProxyConf) CheckForSvr(serverCfg ServerCommonConf) (err error) {
+	if err := cfg.BaseProxyConf.checkForSvr(); err != nil {
+		return err
+	}
+
 	if cfg.Multiplexer != consts.HTTPConnectTCPMultiplexer {
 		return fmt.Errorf("proxy [%s] incorrect multiplexer [%s]", cfg.ProxyName, cfg.Multiplexer)
 	}
@@ -703,6 +731,9 @@ func (cfg *UDPProxyConf) CheckForCli() (err error) {
 }
 
 func (cfg *UDPProxyConf) CheckForSvr(serverCfg ServerCommonConf) error {
+	if err := cfg.BaseProxyConf.checkForSvr(); err != nil {
+		return err
+	}
 	return nil
 }
 
@@ -788,6 +819,10 @@ func (cfg *HTTPProxyConf) CheckForCli() (err error) {
 }
 
 func (cfg *HTTPProxyConf) CheckForSvr(serverCfg ServerCommonConf) (err error) {
+	if err := cfg.BaseProxyConf.checkForSvr(); err != nil {
+		return err
+	}
+
 	if serverCfg.VhostHTTPPort == 0 {
 		return fmt.Errorf("type [http] not support when vhost_http_port is not set")
 	}
@@ -860,6 +895,10 @@ func (cfg *HTTPSProxyConf) CheckForCli() (err error) {
 }
 
 func (cfg *HTTPSProxyConf) CheckForSvr(serverCfg ServerCommonConf) (err error) {
+	if err := cfg.BaseProxyConf.checkForSvr(); err != nil {
+		return err
+	}
+
 	if serverCfg.VhostHTTPSPort == 0 {
 		return fmt.Errorf("type [https] not support when vhost_https_port is not set")
 	}
@@ -932,6 +971,9 @@ func (cfg *SUDPProxyConf) CheckForCli() (err error) {
 }
 
 func (cfg *SUDPProxyConf) CheckForSvr(serverCfg ServerCommonConf) error {
+	if err := cfg.BaseProxyConf.checkForSvr(); err != nil {
+		return err
+	}
 	return nil
 }
 
@@ -998,6 +1040,9 @@ func (cfg *STCPProxyConf) CheckForCli() (err error) {
 }
 
 func (cfg *STCPProxyConf) CheckForSvr(serverCfg ServerCommonConf) error {
+	if err := cfg.BaseProxyConf.checkForSvr(); err != nil {
+		return err
+	}
 	return nil
 }
 
@@ -1064,5 +1109,8 @@ func (cfg *XTCPProxyConf) CheckForCli() (err error) {
 }
 
 func (cfg *XTCPProxyConf) CheckForSvr(serverCfg ServerCommonConf) error {
+	if err := cfg.BaseProxyConf.checkForSvr(); err != nil {
+		return err
+	}
 	return nil
 }

pkg/config/proxy_test.go

@@ -58,6 +58,7 @@ func Test_Proxy_UnmarshalFromIni(t *testing.T) {
 				local_ip = 127.0.0.9
 				local_port = 29
 				bandwidth_limit = 19MB
+				bandwidth_limit_mode = server
 				use_encryption
 				use_compression
 				remote_port = 6009
@@ -71,13 +72,14 @@ func Test_Proxy_UnmarshalFromIni(t *testing.T) {
 				meta_var2 = 234`),
 			expected: &TCPProxyConf{
 				BaseProxyConf: BaseProxyConf{
-					ProxyName:      testProxyPrefix + "ssh",
-					ProxyType:      consts.TCPProxy,
-					UseCompression: true,
-					UseEncryption:  true,
-					Group:          "test_group",
-					GroupKey:       "123456",
-					BandwidthLimit: MustBandwidthQuantity("19MB"),
+					ProxyName:          testProxyPrefix + "ssh",
+					ProxyType:          consts.TCPProxy,
+					UseCompression:     true,
+					UseEncryption:      true,
+					Group:              "test_group",
+					GroupKey:           "123456",
+					BandwidthLimit:     MustBandwidthQuantity("19MB"),
+					BandwidthLimitMode: BandwidthLimitModeServer,
 					Metas: map[string]string{
 						"var1": "123",
 						"var2": "234",
@@ -114,6 +116,7 @@ func Test_Proxy_UnmarshalFromIni(t *testing.T) {
 						LocalIP:   "127.0.0.9",
 						LocalPort: 29,
 					},
+					BandwidthLimitMode: BandwidthLimitModeClient,
 				},
 				RemotePort: 9,
 			},
@@ -139,6 +142,7 @@ func Test_Proxy_UnmarshalFromIni(t *testing.T) {
 						LocalIP:   "114.114.114.114",
 						LocalPort: 59,
 					},
+					BandwidthLimitMode: BandwidthLimitModeClient,
 				},
 				RemotePort: 6009,
 			},
@@ -182,6 +186,7 @@ func Test_Proxy_UnmarshalFromIni(t *testing.T) {
 						HealthCheckIntervalS: 19,
 						HealthCheckURL:       "http://127.0.0.9:89/status",
 					},
+					BandwidthLimitMode: BandwidthLimitModeClient,
 				},
 				DomainConf: DomainConf{
 					CustomDomains: []string{"web02.yourdomain.com"},
@@ -220,6 +225,7 @@ func Test_Proxy_UnmarshalFromIni(t *testing.T) {
 						LocalPort: 8009,
 					},
 					ProxyProtocolVersion: "v2",
+					BandwidthLimitMode:   BandwidthLimitModeClient,
 				},
 				DomainConf: DomainConf{
 					CustomDomains: []string{"web02.yourdomain.com"},
@@ -246,6 +252,7 @@ func Test_Proxy_UnmarshalFromIni(t *testing.T) {
 						LocalIP:   "127.0.0.1",
 						LocalPort: 22,
 					},
+					BandwidthLimitMode: BandwidthLimitModeClient,
 				},
 				Role: "server",
 				Sk:   "abcdefg",
@@ -270,6 +277,7 @@ func Test_Proxy_UnmarshalFromIni(t *testing.T) {
 						LocalIP:   "127.0.0.1",
 						LocalPort: 22,
 					},
+					BandwidthLimitMode: BandwidthLimitModeClient,
 				},
 				Role: "server",
 				Sk:   "abcdefg",
@@ -293,6 +301,7 @@ func Test_Proxy_UnmarshalFromIni(t *testing.T) {
 						LocalIP:   "127.0.0.1",
 						LocalPort: 10701,
 					},
+					BandwidthLimitMode: BandwidthLimitModeClient,
 				},
 				DomainConf: DomainConf{
 					CustomDomains: []string{"tunnel1"},
@@ -347,6 +356,7 @@ func Test_RangeProxy_UnmarshalFromIni(t *testing.T) {
 							LocalIP:   "127.0.0.9",
 							LocalPort: 6010,
 						},
+						BandwidthLimitMode: BandwidthLimitModeClient,
 					},
 					RemotePort: 6010,
 				},
@@ -358,6 +368,7 @@ func Test_RangeProxy_UnmarshalFromIni(t *testing.T) {
 							LocalIP:   "127.0.0.9",
 							LocalPort: 6011,
 						},
+						BandwidthLimitMode: BandwidthLimitModeClient,
 					},
 					RemotePort: 6011,
 				},
@@ -369,6 +380,7 @@ func Test_RangeProxy_UnmarshalFromIni(t *testing.T) {
 							LocalIP:   "127.0.0.9",
 							LocalPort: 6019,
 						},
+						BandwidthLimitMode: BandwidthLimitModeClient,
 					},
 					RemotePort: 6019,
 				},
@@ -396,6 +408,7 @@ func Test_RangeProxy_UnmarshalFromIni(t *testing.T) {
 							LocalIP:   "114.114.114.114",
 							LocalPort: 6000,
 						},
+						BandwidthLimitMode: BandwidthLimitModeClient,
 					},
 					RemotePort: 6000,
 				},
@@ -409,6 +422,7 @@ func Test_RangeProxy_UnmarshalFromIni(t *testing.T) {
 							LocalIP:   "114.114.114.114",
 							LocalPort: 6010,
 						},
+						BandwidthLimitMode: BandwidthLimitModeClient,
 					},
 					RemotePort: 6010,
 				},
@@ -422,6 +436,7 @@ func Test_RangeProxy_UnmarshalFromIni(t *testing.T) {
 							LocalIP:   "114.114.114.114",
 							LocalPort: 6011,
 						},
+						BandwidthLimitMode: BandwidthLimitModeClient,
 					},
 					RemotePort: 6011,
 				},

pkg/config/types.go

@@ -24,6 +24,9 @@ import (
 const (
 	MB = 1024 * 1024
 	KB = 1024
+
+	BandwidthLimitModeClient = "client"
+	BandwidthLimitModeServer = "server"
 )
 
 type BandwidthQuantity struct {