From 6da093a4028eb5b6ef1d5cc6d0f98f047557d8a1 Mon Sep 17 00:00:00 2001 From: fatedier Date: Sun, 3 Nov 2019 01:20:49 +0800 Subject: [PATCH] support bandwith limit for one proxy --- client/proxy/proxy.go | 42 ++++++++++++----- conf/frpc_full.ini | 2 + go.mod | 1 + go.sum | 2 + models/config/proxy.go | 17 +++++-- models/config/types.go | 100 +++++++++++++++++++++++++++++++++++++++++ utils/limit/reader.go | 51 +++++++++++++++++++++ utils/limit/writer.go | 60 +++++++++++++++++++++++++ 8 files changed, 262 insertions(+), 13 deletions(-) create mode 100644 models/config/types.go create mode 100644 utils/limit/reader.go create mode 100644 utils/limit/writer.go diff --git a/client/proxy/proxy.go b/client/proxy/proxy.go index 2da68ce..aca1b94 100644 --- a/client/proxy/proxy.go +++ b/client/proxy/proxy.go @@ -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) diff --git a/conf/frpc_full.ini b/conf/frpc_full.ini index 158ff23..afa0707 100644 --- a/conf/frpc_full.ini +++ b/conf/frpc_full.ini @@ -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 diff --git a/go.mod b/go.mod index 81de8ab..a71a97c 100644 --- a/go.mod +++ b/go.mod @@ -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 ) diff --git a/go.sum b/go.sum index b9bdf2a..26c9004 100644 --- a/go.sum +++ b/go.sum @@ -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= diff --git a/models/config/proxy.go b/models/config/proxy.go index 9ab1ef8..2b129cf 100644 --- a/models/config/proxy.go +++ b/models/config/proxy.go @@ -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.LocalSvrConf.UnmarshalFromIni(prefix, name, section); err != nil { return err } - if err := cfg.HealthCheckConf.UnmarshalFromIni(prefix, name, section); err != nil { + if err = cfg.HealthCheckConf.UnmarshalFromIni(prefix, name, section); err != nil { return err } diff --git a/models/config/types.go b/models/config/types.go new file mode 100644 index 0000000..9f3de66 --- /dev/null +++ b/models/config/types.go @@ -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 +} diff --git a/utils/limit/reader.go b/utils/limit/reader.go new file mode 100644 index 0000000..efa828f --- /dev/null +++ b/utils/limit/reader.go @@ -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 +} diff --git a/utils/limit/writer.go b/utils/limit/writer.go new file mode 100644 index 0000000..5256d1e --- /dev/null +++ b/utils/limit/writer.go @@ -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 +} -- GitLab