go-gost/x
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add scope param for traffic limiter

Browse Source
This commit is contained in:
ginuerzh
2024-08-31 09:30:04 +08:00
parent 9405d1fb48
commit 31c482c649
57 changed files with 1239 additions and 485 deletions
Download Patch File Download Diff File Expand all files Collapse all files
limiter/traffic/plugin/grpc.go
+9 -4
View File
@@ -4,6 +4,7 @@ import (
"context"
"io"
"github.com/go-gost/core/limiter"
"github.com/go-gost/core/limiter/traffic"
"github.com/go-gost/core/logger"
"github.com/go-gost/plugin/limiter/traffic/proto"
@@ -44,18 +45,20 @@ func NewGRPCPlugin(name string, addr string, opts ...plugin.Option) traffic.Traf
return p
}
func (p *grpcPlugin) In(ctx context.Context, key string, opts ...traffic.Option) traffic.Limiter {
func (p *grpcPlugin) In(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
if p.client == nil {
return nil
}
var options traffic.Options
var options limiter.Options
for _, opt := range opts {
opt(&options)
}
r, err := p.client.Limit(ctx,
&proto.LimitRequest{
Service: options.Service,
Scope: options.Scope,
Network: options.Network,
Addr: options.Addr,
Client: options.Client,
@@ -69,18 +72,20 @@ func (p *grpcPlugin) In(ctx context.Context, key string, opts ...traffic.Option)
return xtraffic.NewLimiter(int(r.In))
}
func (p *grpcPlugin) Out(ctx context.Context, key string, opts ...traffic.Option) traffic.Limiter {
func (p *grpcPlugin) Out(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
if p.client == nil {
return nil
}
var options traffic.Options
var options limiter.Options
for _, opt := range opts {
opt(&options)
}
r, err := p.client.Limit(ctx,
&proto.LimitRequest{
Service: options.Service,
Scope: options.Scope,
Network: options.Network,
Addr: options.Addr,
Client: options.Client,
limiter/traffic/plugin/http.go
+21 -4
View File
@@ -6,6 +6,7 @@ import (
"encoding/json"
"net/http"
"github.com/go-gost/core/limiter"
"github.com/go-gost/core/limiter/traffic"
"github.com/go-gost/core/logger"
"github.com/go-gost/x/internal/plugin"
@@ -13,6 +14,8 @@ import (
)
type httpPluginRequest struct {
Service string `json:"service"`
Scope string `json:"scope"`
Network string `json:"network"`
Addr string `json:"addr"`
Client string `json:"client"`
@@ -49,17 +52,19 @@ func NewHTTPPlugin(name string, url string, opts ...plugin.Option) traffic.Traff
}
}
func (p *httpPlugin) In(ctx context.Context, key string, opts ...traffic.Option) traffic.Limiter {
func (p *httpPlugin) In(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
if p.client == nil {
return nil
}
var options traffic.Options
var options limiter.Options
for _, opt := range opts {
opt(&options)
}
rb := httpPluginRequest{
Service: options.Service,
Scope: options.Scope,
Network: options.Network,
Addr: options.Addr,
Client: options.Client,
@@ -67,11 +72,13 @@ func (p *httpPlugin) In(ctx context.Context, key string, opts ...traffic.Option)
}
v, err := json.Marshal(&rb)
if err != nil {
p.log.Error(err)
return nil
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.url, bytes.NewReader(v))
if err != nil {
p.log.Error(err)
return nil
}
@@ -81,32 +88,37 @@ func (p *httpPlugin) In(ctx context.Context, key string, opts ...traffic.Option)
req.Header.Set("Content-Type", "application/json")
resp, err := p.client.Do(req)
if err != nil {
p.log.Error(err)
return nil
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
p.log.Errorf("server return non-200 code: %s", resp.Status)
return nil
}
res := httpPluginResponse{}
if err := json.NewDecoder(resp.Body).Decode(&res); err != nil {
p.log.Error(err)
return nil
}
return xtraffic.NewLimiter(int(res.In))
}
func (p *httpPlugin) Out(ctx context.Context, key string, opts ...traffic.Option) traffic.Limiter {
func (p *httpPlugin) Out(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
if p.client == nil {
return nil
}
var options traffic.Options
var options limiter.Options
for _, opt := range opts {
opt(&options)
}
rb := httpPluginRequest{
Service: options.Service,
Scope: options.Scope,
Network: options.Network,
Addr: options.Addr,
Client: options.Client,
@@ -114,11 +126,13 @@ func (p *httpPlugin) Out(ctx context.Context, key string, opts ...traffic.Option
}
v, err := json.Marshal(&rb)
if err != nil {
p.log.Error(err)
return nil
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.url, bytes.NewReader(v))
if err != nil {
p.log.Error(err)
return nil
}
@@ -128,16 +142,19 @@ func (p *httpPlugin) Out(ctx context.Context, key string, opts ...traffic.Option
req.Header.Set("Content-Type", "application/json")
resp, err := p.client.Do(req)
if err != nil {
p.log.Error(err)
return nil
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
p.log.Errorf("server return non-200 code: %s", resp.Status)
return nil
}
res := httpPluginResponse{}
if err := json.NewDecoder(resp.Body).Decode(&res); err != nil {
p.log.Error(err)
return nil
}
return xtraffic.NewLimiter(int(res.Out))
limiter/traffic/traffic.go
+61 -30
View File
@@ -10,7 +10,8 @@ import (
"time"
"github.com/alecthomas/units"
limiter "github.com/go-gost/core/limiter/traffic"
"github.com/go-gost/core/limiter"
"github.com/go-gost/core/limiter/traffic"
"github.com/go-gost/core/logger"
"github.com/go-gost/x/internal/loader"
"github.com/patrickmn/go-cache"
@@ -93,7 +94,7 @@ type trafficLimiter struct {
options options
}
func NewTrafficLimiter(opts ...Option) limiter.TrafficLimiter {
func NewTrafficLimiter(opts ...Option) traffic.TrafficLimiter {
var options options
for _, opt := range opts {
opt(&options)
@@ -120,20 +121,35 @@ func NewTrafficLimiter(opts ...Option) limiter.TrafficLimiter {
}
// In obtains a traffic input limiter based on key.
// The key should be client connection address.
func (l *trafficLimiter) In(ctx context.Context, key string, opts ...limiter.Option) limiter.Limiter {
var lims []limiter.Limiter
// service level limiter
if lim, ok := l.inLimits.Get(GlobalLimitKey); ok && lim != nil {
lims = append(lims, lim.(limiter.Limiter))
// For connection scope, the key should be client connection address.
func (l *trafficLimiter) In(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
var options limiter.Options
for _, opt := range opts {
opt(&options)
}
switch options.Scope {
case limiter.ScopeService:
if lim, ok := l.inLimits.Get(GlobalLimitKey); ok && lim != nil {
return lim.(traffic.Limiter)
}
return nil
case limiter.ScopeClient:
return nil
case limiter.ScopeConn:
fallthrough
default:
}
var lims []traffic.Limiter
// connection level limiter
if lim, ok := l.connInLimits.Get(key); ok {
if lim != nil {
// cached connection level limiter
lims = append(lims, lim.(limiter.Limiter))
lims = append(lims, lim.(traffic.Limiter))
// reset expiration
l.connInLimits.Set(key, lim, defaultExpiration)
}
@@ -153,7 +169,7 @@ func (l *trafficLimiter) In(ctx context.Context, key string, opts ...limiter.Opt
if lim, ok := l.inLimits.Get(host); ok {
// cached IP limiter
if lim != nil {
lims = append(lims, lim.(limiter.Limiter))
lims = append(lims, lim.(traffic.Limiter))
}
} else {
l.mu.RLock()
@@ -171,7 +187,7 @@ func (l *trafficLimiter) In(ctx context.Context, key string, opts ...limiter.Opt
}
}
var lim limiter.Limiter
var lim traffic.Limiter
if len(lims) > 0 {
lim = newLimiterGroup(lims...)
}
@@ -184,20 +200,35 @@ func (l *trafficLimiter) In(ctx context.Context, key string, opts ...limiter.Opt
}
// Out obtains a traffic output limiter based on key.
// The key should be client connection address.
func (l *trafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Option) limiter.Limiter {
var lims []limiter.Limiter
// service level limiter
if lim, ok := l.outLimits.Get(GlobalLimitKey); ok && lim != nil {
lims = append(lims, lim.(limiter.Limiter))
// For connection scope, the key should be client connection address.
func (l *trafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
var options limiter.Options
for _, opt := range opts {
opt(&options)
}
switch options.Scope {
case limiter.ScopeService:
if lim, ok := l.outLimits.Get(GlobalLimitKey); ok && lim != nil {
return lim.(traffic.Limiter)
}
return nil
case limiter.ScopeClient:
return nil
case limiter.ScopeConn:
fallthrough
default:
}
var lims []traffic.Limiter
// connection level limiter
if lim, ok := l.connOutLimits.Get(key); ok {
if lim != nil {
// cached connection level limiter
lims = append(lims, lim.(limiter.Limiter))
lims = append(lims, lim.(traffic.Limiter))
// reset expiration
l.connOutLimits.Set(key, lim, defaultExpiration)
}
@@ -217,7 +248,7 @@ func (l *trafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Op
if lim, ok := l.outLimits.Get(host); ok {
if lim != nil {
// cached IP level limiter
lims = append(lims, lim.(limiter.Limiter))
lims = append(lims, lim.(traffic.Limiter))
}
} else {
l.mu.RLock()
@@ -235,7 +266,7 @@ func (l *trafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Op
}
}
var lim limiter.Limiter
var lim traffic.Limiter
if len(lims) > 0 {
lim = newLimiterGroup(lims...)
}
@@ -278,7 +309,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
{
value := values[GlobalLimitKey]
if v, _ := l.inLimits.Get(GlobalLimitKey); v != nil {
lim := v.(limiter.Limiter)
lim := v.(traffic.Limiter)
if value.in <= 0 {
l.inLimits.Delete(GlobalLimitKey)
} else {
@@ -291,7 +322,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
}
if v, _ := l.outLimits.Get(GlobalLimitKey); v != nil {
lim := v.(limiter.Limiter)
lim := v.(traffic.Limiter)
if value.out <= 0 {
l.outLimits.Delete(GlobalLimitKey)
} else {
@@ -321,7 +352,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
if in != value.in {
for _, item := range l.connInLimits.Items() {
if v := item.Object; v != nil {
v.(limiter.Limiter).Set(in)
v.(traffic.Limiter).Set(in)
}
}
}
@@ -333,7 +364,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
if out != value.out {
for _, item := range l.connOutLimits.Items() {
if v := item.Object; v != nil {
v.(limiter.Limiter).Set(out)
v.(traffic.Limiter).Set(out)
}
}
}
@@ -361,7 +392,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
}
if v, _ := l.inLimits.Get(key); v != nil {
lim := v.(limiter.Limiter)
lim := v.(traffic.Limiter)
if value.in <= 0 {
l.inLimits.Delete(key)
} else {
@@ -375,7 +406,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
}
if v, _ := l.outLimits.Get(key); v != nil {
lim := v.(limiter.Limiter)
lim := v.(traffic.Limiter)
if value.out <= 0 {
l.outLimits.Delete(key)
} else {
@@ -398,7 +429,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
l.inLimits.Delete(k)
continue
}
lim := v.Object.(limiter.Limiter)
lim := v.Object.(traffic.Limiter)
if lim.Limit() != in {
lim.Set(in)
}
@@ -415,7 +446,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
l.outLimits.Delete(k)
continue
}
lim := v.Object.(limiter.Limiter)
lim := v.Object.(traffic.Limiter)
if lim.Limit() != out {
lim.Set(out)
}
limiter/traffic/wrapper/conn.go
+186 -235
View File
@@ -7,74 +7,43 @@ import (
"io"
"net"
"syscall"
"time"
limiter "github.com/go-gost/core/limiter/traffic"
"github.com/go-gost/core/limiter"
"github.com/go-gost/core/limiter/traffic"
"github.com/go-gost/core/metadata"
xnet "github.com/go-gost/x/internal/net"
"github.com/go-gost/x/internal/net/udp"
"github.com/patrickmn/go-cache"
)
var (
errUnsupport = errors.New("unsupported operation")
)
// serverConn is a server side Conn with traffic limiter supported.
type serverConn struct {
// limitConn is a Conn with traffic limiter supported.
type limitConn struct {
net.Conn
rbuf bytes.Buffer
limiter limiter.TrafficLimiter
limiterIn limiter.Limiter
limiterOut limiter.Limiter
expIn int64
expOut int64
opts []limiter.Option
rbuf bytes.Buffer
limiter traffic.TrafficLimiter
opts []limiter.Option
key string
}
func WrapConn(tlimiter limiter.TrafficLimiter, c net.Conn) net.Conn {
func WrapConn(c net.Conn, tlimiter traffic.TrafficLimiter, key string, opts ...limiter.Option) net.Conn {
if tlimiter == nil {
return c
}
return &serverConn{
return &limitConn{
Conn: c,
limiter: tlimiter,
opts: []limiter.Option{
limiter.NetworkOption(c.LocalAddr().Network()),
limiter.SrcOption(c.RemoteAddr().String()),
limiter.AddrOption(c.LocalAddr().String()),
},
opts: opts,
key: key,
}
}
func (c *serverConn) getInLimiter() limiter.Limiter {
now := time.Now().UnixNano()
// cache the limiter for 60s
if c.limiter != nil && time.Duration(now-c.expIn) > 60*time.Second {
if lim := c.limiter.In(context.Background(), c.RemoteAddr().String()); lim != nil {
c.limiterIn = lim
}
c.expIn = now
}
return c.limiterIn
}
func (c *serverConn) getOutLimiter() limiter.Limiter {
now := time.Now().UnixNano()
// cache the limiter for 60s
if c.limiter != nil && time.Duration(now-c.expOut) > 60*time.Second {
if lim := c.limiter.Out(context.Background(), c.RemoteAddr().String()); lim != nil {
c.limiterOut = lim
}
c.expOut = now
}
return c.limiterOut
}
func (c *serverConn) Read(b []byte) (n int, err error) {
limiter := c.getInLimiter()
if limiter == nil {
func (c *limitConn) Read(b []byte) (n int, err error) {
limiter := c.limiter.In(context.Background(), c.key, c.opts...)
if limiter == nil || limiter.Limit() <= 0 {
return c.Conn.Read(b)
}
@@ -102,9 +71,9 @@ func (c *serverConn) Read(b []byte) (n int, err error) {
return
}
func (c *serverConn) Write(b []byte) (n int, err error) {
limiter := c.getOutLimiter()
if limiter == nil {
func (c *limitConn) Write(b []byte) (n int, err error) {
limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
if limiter == nil || limiter.Limit() <= 0 {
return c.Conn.Write(b)
}
@@ -121,7 +90,7 @@ func (c *serverConn) Write(b []byte) (n int, err error) {
return
}
func (c *serverConn) SyscallConn() (rc syscall.RawConn, err error) {
func (c *limitConn) SyscallConn() (rc syscall.RawConn, err error) {
if sc, ok := c.Conn.(syscall.Conn); ok {
rc, err = sc.SyscallConn()
return
@@ -130,7 +99,7 @@ func (c *serverConn) SyscallConn() (rc syscall.RawConn, err error) {
return
}
func (c *serverConn) Metadata() metadata.Metadata {
func (c *limitConn) Metadata() metadata.Metadata {
if md, ok := c.Conn.(metadata.Metadatable); ok {
return md.Metadata()
}
@@ -139,71 +108,23 @@ func (c *serverConn) Metadata() metadata.Metadata {
type packetConn struct {
net.PacketConn
limiter limiter.TrafficLimiter
inLimits *cache.Cache
outLimits *cache.Cache
limiter traffic.TrafficLimiter
opts []limiter.Option
key string
}
func WrapPacketConn(lim limiter.TrafficLimiter, pc net.PacketConn) net.PacketConn {
func WrapPacketConn(pc net.PacketConn, lim traffic.TrafficLimiter, key string, opts ...limiter.Option) net.PacketConn {
if lim == nil {
return pc
}
return &packetConn{
PacketConn: pc,
limiter: lim,
inLimits: cache.New(time.Second, 10*time.Second),
outLimits: cache.New(time.Second, 10*time.Second),
opts: opts,
key: key,
}
}
func (c *packetConn) getInLimiter(addr net.Addr) limiter.Limiter {
if c.limiter == nil {
return nil
}
lim, ok := func() (lim limiter.Limiter, ok bool) {
v, ok := c.inLimits.Get(addr.String())
if ok {
if v != nil {
lim = v.(limiter.Limiter)
}
}
return
}()
if ok {
return lim
}
lim = c.limiter.In(context.Background(), addr.String())
c.inLimits.Set(addr.String(), lim, 0)
return lim
}
func (c *packetConn) getOutLimiter(addr net.Addr) limiter.Limiter {
if c.limiter == nil {
return nil
}
lim, ok := func() (lim limiter.Limiter, ok bool) {
v, ok := c.outLimits.Get(addr.String())
if ok {
if v != nil {
lim = v.(limiter.Limiter)
}
}
return
}()
if ok {
return lim
}
lim = c.limiter.Out(context.Background(), addr.String())
c.outLimits.Set(addr.String(), lim, 0)
return lim
}
func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
for {
n, addr, err = c.PacketConn.ReadFrom(p)
@@ -211,8 +132,8 @@ func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
return
}
limiter := c.getInLimiter(addr)
if limiter == nil {
limiter := c.limiter.In(context.Background(), c.key, c.opts...)
if limiter == nil || limiter.Limit() <= 0 {
return
}
@@ -227,7 +148,8 @@ func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
func (c *packetConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
// discard when exceed the limit size.
if limiter := c.getOutLimiter(addr); limiter != nil &&
limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
if limiter != nil && limiter.Limit() > 0 &&
limiter.Wait(context.Background(), len(p)) < len(p) {
n = len(p)
return
@@ -245,68 +167,20 @@ func (c *packetConn) Metadata() metadata.Metadata {
type udpConn struct {
net.PacketConn
limiter limiter.TrafficLimiter
inLimits *cache.Cache
outLimits *cache.Cache
limiter traffic.TrafficLimiter
opts []limiter.Option
key string
}
func WrapUDPConn(limiter limiter.TrafficLimiter, pc net.PacketConn) udp.Conn {
func WrapUDPConn(pc net.PacketConn, limiter traffic.TrafficLimiter, key string, opts ...limiter.Option) udp.Conn {
return &udpConn{
PacketConn: pc,
limiter: limiter,
inLimits: cache.New(time.Second, 10*time.Second),
outLimits: cache.New(time.Second, 10*time.Second),
opts: opts,
key: key,
}
}
func (c *udpConn) getInLimiter(addr net.Addr) limiter.Limiter {
if c.limiter == nil {
return nil
}
lim, ok := func() (lim limiter.Limiter, ok bool) {
v, ok := c.inLimits.Get(addr.String())
if ok {
if v != nil {
lim = v.(limiter.Limiter)
}
}
return
}()
if ok {
return lim
}
lim = c.limiter.In(context.Background(), addr.String())
c.inLimits.Set(addr.String(), lim, 0)
return lim
}
func (c *udpConn) getOutLimiter(addr net.Addr) limiter.Limiter {
if c.limiter == nil {
return nil
}
lim, ok := func() (lim limiter.Limiter, ok bool) {
v, ok := c.outLimits.Get(addr.String())
if ok {
if v != nil {
lim = v.(limiter.Limiter)
}
}
return
}()
if ok {
return lim
}
lim = c.limiter.Out(context.Background(), addr.String())
c.outLimits.Set(addr.String(), lim, 0)
return lim
}
func (c *udpConn) RemoteAddr() net.Addr {
if nc, ok := c.PacketConn.(xnet.RemoteAddr); ok {
return nc.RemoteAddr()
@@ -329,12 +203,34 @@ func (c *udpConn) SetWriteBuffer(n int) error {
}
func (c *udpConn) Read(b []byte) (n int, err error) {
if nc, ok := c.PacketConn.(io.Reader); ok {
n, err = nc.Read(b)
nc, ok := c.PacketConn.(io.Reader)
if !ok {
err = errUnsupport
return
}
for {
n, err = nc.Read(b)
if err != nil {
return
}
if c.limiter == nil {
return
}
limiter := c.limiter.In(context.Background(), c.key, c.opts...)
if limiter == nil || limiter.Limit() <= 0 {
return
}
// discard when exceed the limit size.
if limiter.Wait(context.Background(), n) < n {
continue
}
return
}
err = errUnsupport
return
}
func (c *udpConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
@@ -344,105 +240,160 @@ func (c *udpConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
return
}
if c.limiter == nil {
return
}
limiter := c.limiter.In(context.Background(), c.key, c.opts...)
if limiter == nil || limiter.Limit() <= 0 {
return
}
// discard when exceed the limit size.
if limiter := c.getInLimiter(addr); limiter != nil &&
limiter.Wait(context.Background(), n) < n {
if limiter.Wait(context.Background(), n) < n {
continue
}
return
}
}
func (c *udpConn) ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) {
nc, ok := c.PacketConn.(udp.ReadUDP)
if !ok {
err = errUnsupport
return
}
for {
n, addr, err = nc.ReadFromUDP(b)
if err != nil {
return
}
if c.limiter == nil {
return
}
limiter := c.limiter.In(context.Background(), c.key, c.opts...)
if limiter == nil || limiter.Limit() <= 0 {
return
}
// discard when exceed the limit size.
if limiter.Wait(context.Background(), n) < n {
continue
}
return
}
}
func (c *udpConn) ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error) {
nc, ok := c.PacketConn.(udp.ReadUDP)
if !ok {
err = errUnsupport
return
}
for {
n, oobn, flags, addr, err = nc.ReadMsgUDP(b, oob)
if err != nil {
return
}
if c.limiter == nil {
return
}
limiter := c.limiter.In(context.Background(), c.key, c.opts...)
if limiter == nil || limiter.Limit() <= 0 {
return
}
// discard when exceed the limit size.
if limiter.Wait(context.Background(), n) < n {
continue
}
return
}
}
func (c *udpConn) ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) {
if nc, ok := c.PacketConn.(udp.ReadUDP); ok {
for {
n, addr, err = nc.ReadFromUDP(b)
if err != nil {
return
}
// discard when exceed the limit size.
if limiter := c.getInLimiter(addr); limiter != nil &&
limiter.Wait(context.Background(), n) < n {
continue
}
return
}
}
err = errUnsupport
return
}
func (c *udpConn) ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error) {
if nc, ok := c.PacketConn.(udp.ReadUDP); ok {
for {
n, oobn, flags, addr, err = nc.ReadMsgUDP(b, oob)
if err != nil {
return
}
// discard when exceed the limit size.
if limiter := c.getInLimiter(addr); limiter != nil &&
limiter.Wait(context.Background(), n) < n {
continue
}
return
}
}
err = errUnsupport
return
}
func (c *udpConn) Write(b []byte) (n int, err error) {
if nc, ok := c.PacketConn.(io.Writer); ok {
n, err = nc.Write(b)
func (c *udpConn) Write(p []byte) (n int, err error) {
nc, ok := c.PacketConn.(io.Writer)
if !ok {
err = errUnsupport
return
}
err = errUnsupport
if c.limiter != nil {
// discard when exceed the limit size.
limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
if limiter != nil && limiter.Limit() > 0 &&
limiter.Wait(context.Background(), len(p)) < len(p) {
n = len(p)
return
}
}
n, err = nc.Write(p)
return
}
func (c *udpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
// discard when exceed the limit size.
if limiter := c.getOutLimiter(addr); limiter != nil &&
limiter.Wait(context.Background(), len(p)) < len(p) {
n = len(p)
return
if c.limiter != nil {
// discard when exceed the limit size.
limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
if limiter != nil && limiter.Limit() > 0 &&
limiter.Wait(context.Background(), len(p)) < len(p) {
n = len(p)
return
}
}
n, err = c.PacketConn.WriteTo(p, addr)
return
}
func (c *udpConn) WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error) {
// discard when exceed the limit size.
if limiter := c.getOutLimiter(addr); limiter != nil &&
limiter.Wait(context.Background(), len(b)) < len(b) {
n = len(b)
func (c *udpConn) WriteToUDP(p []byte, addr *net.UDPAddr) (n int, err error) {
nc, ok := c.PacketConn.(udp.WriteUDP)
if !ok {
err = errUnsupport
return
}
if nc, ok := c.PacketConn.(udp.WriteUDP); ok {
n, err = nc.WriteToUDP(b, addr)
return
if c.limiter != nil {
// discard when exceed the limit size.
limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
if limiter != nil && limiter.Limit() > 0 &&
limiter.Wait(context.Background(), len(p)) < len(p) {
n = len(p)
return
}
}
err = errUnsupport
n, err = nc.WriteToUDP(p, addr)
return
}
func (c *udpConn) WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error) {
// discard when exceed the limit size.
if limiter := c.getOutLimiter(addr); limiter != nil &&
limiter.Wait(context.Background(), len(b)) < len(b) {
n = len(b)
func (c *udpConn) WriteMsgUDP(p, oob []byte, addr *net.UDPAddr) (n, oobn int, err error) {
nc, ok := c.PacketConn.(udp.WriteUDP)
if !ok {
err = errUnsupport
return
}
if nc, ok := c.PacketConn.(udp.WriteUDP); ok {
n, oobn, err = nc.WriteMsgUDP(b, oob, addr)
return
if c.limiter != nil {
// discard when exceed the limit size.
limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
if limiter != nil && limiter.Limit() > 0 &&
limiter.Wait(context.Background(), len(p)) < len(p) {
n = len(p)
return
}
}
err = errUnsupport
n, oobn, err = nc.WriteMsgUDP(p, oob, addr)
return
}
limiter/traffic/wrapper/io.go
+12 -39
View File
@@ -4,25 +4,21 @@ import (
"bytes"
"context"
"io"
"time"
limiter "github.com/go-gost/core/limiter/traffic"
"github.com/go-gost/core/limiter"
"github.com/go-gost/core/limiter/traffic"
)
// readWriter is an io.ReadWriter with traffic limiter supported.
type readWriter struct {
io.ReadWriter
rbuf bytes.Buffer
limiter limiter.TrafficLimiter
limiterIn limiter.Limiter
limiterOut limiter.Limiter
expIn int64
expOut int64
opts []limiter.Option
key string
rbuf bytes.Buffer
limiter traffic.TrafficLimiter
opts []limiter.Option
key string
}
func WrapReadWriter(limiter limiter.TrafficLimiter, rw io.ReadWriter, opts ...limiter.Option) io.ReadWriter {
func WrapReadWriter(limiter traffic.TrafficLimiter, rw io.ReadWriter, key string, opts ...limiter.Option) io.ReadWriter {
if limiter == nil {
return rw
}
@@ -31,36 +27,13 @@ func WrapReadWriter(limiter limiter.TrafficLimiter, rw io.ReadWriter, opts ...li
ReadWriter: rw,
limiter: limiter,
opts: opts,
key: key,
}
}
func (p *readWriter) getInLimiter() limiter.Limiter {
now := time.Now().UnixNano()
// cache the limiter for 60s
if p.limiter != nil && time.Duration(now-p.expIn) > 60*time.Second {
if lim := p.limiter.In(context.Background(), p.key, p.opts...); lim != nil {
p.limiterIn = lim
}
p.expIn = now
}
return p.limiterIn
}
func (p *readWriter) getOutLimiter() limiter.Limiter {
now := time.Now().UnixNano()
// cache the limiter for 60s
if p.limiter != nil && time.Duration(now-p.expOut) > 60*time.Second {
if lim := p.limiter.Out(context.Background(), p.key, p.opts...); lim != nil {
p.limiterOut = lim
}
p.expOut = now
}
return p.limiterOut
}
func (p *readWriter) Read(b []byte) (n int, err error) {
limiter := p.getInLimiter()
if limiter == nil {
limiter := p.limiter.In(context.Background(), p.key, p.opts...)
if limiter == nil || limiter.Limit() <= 0 {
return p.ReadWriter.Read(b)
}
@@ -89,8 +62,8 @@ func (p *readWriter) Read(b []byte) (n int, err error) {
}
func (p *readWriter) Write(b []byte) (n int, err error) {
limiter := p.getOutLimiter()
if limiter == nil {
limiter := p.limiter.Out(context.Background(), p.key, p.opts...)
if limiter == nil || limiter.Limit() <= 0 {
return p.ReadWriter.Write(b)
}
limiter/traffic/wrapper/listener.go
+12 -5
View File
@@ -3,22 +3,25 @@ package wrapper
import (
"net"
limiter "github.com/go-gost/core/limiter/traffic"
"github.com/go-gost/core/limiter"
"github.com/go-gost/core/limiter/traffic"
)
type listener struct {
net.Listener
limiter limiter.TrafficLimiter
limiter traffic.TrafficLimiter
service string
}
func WrapListener(limiter limiter.TrafficLimiter, ln net.Listener) net.Listener {
func WrapListener(service string, ln net.Listener, limiter traffic.TrafficLimiter) net.Listener {
if limiter == nil {
return ln
}
return &listener{
limiter: limiter,
Listener: ln,
limiter: limiter,
service: service,
}
}
@@ -28,5 +31,9 @@ func (ln *listener) Accept() (net.Conn, error) {
return nil, err
}
return WrapConn(ln.limiter, c), nil
return WrapConn(c, ln.limiter, "",
limiter.ScopeOption(limiter.ScopeService),
limiter.ServiceOption(ln.service),
limiter.NetworkOption(ln.Addr().Network()),
), nil
}