cyonjan/frp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

limiter at proxy level

Browse Source
This commit is contained in:
Craig O'Donnell 2023-01-30 23:15:09 -05:00
parent cb2e1cfffa
commit de7771da54
9 changed files with 67 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
server/proxy/http.go
View File

@ -20,8 +20,10 @@ import (
"strings"
frpIo "github.com/fatedier/golib/io"
"golang.org/x/time/rate"
"github.com/fatedier/frp/pkg/config"
"github.com/fatedier/frp/pkg/util/limit"
frpNet "github.com/fatedier/frp/pkg/util/net"
"github.com/fatedier/frp/pkg/util/util"
"github.com/fatedier/frp/pkg/util/vhost"
@ -135,6 +137,10 @@ func (pxy *HTTPProxy) GetConf() config.ProxyConf {
return pxy.cfg
}
func (pxy *HTTPProxy) GetLimiter() *rate.Limiter {
return pxy.limiter
}
func (pxy *HTTPProxy) GetRealConn(remoteAddr string) (workConn net.Conn, err error) {
xl := pxy.xl
rAddr, errRet := net.ResolveTCPAddr("tcp", remoteAddr)
@ -160,6 +166,13 @@ func (pxy *HTTPProxy) GetRealConn(remoteAddr string) (workConn net.Conn, err err
if pxy.cfg.UseCompression {
rwc = frpIo.WithCompression(rwc)
}
if pxy.GetLimiter() != nil {
rwc = frpIo.WrapReadWriteCloser(limit.NewReader(rwc, pxy.GetLimiter()), limit.NewWriter(rwc, pxy.GetLimiter()), func() error {
return rwc.Close()
})
}
workConn = frpNet.WrapReadWriteCloserToConn(rwc, tmpConn)
workConn = frpNet.WrapStatsConn(workConn, pxy.updateStatsAfterClosedConn)
metrics.Server.OpenConnection(pxy.GetName(), pxy.GetConf().GetBaseInfo().ProxyType)

5
server/proxy/https.go
View File

@ -20,6 +20,7 @@ import (
"github.com/fatedier/frp/pkg/config"
"github.com/fatedier/frp/pkg/util/util"
"github.com/fatedier/frp/pkg/util/vhost"
"golang.org/x/time/rate"
)
type HTTPSProxy struct {
@ -74,6 +75,10 @@ func (pxy *HTTPSProxy) GetConf() config.ProxyConf {
return pxy.cfg
}
func (pxy *HTTPSProxy) GetLimiter() *rate.Limiter {
return pxy.limiter
}
func (pxy *HTTPSProxy) Close() {
pxy.BaseProxy.Close()
}

16
server/proxy/proxy.go
View File

@ -47,6 +47,7 @@ type Proxy interface {
GetUsedPortsNum() int
GetResourceController() *controller.ResourceController
GetUserInfo() plugin.UserInfo
GetLimiter() *rate.Limiter
Close()
}
@ -58,6 +59,7 @@ type BaseProxy struct {
poolCount int
getWorkConnFn GetWorkConnFn
serverCfg config.ServerCommonConf
limiter *rate.Limiter
userInfo plugin.UserInfo
mu sync.RWMutex
@ -189,6 +191,13 @@ func NewProxy(ctx context.Context, userInfo plugin.UserInfo, rc *controller.Reso
getWorkConnFn GetWorkConnFn, pxyConf config.ProxyConf, serverCfg config.ServerCommonConf,
) (pxy Proxy, err error) {
xl := xlog.FromContextSafe(ctx).Spawn().AppendPrefix(pxyConf.GetBaseInfo().ProxyName)
var limiter *rate.Limiter
limitBytes := pxyConf.GetBaseInfo().BandwidthLimit.Bytes()
if limitBytes > 0 {
limiter = rate.NewLimiter(rate.Limit(float64(limitBytes)), int(limitBytes))
}
basePxy := BaseProxy{
name: pxyConf.GetBaseInfo().ProxyName,
rc: rc,
@ -196,6 +205,7 @@ func NewProxy(ctx context.Context, userInfo plugin.UserInfo, rc *controller.Reso
poolCount: poolCount,
getWorkConnFn: getWorkConnFn,
serverCfg: serverCfg,
limiter: limiter,
xl: xl,
ctx: xlog.NewContext(ctx, xl),
userInfo: userInfo,
@ -290,10 +300,8 @@ func HandleUserTCPConnection(pxy Proxy, userConn net.Conn, serverCfg config.Serv
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 {
if pxy.GetLimiter() != nil {
local = frpIo.WrapReadWriteCloser(limit.NewReader(local, pxy.GetLimiter()), limit.NewWriter(local, pxy.GetLimiter()), func() error {
return local.Close()
})
}

5
server/proxy/stcp.go
View File

@ -16,6 +16,7 @@ package proxy
import (
"github.com/fatedier/frp/pkg/config"
"golang.org/x/time/rate"
)
type STCPProxy struct {
@ -41,6 +42,10 @@ func (pxy *STCPProxy) GetConf() config.ProxyConf {
return pxy.cfg
}
func (pxy *STCPProxy) GetLimiter() *rate.Limiter {
return pxy.limiter
}
func (pxy *STCPProxy) Close() {
pxy.BaseProxy.Close()
pxy.rc.VisitorManager.CloseListener(pxy.GetName())

5
server/proxy/sudp.go
View File

@ -16,6 +16,7 @@ package proxy
import (
"github.com/fatedier/frp/pkg/config"
"golang.org/x/time/rate"
)
type SUDPProxy struct {
@ -42,6 +43,10 @@ func (pxy *SUDPProxy) GetConf() config.ProxyConf {
return pxy.cfg
}
func (pxy *SUDPProxy) GetLimiter() *rate.Limiter {
return pxy.limiter
}
func (pxy *SUDPProxy) Close() {
pxy.BaseProxy.Close()
pxy.rc.VisitorManager.CloseListener(pxy.GetName())

5
server/proxy/tcp.go
View File

@ -20,6 +20,7 @@ import (
"strconv"
"github.com/fatedier/frp/pkg/config"
"golang.org/x/time/rate"
)
type TCPProxy struct {
@ -74,6 +75,10 @@ func (pxy *TCPProxy) GetConf() config.ProxyConf {
return pxy.cfg
}
func (pxy *TCPProxy) GetLimiter() *rate.Limiter {
return pxy.limiter
}
func (pxy *TCPProxy) Close() {
pxy.BaseProxy.Close()
if pxy.cfg.Group == "" {

5
server/proxy/tcpmux.go
View File

@ -23,6 +23,7 @@ import (
"github.com/fatedier/frp/pkg/consts"
"github.com/fatedier/frp/pkg/util/util"
"github.com/fatedier/frp/pkg/util/vhost"
"golang.org/x/time/rate"
)
type TCPMuxProxy struct {
@ -94,6 +95,10 @@ func (pxy *TCPMuxProxy) GetConf() config.ProxyConf {
return pxy.cfg
}
func (pxy *TCPMuxProxy) GetLimiter() *rate.Limiter {
return pxy.limiter
}
func (pxy *TCPMuxProxy) Close() {
pxy.BaseProxy.Close()
}

12
server/proxy/udp.go
View File

@ -24,10 +24,12 @@ import (
"github.com/fatedier/golib/errors"
frpIo "github.com/fatedier/golib/io"
"golang.org/x/time/rate"
"github.com/fatedier/frp/pkg/config"
"github.com/fatedier/frp/pkg/msg"
"github.com/fatedier/frp/pkg/proto/udp"
"github.com/fatedier/frp/pkg/util/limit"
frpNet "github.com/fatedier/frp/pkg/util/net"
"github.com/fatedier/frp/server/metrics"
)
@ -198,6 +200,12 @@ func (pxy *UDPProxy) Run() (remoteAddr string, err error) {
rwc = frpIo.WithCompression(rwc)
}
if pxy.GetLimiter() != nil {
rwc = frpIo.WrapReadWriteCloser(limit.NewReader(rwc, pxy.GetLimiter()), limit.NewWriter(rwc, pxy.GetLimiter()), func() error {
return rwc.Close()
})
}
pxy.workConn = frpNet.WrapReadWriteCloserToConn(rwc, workConn)
ctx, cancel := context.WithCancel(context.Background())
go workConnReaderFn(pxy.workConn)
@ -225,6 +233,10 @@ func (pxy *UDPProxy) GetConf() config.ProxyConf {
return pxy.cfg
}
func (pxy *UDPProxy) GetLimiter() *rate.Limiter {
return pxy.limiter
}
func (pxy *UDPProxy) Close() {
pxy.mu.Lock()
defer pxy.mu.Unlock()

5
server/proxy/xtcp.go
View File

@ -18,6 +18,7 @@ import (
"fmt"
"github.com/fatedier/golib/errors"
"golang.org/x/time/rate"
"github.com/fatedier/frp/pkg/config"
"github.com/fatedier/frp/pkg/msg"
@ -88,6 +89,10 @@ func (pxy *XTCPProxy) GetConf() config.ProxyConf {
return pxy.cfg
}
func (pxy *XTCPProxy) GetLimiter() *rate.Limiter {
return pxy.limiter
}
func (pxy *XTCPProxy) Close() {
pxy.BaseProxy.Close()
pxy.rc.NatHoleController.CloseClient(pxy.GetName())