support bandwidth_limit set by server plugin

This commit is contained in:
Craig O'Donnell 2023-01-25 23:42:59 -05:00
parent 804f2910fd
commit a4ffda814f
3 changed files with 17 additions and 1 deletions

View File

@ -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

View File

@ -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"`

View File

@ -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())