From a4ffda814fdad4cf7419e40e2a7f37036b6c98b9 Mon Sep 17 00:00:00 2001
From: Craig O'Donnell <cbodonnell@comcast.net>
Date: Wed, 25 Jan 2023 23:42:59 -0500
Subject: [PATCH] support bandwidth_limit set by server plugin

---
 pkg/config/proxy.go   |  2 ++
 pkg/msg/msg.go        |  5 ++++-
 server/proxy/proxy.go | 11 +++++++++++
 3 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/pkg/config/proxy.go b/pkg/config/proxy.go
index 92b499f2..1c9d2ad4 100644
--- a/pkg/config/proxy.go
+++ b/pkg/config/proxy.go
@@ -389,6 +389,7 @@ 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.Group = cfg.Group
 	pMsg.GroupKey = cfg.GroupKey
 	pMsg.Metas = cfg.Metas
@@ -399,6 +400,7 @@ 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.Group = pMsg.Group
 	cfg.GroupKey = pMsg.GroupKey
 	cfg.Metas = pMsg.Metas
diff --git a/pkg/msg/msg.go b/pkg/msg/msg.go
index 4b2823c7..194f0423 100644
--- a/pkg/msg/msg.go
+++ b/pkg/msg/msg.go
@@ -14,7 +14,9 @@
 
 package msg
 
-import "net"
+import (
+	"net"
+)
 
 const (
 	TypeLogin                 = 'o'
@@ -87,6 +89,7 @@ type NewProxy struct {
 	ProxyType      string            `json:"proxy_type,omitempty"`
 	UseEncryption  bool              `json:"use_encryption,omitempty"`
 	UseCompression bool              `json:"use_compression,omitempty"`
+	BandwidthLimit string            `json:"bandwidth_limit,omitempty"`
 	Group          string            `json:"group,omitempty"`
 	GroupKey       string            `json:"group_key,omitempty"`
 	Metas          map[string]string `json:"metas,omitempty"`
diff --git a/server/proxy/proxy.go b/server/proxy/proxy.go
index 808d9316..2bc81df6 100644
--- a/server/proxy/proxy.go
+++ b/server/proxy/proxy.go
@@ -24,10 +24,12 @@ import (
 	"time"
 
 	frpIo "github.com/fatedier/golib/io"
+	"golang.org/x/time/rate"
 
 	"github.com/fatedier/frp/pkg/config"
 	"github.com/fatedier/frp/pkg/msg"
 	plugin "github.com/fatedier/frp/pkg/plugin/server"
+	"github.com/fatedier/frp/pkg/util/limit"
 	frpNet "github.com/fatedier/frp/pkg/util/net"
 	"github.com/fatedier/frp/pkg/util/xlog"
 	"github.com/fatedier/frp/server/controller"
@@ -287,6 +289,15 @@ func HandleUserTCPConnection(pxy Proxy, userConn net.Conn, serverCfg config.Serv
 	if cfg.UseCompression {
 		local = frpIo.WithCompression(local)
 	}
+
+	limitBytes := cfg.BandwidthLimit.Bytes()
+	if limitBytes > 0 {
+		limiter := rate.NewLimiter(rate.Limit(float64(limitBytes)), int(limitBytes))
+		local = frpIo.WrapReadWriteCloser(limit.NewReader(local, limiter), limit.NewWriter(local, limiter), func() error {
+			return local.Close()
+		})
+	}
+
 	xl.Debug("join connections, workConn(l[%s] r[%s]) userConn(l[%s] r[%s])", workConn.LocalAddr().String(),
 		workConn.RemoteAddr().String(), userConn.LocalAddr().String(), userConn.RemoteAddr().String())