guoswork/frp
1
0
Fork 0
mirror of https://gitee.com/IrisVega/frp.git synced 2024-11-01 22:31:29 +08:00
Code Issues Packages Projects Releases Wiki Activity

support bandwith limit for one proxy

Browse Source
This commit is contained in:
fatedier 2019-11-03 01:20:49 +08:00
parent 22a79710d8
commit 6da093a402
8 changed files with 262 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

42
client/proxy/proxy.go
View File

@ -30,6 +30,7 @@ import (
"github.com/fatedier/frp/models/msg"
"github.com/fatedier/frp/models/plugin"
"github.com/fatedier/frp/models/proto/udp"
"github.com/fatedier/frp/utils/limit"
frpNet "github.com/fatedier/frp/utils/net"
"github.com/fatedier/frp/utils/xlog"
@ -38,6 +39,7 @@ import (
"github.com/fatedier/golib/pool"
fmux "github.com/hashicorp/yamux"
pp "github.com/pires/go-proxyproto"
"golang.org/x/time/rate"
)
// Proxy defines how to handle work connections for different proxy type.
@ -51,9 +53,16 @@ type Proxy interface {
}
func NewProxy(ctx context.Context, pxyConf config.ProxyConf, clientCfg config.ClientCommonConf, serverUDPPort int) (pxy Proxy) {
var limiter *rate.Limiter
limitBytes := pxyConf.GetBaseInfo().BandwithLimit.Bytes()
if limitBytes > 0 {
limiter = rate.NewLimiter(rate.Limit(float64(limitBytes)), int(limitBytes))
}
baseProxy := BaseProxy{
clientCfg: clientCfg,
serverUDPPort: serverUDPPort,
limiter: limiter,
xl: xlog.FromContextSafe(ctx),
ctx: ctx,
}
@ -96,6 +105,7 @@ type BaseProxy struct {
closed bool
clientCfg config.ClientCommonConf
serverUDPPort int
limiter *rate.Limiter
mu sync.RWMutex
xl *xlog.Logger
@ -127,8 +137,8 @@ func (pxy *TcpProxy) Close() {
}
func (pxy *TcpProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
HandleTcpWorkConnection(pxy.ctx, &pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf, conn,
[]byte(pxy.clientCfg.Token), m)
HandleTcpWorkConnection(pxy.ctx, &pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf, pxy.limiter,
conn, []byte(pxy.clientCfg.Token), m)
}
// HTTP
@ -156,8 +166,8 @@ func (pxy *HttpProxy) Close() {
}
func (pxy *HttpProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
HandleTcpWorkConnection(pxy.ctx, &pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf, conn,
[]byte(pxy.clientCfg.Token), m)
HandleTcpWorkConnection(pxy.ctx, &pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf, pxy.limiter,
conn, []byte(pxy.clientCfg.Token), m)
}
// HTTPS
@ -185,8 +195,8 @@ func (pxy *HttpsProxy) Close() {
}
func (pxy *HttpsProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
HandleTcpWorkConnection(pxy.ctx, &pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf, conn,
[]byte(pxy.clientCfg.Token), m)
HandleTcpWorkConnection(pxy.ctx, &pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf, pxy.limiter,
conn, []byte(pxy.clientCfg.Token), m)
}
// STCP
@ -214,8 +224,8 @@ func (pxy *StcpProxy) Close() {
}
func (pxy *StcpProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
HandleTcpWorkConnection(pxy.ctx, &pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf, conn,
[]byte(pxy.clientCfg.Token), m)
HandleTcpWorkConnection(pxy.ctx, &pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf, pxy.limiter,
conn, []byte(pxy.clientCfg.Token), m)
}
// XTCP
@ -360,7 +370,7 @@ func (pxy *XtcpProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
return
}
HandleTcpWorkConnection(pxy.ctx, &pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf,
HandleTcpWorkConnection(pxy.ctx, &pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf, pxy.limiter,
muxConn, []byte(pxy.cfg.Sk), m)
}
@ -429,6 +439,13 @@ func (pxy *UdpProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
// close resources releated with old workConn
pxy.Close()
if pxy.limiter != nil {
rwc := frpIo.WrapReadWriteCloser(limit.NewReader(conn, pxy.limiter), limit.NewWriter(conn, pxy.limiter), func() error {
return conn.Close()
})
conn = frpNet.WrapReadWriteCloserToConn(rwc, conn)
}
pxy.mu.Lock()
pxy.workConn = conn
pxy.readCh = make(chan *msg.UdpPacket, 1024)
@ -491,13 +508,18 @@ func (pxy *UdpProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
// Common handler for tcp work connections.
func HandleTcpWorkConnection(ctx context.Context, localInfo *config.LocalSvrConf, proxyPlugin plugin.Plugin,
baseInfo *config.BaseProxyConf, workConn net.Conn, encKey []byte, m *msg.StartWorkConn) {
baseInfo *config.BaseProxyConf, limiter *rate.Limiter, workConn net.Conn, encKey []byte, m *msg.StartWorkConn) {
xl := xlog.FromContextSafe(ctx)
var (
remote io.ReadWriteCloser
err error
)
remote = workConn
if limiter != nil {
remote = frpIo.WrapReadWriteCloser(limit.NewReader(workConn, limiter), limit.NewWriter(workConn, limiter), func() error {
return workConn.Close()
})
}
xl.Trace("handle tcp work connection, use_encryption: %t, use_compression: %t",
baseInfo.UseEncryption, baseInfo.UseCompression)

2
conf/frpc_full.ini
View File

@ -71,6 +71,8 @@ tls_enable = true
type = tcp
local_ip = 127.0.0.1
local_port = 22
# limit bandwith for this proxy, unit is KB and MB
bandwith_limit = 1MB
# true or false, if true, messages between frps and frpc will be encrypted, default is false
use_encryption = false
# if true, message will be compressed

1
go.mod
View File

@ -29,4 +29,5 @@ require (
github.com/xtaci/lossyconn v0.0.0-20190602105132-8df528c0c9ae // indirect
golang.org/x/net v0.0.0-20190724013045-ca1201d0de80
golang.org/x/text v0.3.2 // indirect
golang.org/x/time v0.0.0-20191024005414-555d28b269f0
)

2
go.sum
View File

@ -46,4 +46,6 @@ golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5h
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
golang.org/x/text v0.3.2 h1:tW2bmiBqwgJj/UpqtC8EpXEZVYOwU0yG4iWbprSVAcs=
golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
golang.org/x/time v0.0.0-20191024005414-555d28b269f0 h1:/5xXl8Y5W96D+TtHSlonuFqGHIWVuyCkGJLwGh9JJFs=
golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=

17
models/config/proxy.go
View File

@ -125,6 +125,11 @@ type BaseProxyConf struct {
// values include "v1", "v2", and "". If the value is "", a protocol
// version will be automatically selected. By default, this value is "".
ProxyProtocolVersion string `json:"proxy_protocol_version"`
// BandwithLimit limit the proxy bandwith
// 0 means no limit
BandwithLimit BandwithQuantity `json:"bandwith_limit"`
LocalSvrConf
HealthCheckConf
}
@ -140,7 +145,8 @@ func (cfg *BaseProxyConf) compare(cmp *BaseProxyConf) bool {
cfg.UseCompression != cmp.UseCompression ||
cfg.Group != cmp.Group ||
cfg.GroupKey != cmp.GroupKey ||
cfg.ProxyProtocolVersion != cmp.ProxyProtocolVersion {
cfg.ProxyProtocolVersion != cmp.ProxyProtocolVersion ||
cfg.BandwithLimit.Equal(&cmp.BandwithLimit) {
return false
}
if !cfg.LocalSvrConf.compare(&cmp.LocalSvrConf) {
@ -165,6 +171,7 @@ func (cfg *BaseProxyConf) UnmarshalFromIni(prefix string, name string, section i
var (
tmpStr string
ok bool
err error
)
cfg.ProxyName = prefix + name
cfg.ProxyType = section["type"]
@ -183,11 +190,15 @@ func (cfg *BaseProxyConf) UnmarshalFromIni(prefix string, name string, section i
cfg.GroupKey = section["group_key"]
cfg.ProxyProtocolVersion = section["proxy_protocol_version"]
if err := cfg.LocalSvrConf.UnmarshalFromIni(prefix, name, section); err != nil {
if cfg.BandwithLimit, err = NewBandwithQuantity(section["bandwidth_limit"]); err != nil {
return err
}
if err := cfg.HealthCheckConf.UnmarshalFromIni(prefix, name, section); err != nil {
if err = cfg.LocalSvrConf.UnmarshalFromIni(prefix, name, section); err != nil {
return err
}
if err = cfg.HealthCheckConf.UnmarshalFromIni(prefix, name, section); err != nil {
return err
}

100
models/config/types.go Normal file
View File

@ -0,0 +1,100 @@
// Copyright 2019 fatedier, fatedier@gmail.com
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package config
import (
"errors"
"strconv"
"strings"
)
const (
MB = 1024 * 1024
KB = 1024
)
type BandwithQuantity struct {
s string // MB or KB
i int64 // bytes
}
func NewBandwithQuantity(s string) (BandwithQuantity, error) {
q := BandwithQuantity{}
err := q.UnmarshalString(s)
if err != nil {
return q, err
}
return q, nil
}
func (q *BandwithQuantity) Equal(u *BandwithQuantity) bool {
if q == nil && u == nil {
return true
}
if q != nil && u != nil {
return q.i == u.i
}
return false
}
func (q *BandwithQuantity) String() string {
return q.s
}
func (q *BandwithQuantity) UnmarshalString(s string) error {
q.s = strings.TrimSpace(s)
if q.s == "" {
return nil
}
var (
base int64
f float64
err error
)
if strings.HasSuffix(s, "MB") {
base = MB
s = strings.TrimSuffix(s, "MB")
f, err = strconv.ParseFloat(s, 64)
if err != nil {
return err
}
} else if strings.HasSuffix(s, "KB") {
base = KB
s = strings.TrimSuffix(s, "KB")
f, err = strconv.ParseFloat(s, 64)
if err != nil {
return err
}
} else {
return errors.New("unit not support")
}
q.i = int64(f * float64(base))
return nil
}
func (q *BandwithQuantity) UnmarshalJSON(b []byte) error {
return q.UnmarshalString(string(b))
}
func (q *BandwithQuantity) MarshalJSON() ([]byte, error) {
return []byte(q.s), nil
}
func (q *BandwithQuantity) Bytes() int64 {
return q.i
}

51
utils/limit/reader.go Normal file
View File

@ -0,0 +1,51 @@
// Copyright 2019 fatedier, fatedier@gmail.com
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package limit
import (
"context"
"io"
"golang.org/x/time/rate"
)
type Reader struct {
r io.Reader
limiter *rate.Limiter
}
func NewReader(r io.Reader, limiter *rate.Limiter) *Reader {
return &Reader{
r: r,
limiter: limiter,
}
}
func (r *Reader) Read(p []byte) (n int, err error) {
b := r.limiter.Burst()
if b < len(p) {
p = p[:b]
}
n, err = r.r.Read(p)
if err != nil {
return
}
err = r.limiter.WaitN(context.Background(), n)
if err != nil {
return
}
return
}

60
utils/limit/writer.go Normal file
View File

@ -0,0 +1,60 @@
// Copyright 2019 fatedier, fatedier@gmail.com
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package limit
import (
"context"
"io"
"golang.org/x/time/rate"
)
type Writer struct {
w io.Writer
limiter *rate.Limiter
}
func NewWriter(w io.Writer, limiter *rate.Limiter) *Writer {
return &Writer{
w: w,
limiter: limiter,
}
}
func (w *Writer) Write(p []byte) (n int, err error) {
var nn int
b := w.limiter.Burst()
for {
end := len(p)
if end == 0 {
break
}
if b < len(p) {
end = b
}
err = w.limiter.WaitN(context.Background(), end)
if err != nil {
return
}
nn, err = w.w.Write(p[:end])
n += nn
if err != nil {
return
}
p = p[end:]
}
return
}