From 4c2131ca6df291351479620c89c49c05f2e25b11 Mon Sep 17 00:00:00 2001 From: ginuerzh Date: Mon, 5 Sep 2022 22:47:51 +0800 Subject: [PATCH] add rate limiter --- admission/wrapper/conn.go | 219 +++++++++++++++ admission/wrapper/listener.go | 37 +++ config/config.go | 17 +- config/parsing/parse.go | 55 ++-- internal/net/net.go | 23 ++ internal/net/udp/conn.go | 29 ++ limiter/generator.go | 61 +++++ limiter/limiter.go | 30 ++ limiter/rate.go | 439 ++++++++++++++++++++++++++---- limiter/wrapper/conn.go | 340 +++++++++++++++++++++++ limiter/wrapper/listener.go | 32 +++ listener/dns/listener.go | 7 +- listener/ftcp/listener.go | 6 +- listener/grpc/listener.go | 6 +- listener/http2/h2/listener.go | 6 +- listener/http2/listener.go | 6 +- listener/http3/listener.go | 10 +- listener/icmp/listener.go | 6 +- listener/kcp/listener.go | 6 +- listener/mtls/listener.go | 6 +- listener/mws/listener.go | 6 +- listener/obfs/http/listener.go | 7 +- listener/obfs/tls/listener.go | 6 +- listener/pht/listener.go | 6 +- listener/quic/listener.go | 6 +- listener/redirect/tcp/listener.go | 9 +- listener/redirect/udp/listener.go | 6 +- listener/rtcp/listener.go | 6 +- listener/rudp/listener.go | 8 +- listener/ssh/listener.go | 6 +- listener/sshd/listener.go | 6 +- listener/tap/listener.go | 4 +- listener/tcp/listener.go | 6 +- listener/tls/listener.go | 6 +- listener/tun/listener.go | 4 +- listener/udp/listener.go | 6 +- listener/ws/listener.go | 6 +- metrics/wrapper/conn.go | 272 ++++++++++++++++++ metrics/wrapper/listener.go | 32 +++ registry/limiter.go | 8 +- 40 files changed, 1622 insertions(+), 135 deletions(-) create mode 100644 admission/wrapper/conn.go create mode 100644 admission/wrapper/listener.go create mode 100644 internal/net/udp/conn.go create mode 100644 limiter/generator.go create mode 100644 limiter/limiter.go create mode 100644 limiter/wrapper/conn.go create mode 100644 limiter/wrapper/listener.go create mode 100644 metrics/wrapper/conn.go create mode 100644 metrics/wrapper/listener.go diff --git a/admission/wrapper/conn.go b/admission/wrapper/conn.go new file mode 100644 index 0000000..e8a7953 --- /dev/null +++ b/admission/wrapper/conn.go @@ -0,0 +1,219 @@ +package wrapper + +import ( + "errors" + "io" + "net" + "syscall" + + "github.com/go-gost/core/admission" + xnet "github.com/go-gost/x/internal/net" + "github.com/go-gost/x/internal/net/udp" +) + +var ( + errUnsupport = errors.New("unsupported operation") +) + +type serverConn struct { + net.Conn + admission admission.Admission +} + +func WrapConn(admission admission.Admission, c net.Conn) net.Conn { + if admission == nil { + return c + } + return &serverConn{ + Conn: c, + admission: admission, + } +} + +func (c *serverConn) Read(b []byte) (n int, err error) { + if c.admission != nil && + !c.admission.Admit(c.RemoteAddr().String()) { + err = io.EOF + return + } + return c.Conn.Read(b) +} + +func (c *serverConn) SyscallConn() (rc syscall.RawConn, err error) { + if sc, ok := c.Conn.(syscall.Conn); ok { + rc, err = sc.SyscallConn() + return + } + err = errUnsupport + return +} + +type packetConn struct { + net.PacketConn + admission admission.Admission +} + +func WrapPacketConn(admission admission.Admission, pc net.PacketConn) net.PacketConn { + if admission == nil { + return pc + } + return &packetConn{ + PacketConn: pc, + admission: admission, + } +} + +func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) { + for { + n, addr, err = c.PacketConn.ReadFrom(p) + if err != nil { + return + } + + if c.admission != nil && + !c.admission.Admit(addr.String()) { + continue + } + + return + } +} + +type udpConn struct { + net.PacketConn + admission admission.Admission +} + +func WrapUDPConn(admission admission.Admission, pc net.PacketConn) udp.Conn { + return &udpConn{ + PacketConn: pc, + admission: admission, + } +} + +func (c *udpConn) RemoteAddr() net.Addr { + if nc, ok := c.PacketConn.(xnet.RemoteAddr); ok { + return nc.RemoteAddr() + } + return nil +} + +func (c *udpConn) SetReadBuffer(n int) error { + if nc, ok := c.PacketConn.(xnet.SetBuffer); ok { + return nc.SetReadBuffer(n) + } + return errUnsupport +} + +func (c *udpConn) SetWriteBuffer(n int) error { + if nc, ok := c.PacketConn.(xnet.SetBuffer); ok { + return nc.SetWriteBuffer(n) + } + return errUnsupport +} + +func (c *udpConn) Read(b []byte) (n int, err error) { + if nc, ok := c.PacketConn.(io.Reader); ok { + n, err = nc.Read(b) + return + } + err = errUnsupport + return +} + +func (c *udpConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) { + for { + n, addr, err = c.PacketConn.ReadFrom(p) + if err != nil { + return + } + if c.admission != nil && + !c.admission.Admit(addr.String()) { + 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 + } + if c.admission != nil && + !c.admission.Admit(addr.String()) { + 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 + } + if c.admission != nil && + !c.admission.Admit(addr.String()) { + 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) + return + } + err = errUnsupport + return +} + +func (c *udpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) { + n, err = c.PacketConn.WriteTo(p, addr) + return +} + +func (c *udpConn) WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error) { + if nc, ok := c.PacketConn.(udp.WriteUDP); ok { + n, err = nc.WriteToUDP(b, addr) + return + } + err = errUnsupport + return +} + +func (c *udpConn) WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error) { + if nc, ok := c.PacketConn.(udp.WriteUDP); ok { + n, oobn, err = nc.WriteMsgUDP(b, oob, addr) + return + } + err = errUnsupport + return +} + +func (c *udpConn) SyscallConn() (rc syscall.RawConn, err error) { + if nc, ok := c.PacketConn.(xnet.SyscallConn); ok { + return nc.SyscallConn() + } + err = errUnsupport + return +} + +func (c *udpConn) SetDSCP(n int) error { + if nc, ok := c.PacketConn.(xnet.SetDSCP); ok { + return nc.SetDSCP(n) + } + return nil +} diff --git a/admission/wrapper/listener.go b/admission/wrapper/listener.go new file mode 100644 index 0000000..05d4dee --- /dev/null +++ b/admission/wrapper/listener.go @@ -0,0 +1,37 @@ +package wrapper + +import ( + "net" + + "github.com/go-gost/core/admission" +) + +type listener struct { + net.Listener + admission admission.Admission +} + +func WrapListener(admission admission.Admission, ln net.Listener) net.Listener { + if admission == nil { + return ln + } + return &listener{ + Listener: ln, + admission: admission, + } +} + +func (ln *listener) Accept() (net.Conn, error) { + for { + c, err := ln.Listener.Accept() + if err != nil { + return nil, err + } + if ln.admission != nil && + !ln.admission.Admit(c.RemoteAddr().String()) { + c.Close() + continue + } + return c, err + } +} diff --git a/config/config.go b/config/config.go index 79fc72c..4a7b656 100644 --- a/config/config.go +++ b/config/config.go @@ -185,19 +185,20 @@ type RecorderObject struct { } type LimiterConfig struct { - Name string `json:"name"` - RateLimit *RateLimitConfig `yaml:"rate" json:"rate"` + Name string `json:"name"` + Rate *RateLimiterConfig `yaml:"rate" json:"rate"` } -type RateLimitConfig struct { - Input string `yaml:",omitempty" json:"input,omitempty"` - Output string `yaml:",omitempty" json:"output,omitempty"` - Conn *LimitConfig `yaml:",omitempty" json:"conn,omitempty"` +type RateLimiterConfig struct { + Limits []string `yaml:",omitempty" json:"limits,omitempty"` + Reload time.Duration `yaml:",omitempty" json:"reload,omitempty"` + File *FileLoader `yaml:",omitempty" json:"file,omitempty"` + Redis *RedisLoader `yaml:",omitempty" json:"redis,omitempty"` } type LimitConfig struct { - Input string `yaml:",omitempty" json:"input,omitempty"` - Output string `yaml:",omitempty" json:"output,omitempty"` + In string `yaml:",omitempty" json:"in,omitempty"` + Out string `yaml:",omitempty" json:"out,omitempty"` } type ListenerConfig struct { diff --git a/config/parsing/parse.go b/config/parsing/parse.go index 084b69a..b980c1c 100644 --- a/config/parsing/parse.go +++ b/config/parsing/parse.go @@ -4,7 +4,6 @@ import ( "net" "net/url" - "github.com/alecthomas/units" "github.com/go-gost/core/admission" "github.com/go-gost/core/auth" "github.com/go-gost/core/bypass" @@ -323,33 +322,41 @@ func defaultChainSelector() selector.Selector[chain.Chainer] { } func ParseRateLimiter(cfg *config.LimiterConfig) (lim limiter.RateLimiter) { - if cfg == nil || cfg.RateLimit == nil { + if cfg == nil || cfg.Rate == nil { return nil } - var rlimiters []limiter.Limiter - var wlimiters []limiter.Limiter - if cfg.RateLimit.Conn != nil { - if v, _ := units.ParseBase2Bytes(cfg.RateLimit.Conn.Input); v > 0 { - rlimiters = append(rlimiters, xlimiter.Limiter(int(v))) - } - if v, _ := units.ParseBase2Bytes(cfg.RateLimit.Conn.Output); v > 0 { - wlimiters = append(wlimiters, xlimiter.Limiter(int(v))) - } - } - if v, _ := units.ParseBase2Bytes(cfg.RateLimit.Input); v > 0 { - rlimiters = append(rlimiters, xlimiter.Limiter(int(v))) - } - if v, _ := units.ParseBase2Bytes(cfg.RateLimit.Output); v > 0 { - wlimiters = append(wlimiters, xlimiter.Limiter(int(v))) - } + var opts []xlimiter.Option - var input, output limiter.Limiter - if len(rlimiters) > 0 { - input = xlimiter.MultiLimiter(rlimiters...) + if cfg.Rate.File != nil && cfg.Rate.File.Path != "" { + opts = append(opts, xlimiter.FileLoaderOption(loader.FileLoader(cfg.Rate.File.Path))) } - if len(wlimiters) > 0 { - output = xlimiter.MultiLimiter(wlimiters...) + if cfg.Rate.Redis != nil && cfg.Rate.Redis.Addr != "" { + switch cfg.Rate.Redis.Type { + case "list": // redis list + opts = append(opts, xlimiter.RedisLoaderOption(loader.RedisListLoader( + cfg.Rate.Redis.Addr, + loader.DBRedisLoaderOption(cfg.Rate.Redis.DB), + loader.PasswordRedisLoaderOption(cfg.Rate.Redis.Password), + loader.KeyRedisLoaderOption(cfg.Rate.Redis.Key), + ))) + default: // redis set + opts = append(opts, xlimiter.RedisLoaderOption(loader.RedisSetLoader( + cfg.Rate.Redis.Addr, + loader.DBRedisLoaderOption(cfg.Rate.Redis.DB), + loader.PasswordRedisLoaderOption(cfg.Rate.Redis.Password), + loader.KeyRedisLoaderOption(cfg.Rate.Redis.Key), + ))) + } } - return xlimiter.RateLimiter(input, output) + opts = append(opts, + xlimiter.LimitsOption(cfg.Rate.Limits...), + xlimiter.ReloadPeriodOption(cfg.Rate.Reload), + xlimiter.LoggerOption(logger.Default().WithFields(map[string]any{ + "kind": "limiter", + "hosts": cfg.Name, + })), + ) + + return xlimiter.NewRateLimiter(opts...) } diff --git a/internal/net/net.go b/internal/net/net.go index ba90709..364dafa 100644 --- a/internal/net/net.go +++ b/internal/net/net.go @@ -1,5 +1,28 @@ package net +import ( + "net" + "syscall" +) + +type SetBuffer interface { + SetReadBuffer(bytes int) error + SetWriteBuffer(bytes int) error +} + +type SyscallConn interface { + SyscallConn() (syscall.RawConn, error) +} + +type RemoteAddr interface { + RemoteAddr() net.Addr +} + +// tcpraw.TCPConn +type SetDSCP interface { + SetDSCP(int) error +} + func IsIPv4(address string) bool { return address != "" && address[0] != ':' && address[0] != '[' } diff --git a/internal/net/udp/conn.go b/internal/net/udp/conn.go new file mode 100644 index 0000000..c670b78 --- /dev/null +++ b/internal/net/udp/conn.go @@ -0,0 +1,29 @@ +package udp + +import ( + "io" + "net" + + xnet "github.com/go-gost/x/internal/net" +) + +type Conn interface { + net.PacketConn + io.Reader + io.Writer + ReadUDP + WriteUDP + xnet.SetBuffer + xnet.SyscallConn + xnet.RemoteAddr +} + +type ReadUDP interface { + ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) + ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error) +} + +type WriteUDP interface { + WriteToUDP(b []byte, addr *net.UDPAddr) (int, error) + WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error) +} diff --git a/limiter/generator.go b/limiter/generator.go new file mode 100644 index 0000000..ad2bf3c --- /dev/null +++ b/limiter/generator.go @@ -0,0 +1,61 @@ +package limiter + +import ( + "github.com/go-gost/core/limiter" +) + +type RateLimitGenerator interface { + In() limiter.Limiter + Out() limiter.Limiter +} + +type rateLimitGenerator struct { + in int + out int +} + +func NewRateLimitGenerator(in, out int) RateLimitGenerator { + return &rateLimitGenerator{ + in: in, + out: out, + } +} + +func (p *rateLimitGenerator) In() limiter.Limiter { + if p == nil || p.in <= 0 { + return nil + } + return NewLimiter(p.in) +} + +func (p *rateLimitGenerator) Out() limiter.Limiter { + if p == nil || p.out <= 0 { + return nil + } + return NewLimiter(p.out) +} + +type rateLimitSingleGenerator struct { + in limiter.Limiter + out limiter.Limiter +} + +func NewRateLimitSingleGenerator(in, out int) RateLimitGenerator { + p := &rateLimitSingleGenerator{} + if in > 0 { + p.in = NewLimiter(in) + } + if out > 0 { + p.out = NewLimiter(out) + } + + return p +} + +func (p *rateLimitSingleGenerator) In() limiter.Limiter { + return p.in +} + +func (p *rateLimitSingleGenerator) Out() limiter.Limiter { + return p.out +} diff --git a/limiter/limiter.go b/limiter/limiter.go new file mode 100644 index 0000000..ea18409 --- /dev/null +++ b/limiter/limiter.go @@ -0,0 +1,30 @@ +package limiter + +import ( + "context" + + "github.com/go-gost/core/limiter" + "golang.org/x/time/rate" +) + +type llimiter struct { + limiter *rate.Limiter +} + +func NewLimiter(r int) limiter.Limiter { + return &llimiter{ + limiter: rate.NewLimiter(rate.Limit(r), r), + } +} + +func (l *llimiter) Wait(ctx context.Context, n int) int { + if l.limiter.Burst() < n { + n = l.limiter.Burst() + } + l.limiter.WaitN(ctx, n) + return n +} + +func (l *llimiter) Limit() int { + return l.limiter.Burst() +} diff --git a/limiter/rate.go b/limiter/rate.go index 120710b..dfbc71f 100644 --- a/limiter/rate.go +++ b/limiter/rate.go @@ -1,83 +1,412 @@ package limiter import ( + "bufio" "context" + "io" + "net" + "sort" + "strings" + "sync" + "time" + "github.com/alecthomas/units" "github.com/go-gost/core/limiter" - "golang.org/x/time/rate" + "github.com/go-gost/core/logger" + "github.com/go-gost/x/internal/loader" + "github.com/yl2chen/cidranger" ) -type llimiter struct { - limiter *rate.Limiter -} +const ( + GlobalLimitKey = "$" + ConnLimitKey = "$$" +) -func Limiter(r int) limiter.Limiter { - return &llimiter{ - limiter: rate.NewLimiter(rate.Limit(r), r), - } -} - -func (l *llimiter) Limit(b int) int { - if l.limiter.Burst() < b { - b = l.limiter.Burst() - } - l.limiter.WaitN(context.Background(), b) - return b -} - -type Generator interface { - Generate() limiter.Limiter -} - -type limiterGenerator struct { - limit int -} - -func NewGenerator(r int) Generator { - return &limiterGenerator{limit: r} -} - -// Generate creates a new Limiter. -func (g *limiterGenerator) Generate() limiter.Limiter { - return Limiter(g.limit) -} - -type multiLimiter struct { +type limiterGroup struct { limiters []limiter.Limiter } -func MultiLimiter(limiters ...limiter.Limiter) limiter.Limiter { - return &multiLimiter{ - limiters: limiters, +func newLimiterGroup(limiters ...limiter.Limiter) *limiterGroup { + sort.Slice(limiters, func(i, j int) bool { + return limiters[i].Limit() < limiters[j].Limit() + }) + return &limiterGroup{limiters: limiters} +} + +func (l *limiterGroup) Wait(ctx context.Context, n int) int { + for i := range l.limiters { + if v := l.limiters[i].Wait(ctx, n); v < n { + n = v + } + } + return n +} + +func (l *limiterGroup) Limit() int { + if len(l.limiters) == 0 { + return 0 + } + + return l.limiters[0].Limit() +} + +type options struct { + limits []string + fileLoader loader.Loader + redisLoader loader.Loader + period time.Duration + logger logger.Logger +} + +type Option func(opts *options) + +func LimitsOption(limits ...string) Option { + return func(opts *options) { + opts.limits = limits } } -func (l *multiLimiter) Limit(b int) int { - for i := range l.limiters { - b = l.limiters[i].Limit(b) +func ReloadPeriodOption(period time.Duration) Option { + return func(opts *options) { + opts.period = period + } +} + +func FileLoaderOption(fileLoader loader.Loader) Option { + return func(opts *options) { + opts.fileLoader = fileLoader + } +} + +func RedisLoaderOption(redisLoader loader.Loader) Option { + return func(opts *options) { + opts.redisLoader = redisLoader + } +} + +func LoggerOption(logger logger.Logger) Option { + return func(opts *options) { + opts.logger = logger } - return b } type rateLimiter struct { - input limiter.Limiter - output limiter.Limiter + ipLimits map[string]RateLimitGenerator + cidrLimits cidranger.Ranger + inLimits map[string]limiter.Limiter + outLimits map[string]limiter.Limiter + mu sync.RWMutex + cancelFunc context.CancelFunc + options options } -func RateLimiter(input, output limiter.Limiter) limiter.RateLimiter { - if input == nil || output == nil { - return nil +func NewRateLimiter(opts ...Option) limiter.RateLimiter { + var options options + for _, opt := range opts { + opt(&options) } - return &rateLimiter{ - input: input, - output: output, + + ctx, cancel := context.WithCancel(context.TODO()) + lim := &rateLimiter{ + ipLimits: make(map[string]RateLimitGenerator), + cidrLimits: cidranger.NewPCTrieRanger(), + inLimits: make(map[string]limiter.Limiter), + outLimits: make(map[string]limiter.Limiter), + options: options, + cancelFunc: cancel, + } + + if err := lim.reload(ctx); err != nil { + options.logger.Warnf("reload: %v", err) + } + if lim.options.period > 0 { + go lim.periodReload(ctx) + } + return lim +} + +func (l *rateLimiter) In(key string) limiter.Limiter { + l.mu.Lock() + defer l.mu.Unlock() + + if lim, ok := l.inLimits[key]; ok { + return lim + } + + var lims []limiter.Limiter + + if ip := net.ParseIP(key); ip != nil { + found := false + if p := l.ipLimits[key]; p != nil { + if lim := p.In(); lim != nil { + lims = append(lims, lim) + found = true + } + } + if !found { + if p, _ := l.cidrLimits.ContainingNetworks(ip); len(p) > 0 { + if v, _ := p[0].(*cidrLimitEntry); v != nil { + if lim := v.limit.In(); lim != nil { + lims = append(lims, lim) + } + } + } + } + } + + if p := l.ipLimits[ConnLimitKey]; p != nil { + if lim := p.In(); lim != nil { + lims = append(lims, lim) + } + } + if p := l.ipLimits[GlobalLimitKey]; p != nil { + if lim := p.In(); lim != nil { + lims = append(lims, lim) + } + } + + var lim limiter.Limiter + if len(lims) > 0 { + lim = newLimiterGroup(lims...) + } + l.inLimits[key] = lim + + if lim != nil && l.options.logger != nil { + l.options.logger.Debugf("input limit for %s: %d", key, lim.Limit()) + } + + return lim +} + +func (l *rateLimiter) Out(key string) limiter.Limiter { + l.mu.Lock() + defer l.mu.Unlock() + + if lim, ok := l.outLimits[key]; ok { + return lim + } + + var lims []limiter.Limiter + + if ip := net.ParseIP(key); ip != nil { + found := false + if p := l.ipLimits[key]; p != nil { + if lim := p.Out(); lim != nil { + lims = append(lims, lim) + found = true + } + } + if !found { + if p, _ := l.cidrLimits.ContainingNetworks(ip); len(p) > 0 { + if v, _ := p[0].(*cidrLimitEntry); v != nil { + if lim := v.limit.Out(); lim != nil { + lims = append(lims, lim) + } + } + } + } + } + + if p := l.ipLimits[ConnLimitKey]; p != nil { + if lim := p.Out(); lim != nil { + lims = append(lims, lim) + } + } + if p := l.ipLimits[GlobalLimitKey]; p != nil { + if lim := p.Out(); lim != nil { + lims = append(lims, lim) + } + } + + var lim limiter.Limiter + if len(lims) > 0 { + lim = newLimiterGroup(lims...) + } + l.outLimits[key] = lim + + if lim != nil && l.options.logger != nil { + l.options.logger.Debugf("output limit for %s: %d", key, lim.Limit()) + } + + return lim +} + +func (l *rateLimiter) periodReload(ctx context.Context) error { + period := l.options.period + if period < time.Second { + period = time.Second + } + ticker := time.NewTicker(period) + defer ticker.Stop() + + for { + select { + case <-ticker.C: + if err := l.reload(ctx); err != nil { + l.options.logger.Warnf("reload: %v", err) + // return err + } + case <-ctx.Done(): + return ctx.Err() + } } } -func (l *rateLimiter) Input() limiter.Limiter { - return l.input +func (l *rateLimiter) reload(ctx context.Context) error { + v, err := l.load(ctx) + if err != nil { + return err + } + + lines := append(l.options.limits, v...) + + ipLimits := make(map[string]RateLimitGenerator) + cidrLimits := cidranger.NewPCTrieRanger() + + for _, s := range lines { + key, in, out := l.parseLimit(s) + if key == "" { + continue + } + switch key { + case GlobalLimitKey: + ipLimits[key] = NewRateLimitSingleGenerator(in, out) + case ConnLimitKey: + ipLimits[key] = NewRateLimitGenerator(in, out) + default: + if ip := net.ParseIP(key); ip != nil { + ipLimits[key] = NewRateLimitGenerator(in, out) + break + } + if _, ipNet, _ := net.ParseCIDR(key); ipNet != nil { + cidrLimits.Insert(&cidrLimitEntry{ + ipNet: *ipNet, + limit: NewRateLimitGenerator(in, out), + }) + } + } + } + + l.mu.Lock() + defer l.mu.Unlock() + + l.ipLimits = ipLimits + l.cidrLimits = cidrLimits + l.inLimits = make(map[string]limiter.Limiter) + l.outLimits = make(map[string]limiter.Limiter) + + return nil } -func (l *rateLimiter) Output() limiter.Limiter { - return l.output +func (l *rateLimiter) load(ctx context.Context) (patterns []string, err error) { + if l.options.fileLoader != nil { + if lister, ok := l.options.fileLoader.(loader.Lister); ok { + list, er := lister.List(ctx) + if er != nil { + l.options.logger.Warnf("file loader: %v", er) + } + for _, s := range list { + if line := l.parseLine(s); line != "" { + patterns = append(patterns, line) + } + } + } else { + r, er := l.options.fileLoader.Load(ctx) + if er != nil { + l.options.logger.Warnf("file loader: %v", er) + } + if v, _ := l.parsePatterns(r); v != nil { + patterns = append(patterns, v...) + } + } + } + if l.options.redisLoader != nil { + if lister, ok := l.options.redisLoader.(loader.Lister); ok { + list, er := lister.List(ctx) + if er != nil { + l.options.logger.Warnf("redis loader: %v", er) + } + patterns = append(patterns, list...) + } else { + r, er := l.options.redisLoader.Load(ctx) + if er != nil { + l.options.logger.Warnf("redis loader: %v", er) + } + if v, _ := l.parsePatterns(r); v != nil { + patterns = append(patterns, v...) + } + } + } + + l.options.logger.Debugf("load items %d", len(patterns)) + return +} + +func (l *rateLimiter) parsePatterns(r io.Reader) (patterns []string, err error) { + if r == nil { + return + } + + scanner := bufio.NewScanner(r) + for scanner.Scan() { + if line := l.parseLine(scanner.Text()); line != "" { + patterns = append(patterns, line) + } + } + + err = scanner.Err() + return +} + +func (l *rateLimiter) parseLine(s string) string { + if n := strings.IndexByte(s, '#'); n >= 0 { + s = s[:n] + } + return strings.TrimSpace(s) +} + +func (l *rateLimiter) parseLimit(s string) (key string, in, out int) { + s = strings.Replace(s, "\t", " ", -1) + s = strings.TrimSpace(s) + var ss []string + for _, v := range strings.Split(s, " ") { + if v != "" { + ss = append(ss, v) + } + } + if len(ss) < 2 { + return + } + + key = ss[0] + if v, _ := units.ParseBase2Bytes(ss[1]); v > 0 { + in = int(v) + } + if len(ss) > 2 { + if v, _ := units.ParseBase2Bytes(ss[2]); v > 0 { + out = int(v) + } + } + + return +} + +func (l *rateLimiter) Close() error { + l.cancelFunc() + if l.options.fileLoader != nil { + l.options.fileLoader.Close() + } + if l.options.redisLoader != nil { + l.options.redisLoader.Close() + } + return nil +} + +type cidrLimitEntry struct { + ipNet net.IPNet + limit RateLimitGenerator +} + +func (p *cidrLimitEntry) Network() net.IPNet { + return p.ipNet } diff --git a/limiter/wrapper/conn.go b/limiter/wrapper/conn.go new file mode 100644 index 0000000..119c1b3 --- /dev/null +++ b/limiter/wrapper/conn.go @@ -0,0 +1,340 @@ +package wrapper + +import ( + "bytes" + "context" + "errors" + "io" + "net" + "syscall" + + "github.com/go-gost/core/limiter" + xnet "github.com/go-gost/x/internal/net" + "github.com/go-gost/x/internal/net/udp" +) + +var ( + errUnsupport = errors.New("unsupported operation") +) + +// serverConn is a server side Conn with metrics supported. +type serverConn struct { + net.Conn + rbuf bytes.Buffer + raddr string + rlimiter limiter.RateLimiter +} + +func WrapConn(rlimiter limiter.RateLimiter, c net.Conn) net.Conn { + if rlimiter == nil { + return c + } + host, _, _ := net.SplitHostPort(c.RemoteAddr().String()) + return &serverConn{ + Conn: c, + rlimiter: rlimiter, + raddr: host, + } +} + +func (c *serverConn) Read(b []byte) (n int, err error) { + if c.rlimiter == nil || + c.rlimiter.In(c.raddr) == nil { + return c.Conn.Read(b) + } + + limiter := c.rlimiter.In(c.raddr) + + if c.rbuf.Len() > 0 { + burst := len(b) + if c.rbuf.Len() < burst { + burst = c.rbuf.Len() + } + lim := limiter.Wait(context.Background(), burst) + return c.rbuf.Read(b[:lim]) + } + + nn, err := c.Conn.Read(b) + if err != nil { + return nn, err + } + + n = limiter.Wait(context.Background(), nn) + if n < nn { + if _, err = c.rbuf.Write(b[n:nn]); err != nil { + return 0, err + } + } + + return +} + +func (c *serverConn) Write(b []byte) (n int, err error) { + if c.rlimiter == nil || + c.rlimiter.Out(c.raddr) == nil { + return c.Conn.Write(b) + } + + limiter := c.rlimiter.Out(c.raddr) + nn := 0 + for len(b) > 0 { + nn, err = c.Conn.Write(b[:limiter.Wait(context.Background(), len(b))]) + n += nn + if err != nil { + return + } + b = b[nn:] + } + + return +} + +func (c *serverConn) SyscallConn() (rc syscall.RawConn, err error) { + if sc, ok := c.Conn.(syscall.Conn); ok { + rc, err = sc.SyscallConn() + return + } + err = errUnsupport + return +} + +type packetConn struct { + net.PacketConn + rlimiter limiter.RateLimiter +} + +func WrapPacketConn(rlimiter limiter.RateLimiter, pc net.PacketConn) net.PacketConn { + if rlimiter == nil { + return pc + } + return &packetConn{ + PacketConn: pc, + rlimiter: rlimiter, + } +} + +func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) { + for { + n, addr, err = c.PacketConn.ReadFrom(p) + if err != nil { + return + } + + host, _, _ := net.SplitHostPort(addr.String()) + + if c.rlimiter == nil || c.rlimiter.In(host) == nil { + return + } + + limiter := c.rlimiter.In(host) + // discard when exceed the limit size. + if limiter.Wait(context.Background(), n) < n { + continue + } + + return + } +} + +func (c *packetConn) WriteTo(p []byte, addr net.Addr) (n int, err error) { + if c.rlimiter != nil { + host, _, _ := net.SplitHostPort(addr.String()) + // discard when exceed the limit size. + if limiter := c.rlimiter.Out(host); limiter != nil && + limiter.Wait(context.Background(), len(p)) < len(p) { + n = len(p) + return + } + } + + return c.PacketConn.WriteTo(p, addr) +} + +type udpConn struct { + net.PacketConn + rlimiter limiter.RateLimiter +} + +func WrapUDPConn(rlimiter limiter.RateLimiter, pc net.PacketConn) udp.Conn { + return &udpConn{ + PacketConn: pc, + rlimiter: rlimiter, + } +} + +func (c *udpConn) RemoteAddr() net.Addr { + if nc, ok := c.PacketConn.(xnet.RemoteAddr); ok { + return nc.RemoteAddr() + } + return nil +} + +func (c *udpConn) SetReadBuffer(n int) error { + if nc, ok := c.PacketConn.(xnet.SetBuffer); ok { + return nc.SetReadBuffer(n) + } + return errUnsupport +} + +func (c *udpConn) SetWriteBuffer(n int) error { + if nc, ok := c.PacketConn.(xnet.SetBuffer); ok { + return nc.SetWriteBuffer(n) + } + return errUnsupport +} + +func (c *udpConn) Read(b []byte) (n int, err error) { + if nc, ok := c.PacketConn.(io.Reader); ok { + n, err = nc.Read(b) + return + } + err = errUnsupport + return +} + +func (c *udpConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) { + for { + n, addr, err = c.PacketConn.ReadFrom(p) + if err != nil { + return + } + host, _, _ := net.SplitHostPort(addr.String()) + + if c.rlimiter == nil || c.rlimiter.In(host) == nil { + return + } + limiter := c.rlimiter.In(host) + // 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 + } + + host, _, _ := net.SplitHostPort(addr.String()) + + if c.rlimiter == nil || c.rlimiter.In(host) == nil { + return + } + limiter := c.rlimiter.In(host) + // discard when exceed the limit size. + if 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 + } + + host, _, _ := net.SplitHostPort(addr.String()) + + if c.rlimiter == nil || c.rlimiter.In(host) == nil { + return + } + limiter := c.rlimiter.In(host) + // discard when exceed the limit size. + if 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) + return + } + err = errUnsupport + return +} + +func (c *udpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) { + if c.rlimiter != nil { + host, _, _ := net.SplitHostPort(addr.String()) + // discard when exceed the limit size. + if limiter := c.rlimiter.Out(host); limiter != nil && + 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) { + if c.rlimiter != nil { + host, _, _ := net.SplitHostPort(addr.String()) + // discard when exceed the limit size. + if limiter := c.rlimiter.Out(host); limiter != nil && + limiter.Wait(context.Background(), len(b)) < len(b) { + n = len(b) + return + } + } + + if nc, ok := c.PacketConn.(udp.WriteUDP); ok { + n, err = nc.WriteToUDP(b, addr) + return + } + err = errUnsupport + return +} + +func (c *udpConn) WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error) { + if c.rlimiter != nil { + host, _, _ := net.SplitHostPort(addr.String()) + // discard when exceed the limit size. + if limiter := c.rlimiter.Out(host); limiter != nil && + limiter.Wait(context.Background(), len(b)) < len(b) { + n = len(b) + return + } + } + + if nc, ok := c.PacketConn.(udp.WriteUDP); ok { + n, oobn, err = nc.WriteMsgUDP(b, oob, addr) + return + } + err = errUnsupport + return +} + +func (c *udpConn) SyscallConn() (rc syscall.RawConn, err error) { + if nc, ok := c.PacketConn.(xnet.SyscallConn); ok { + return nc.SyscallConn() + } + err = errUnsupport + return +} + +func (c *udpConn) SetDSCP(n int) error { + if nc, ok := c.PacketConn.(xnet.SetDSCP); ok { + return nc.SetDSCP(n) + } + return nil +} diff --git a/limiter/wrapper/listener.go b/limiter/wrapper/listener.go new file mode 100644 index 0000000..42a2c8a --- /dev/null +++ b/limiter/wrapper/listener.go @@ -0,0 +1,32 @@ +package wrapper + +import ( + "net" + + "github.com/go-gost/core/limiter" +) + +type listener struct { + net.Listener + rlimiter limiter.RateLimiter +} + +func WrapListener(rlimiter limiter.RateLimiter, ln net.Listener) net.Listener { + if rlimiter == nil { + return ln + } + + return &listener{ + rlimiter: rlimiter, + Listener: ln, + } +} + +func (ln *listener) Accept() (net.Conn, error) { + c, err := ln.Listener.Accept() + if err != nil { + return nil, err + } + + return WrapConn(ln.rlimiter, c), nil +} diff --git a/listener/dns/listener.go b/listener/dns/listener.go index 35a53b7..ea27ec4 100644 --- a/listener/dns/listener.go +++ b/listener/dns/listener.go @@ -9,10 +9,13 @@ import ( "net/http" "strings" + admission "github.com/go-gost/x/admission/wrapper" + limiter "github.com/go-gost/x/limiter/wrapper" + "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" "github.com/miekg/dns" ) @@ -114,6 +117,8 @@ func (l *dnsListener) Accept() (conn net.Conn, err error) { select { case conn = <-l.cqueue: conn = metrics.WrapConn(l.options.Service, conn) + conn = admission.WrapConn(l.options.Admission, conn) + conn = limiter.WrapConn(l.options.RateLimiter, conn) case err, ok = <-l.errChan: if !ok { err = listener.ErrClosed diff --git a/listener/ftcp/listener.go b/listener/ftcp/listener.go index 5c6c83c..3dfabf6 100644 --- a/listener/ftcp/listener.go +++ b/listener/ftcp/listener.go @@ -7,8 +7,10 @@ import ( "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" "github.com/xtaci/tcpraw" ) @@ -50,6 +52,8 @@ func (l *ftcpListener) Init(md md.Metadata) (err error) { return } conn = metrics.WrapPacketConn(l.options.Service, conn) + conn = admission.WrapPacketConn(l.options.Admission, conn) + conn = limiter.WrapPacketConn(l.options.RateLimiter, conn) l.ln = udp.NewListener( conn, diff --git a/listener/grpc/listener.go b/listener/grpc/listener.go index 5ee75c9..d8969da 100644 --- a/listener/grpc/listener.go +++ b/listener/grpc/listener.go @@ -3,13 +3,14 @@ package grpc import ( "net" - admission "github.com/go-gost/core/admission/wrapper" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" pb "github.com/go-gost/x/internal/util/grpc/proto" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" "google.golang.org/grpc" "google.golang.org/grpc/credentials" @@ -55,6 +56,7 @@ func (l *grpcListener) Init(md md.Metadata) (err error) { } ln = metrics.WrapListener(l.options.Service, ln) ln = admission.WrapListener(l.options.Admission, ln) + ln = limiter.WrapListener(l.options.RateLimiter, ln) var opts []grpc.ServerOption if !l.md.insecure { diff --git a/listener/http2/h2/listener.go b/listener/http2/h2/listener.go index c80fdce..2dde980 100644 --- a/listener/http2/h2/listener.go +++ b/listener/http2/h2/listener.go @@ -7,12 +7,13 @@ import ( "net/http" "net/http/httputil" - admission "github.com/go-gost/core/admission/wrapper" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" "golang.org/x/net/http2" "golang.org/x/net/http2/h2c" @@ -77,6 +78,7 @@ func (l *h2Listener) Init(md md.Metadata) (err error) { l.addr = ln.Addr() ln = metrics.WrapListener(l.options.Service, ln) ln = admission.WrapListener(l.options.Admission, ln) + ln = limiter.WrapListener(l.options.RateLimiter, ln) if l.h2c { l.server.Handler = h2c.NewHandler( diff --git a/listener/http2/listener.go b/listener/http2/listener.go index 434b677..c3d1be6 100644 --- a/listener/http2/listener.go +++ b/listener/http2/listener.go @@ -5,13 +5,14 @@ import ( "net" "net/http" - admission "github.com/go-gost/core/admission/wrapper" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" + limiter "github.com/go-gost/x/limiter/wrapper" mdx "github.com/go-gost/x/metadata" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" "golang.org/x/net/http2" ) @@ -66,6 +67,7 @@ func (l *http2Listener) Init(md md.Metadata) (err error) { l.addr = ln.Addr() ln = metrics.WrapListener(l.options.Service, ln) ln = admission.WrapListener(l.options.Admission, ln) + ln = limiter.WrapListener(l.options.RateLimiter, ln) ln = tls.NewListener( ln, diff --git a/listener/http3/listener.go b/listener/http3/listener.go index 5d6b7ee..1f8e41f 100644 --- a/listener/http3/listener.go +++ b/listener/http3/listener.go @@ -6,9 +6,11 @@ import ( "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" pht_util "github.com/go-gost/x/internal/util/pht" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" "github.com/lucas-clemente/quic-go" ) @@ -74,7 +76,11 @@ func (l *http3Listener) Accept() (conn net.Conn, err error) { if err != nil { return } - return metrics.WrapConn(l.options.Service, conn), nil + + conn = metrics.WrapConn(l.options.Service, conn) + conn = admission.WrapConn(l.options.Admission, conn) + conn = limiter.WrapConn(l.options.RateLimiter, conn) + return conn, nil } func (l *http3Listener) Addr() net.Addr { diff --git a/listener/icmp/listener.go b/listener/icmp/listener.go index 40cf52d..23dad6d 100644 --- a/listener/icmp/listener.go +++ b/listener/icmp/listener.go @@ -4,12 +4,13 @@ import ( "context" "net" - admission "github.com/go-gost/core/admission/wrapper" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" icmp_pkg "github.com/go-gost/x/internal/util/icmp" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" "github.com/lucas-clemente/quic-go" "golang.org/x/net/icmp" @@ -57,6 +58,7 @@ func (l *icmpListener) Init(md md.Metadata) (err error) { conn = icmp_pkg.ServerConn(conn) conn = metrics.WrapPacketConn(l.options.Service, conn) conn = admission.WrapPacketConn(l.options.Admission, conn) + conn = limiter.WrapPacketConn(l.options.RateLimiter, conn) config := &quic.Config{ KeepAlivePeriod: l.md.keepAlivePeriod, diff --git a/listener/kcp/listener.go b/listener/kcp/listener.go index 159113e..972179c 100644 --- a/listener/kcp/listener.go +++ b/listener/kcp/listener.go @@ -4,13 +4,14 @@ import ( "net" "time" - admission "github.com/go-gost/core/admission/wrapper" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" kcp_util "github.com/go-gost/x/internal/util/kcp" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" "github.com/xtaci/kcp-go/v5" "github.com/xtaci/smux" @@ -75,6 +76,7 @@ func (l *kcpListener) Init(md md.Metadata) (err error) { conn = metrics.WrapUDPConn(l.options.Service, conn) conn = admission.WrapUDPConn(l.options.Admission, conn) + conn = limiter.WrapUDPConn(l.options.RateLimiter, conn) ln, err := kcp.ServeConn( kcp_util.BlockCrypt(config.Key, config.Crypt, kcp_util.DefaultSalt), diff --git a/listener/mtls/listener.go b/listener/mtls/listener.go index 4c5aa15..5db8e87 100644 --- a/listener/mtls/listener.go +++ b/listener/mtls/listener.go @@ -4,12 +4,13 @@ import ( "crypto/tls" "net" - admission "github.com/go-gost/core/admission/wrapper" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" "github.com/xtaci/smux" ) @@ -54,6 +55,7 @@ func (l *mtlsListener) Init(md md.Metadata) (err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = admission.WrapListener(l.options.Admission, ln) + ln = limiter.WrapListener(l.options.RateLimiter, ln) l.Listener = tls.NewListener(ln, l.options.TLSConfig) l.cqueue = make(chan net.Conn, l.md.backlog) diff --git a/listener/mws/listener.go b/listener/mws/listener.go index efd21aa..bd89ab7 100644 --- a/listener/mws/listener.go +++ b/listener/mws/listener.go @@ -6,13 +6,14 @@ import ( "net/http" "net/http/httputil" - admission "github.com/go-gost/core/admission/wrapper" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" ws_util "github.com/go-gost/x/internal/util/ws" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" "github.com/gorilla/websocket" "github.com/xtaci/smux" @@ -96,6 +97,7 @@ func (l *mwsListener) Init(md md.Metadata) (err error) { } ln = metrics.WrapListener(l.options.Service, ln) ln = admission.WrapListener(l.options.Admission, ln) + ln = limiter.WrapListener(l.options.RateLimiter, ln) if l.tlsEnabled { ln = tls.NewListener(ln, l.options.TLSConfig) diff --git a/listener/obfs/http/listener.go b/listener/obfs/http/listener.go index 385d19d..53d8f2f 100644 --- a/listener/obfs/http/listener.go +++ b/listener/obfs/http/listener.go @@ -3,11 +3,13 @@ package http import ( "net" - admission "github.com/go-gost/core/admission/wrapper" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" + xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/registry" ) @@ -49,6 +51,7 @@ func (l *obfsListener) Init(md md.Metadata) (err error) { } ln = metrics.WrapListener(l.options.Service, ln) ln = admission.WrapListener(l.options.Admission, ln) + ln = limiter.WrapListener(l.options.RateLimiter, ln) l.Listener = ln return diff --git a/listener/obfs/tls/listener.go b/listener/obfs/tls/listener.go index 3a6c565..1c4576b 100644 --- a/listener/obfs/tls/listener.go +++ b/listener/obfs/tls/listener.go @@ -3,12 +3,13 @@ package tls import ( "net" - admission "github.com/go-gost/core/admission/wrapper" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" ) @@ -49,6 +50,7 @@ func (l *obfsListener) Init(md md.Metadata) (err error) { } ln = metrics.WrapListener(l.options.Service, ln) ln = admission.WrapListener(l.options.Admission, ln) + ln = limiter.WrapListener(l.options.RateLimiter, ln) l.Listener = ln return diff --git a/listener/pht/listener.go b/listener/pht/listener.go index 868aabd..96339a2 100644 --- a/listener/pht/listener.go +++ b/listener/pht/listener.go @@ -8,9 +8,11 @@ import ( "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" pht_util "github.com/go-gost/x/internal/util/pht" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" ) @@ -89,6 +91,8 @@ func (l *phtListener) Accept() (conn net.Conn, err error) { return } conn = metrics.WrapConn(l.options.Service, conn) + conn = admission.WrapConn(l.options.Admission, conn) + conn = limiter.WrapConn(l.options.RateLimiter, conn) return } diff --git a/listener/quic/listener.go b/listener/quic/listener.go index 9b264a6..6ec3a4e 100644 --- a/listener/quic/listener.go +++ b/listener/quic/listener.go @@ -7,9 +7,11 @@ import ( "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" quic_util "github.com/go-gost/x/internal/util/quic" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" "github.com/lucas-clemente/quic-go" ) @@ -99,6 +101,8 @@ func (l *quicListener) Accept() (conn net.Conn, err error) { select { case conn = <-l.cqueue: conn = metrics.WrapConn(l.options.Service, conn) + conn = admission.WrapConn(l.options.Admission, conn) + conn = limiter.WrapConn(l.options.RateLimiter, conn) case err, ok = <-l.errChan: if !ok { err = listener.ErrClosed diff --git a/listener/redirect/tcp/listener.go b/listener/redirect/tcp/listener.go index 6c3e780..57f6f5b 100644 --- a/listener/redirect/tcp/listener.go +++ b/listener/redirect/tcp/listener.go @@ -7,8 +7,10 @@ import ( "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" ) @@ -54,7 +56,10 @@ func (l *redirectListener) Init(md md.Metadata) (err error) { return err } - l.ln = metrics.WrapListener(l.options.Service, ln) + ln = metrics.WrapListener(l.options.Service, ln) + ln = admission.WrapListener(l.options.Admission, ln) + ln = limiter.WrapListener(l.options.RateLimiter, ln) + l.ln = ln return } diff --git a/listener/redirect/udp/listener.go b/listener/redirect/udp/listener.go index 255921b..abc1cf9 100644 --- a/listener/redirect/udp/listener.go +++ b/listener/redirect/udp/listener.go @@ -6,7 +6,9 @@ import ( "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" ) @@ -52,6 +54,8 @@ func (l *redirectListener) Accept() (conn net.Conn, err error) { return } conn = metrics.WrapConn(l.options.Service, conn) + conn = admission.WrapConn(l.options.Admission, conn) + conn = limiter.WrapConn(l.options.RateLimiter, conn) return } diff --git a/listener/rtcp/listener.go b/listener/rtcp/listener.go index b327dd8..d41cd7c 100644 --- a/listener/rtcp/listener.go +++ b/listener/rtcp/listener.go @@ -8,8 +8,10 @@ import ( "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" ) @@ -76,6 +78,8 @@ func (l *rtcpListener) Accept() (conn net.Conn, err error) { return nil, listener.NewAcceptError(err) } l.ln = metrics.WrapListener(l.options.Service, l.ln) + l.ln = admission.WrapListener(l.options.Admission, l.ln) + l.ln = limiter.WrapListener(l.options.RateLimiter, l.ln) } conn, err = l.ln.Accept() if err != nil { diff --git a/listener/rudp/listener.go b/listener/rudp/listener.go index 65c8cb2..981dbcc 100644 --- a/listener/rudp/listener.go +++ b/listener/rudp/listener.go @@ -8,8 +8,10 @@ import ( "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" ) @@ -88,7 +90,9 @@ func (l *rudpListener) Accept() (conn net.Conn, err error) { } if pc, ok := conn.(net.PacketConn); ok { - conn = metrics.WrapUDPConn(l.options.Service, pc) + uc := metrics.WrapUDPConn(l.options.Service, pc) + uc = admission.WrapUDPConn(l.options.Admission, uc) + conn = limiter.WrapUDPConn(l.options.RateLimiter, uc) } return diff --git a/listener/ssh/listener.go b/listener/ssh/listener.go index 5ad5601..b505960 100644 --- a/listener/ssh/listener.go +++ b/listener/ssh/listener.go @@ -5,13 +5,14 @@ import ( "net" "time" - admission "github.com/go-gost/core/admission/wrapper" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" ssh_util "github.com/go-gost/x/internal/util/ssh" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" "golang.org/x/crypto/ssh" ) @@ -57,6 +58,7 @@ func (l *sshListener) Init(md md.Metadata) (err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = admission.WrapListener(l.options.Admission, ln) + ln = limiter.WrapListener(l.options.RateLimiter, ln) l.Listener = ln config := &ssh.ServerConfig{ diff --git a/listener/sshd/listener.go b/listener/sshd/listener.go index e3aab90..63f26dc 100644 --- a/listener/sshd/listener.go +++ b/listener/sshd/listener.go @@ -7,14 +7,15 @@ import ( "strconv" "time" - admission "github.com/go-gost/core/admission/wrapper" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" ssh_util "github.com/go-gost/x/internal/util/ssh" sshd_util "github.com/go-gost/x/internal/util/sshd" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" "golang.org/x/crypto/ssh" ) @@ -66,6 +67,7 @@ func (l *sshdListener) Init(md md.Metadata) (err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = admission.WrapListener(l.options.Admission, ln) + ln = limiter.WrapListener(l.options.RateLimiter, ln) l.Listener = ln config := &ssh.ServerConfig{ diff --git a/listener/tap/listener.go b/listener/tap/listener.go index 44b6584..fe1ed1c 100644 --- a/listener/tap/listener.go +++ b/listener/tap/listener.go @@ -6,9 +6,10 @@ import ( "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" mdata "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" xnet "github.com/go-gost/x/internal/net" + limiter "github.com/go-gost/x/limiter/wrapper" mdx "github.com/go-gost/x/metadata" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" ) @@ -79,6 +80,7 @@ func (l *tapListener) Init(md mdata.Metadata) (err error) { raddr: &net.IPAddr{IP: ip}, } c = metrics.WrapConn(l.options.Service, c) + c = limiter.WrapConn(l.options.RateLimiter, c) c = withMetadata(mdx.NewMetadata(map[string]any{ "config": l.md.config, }), c) diff --git a/listener/tcp/listener.go b/listener/tcp/listener.go index 366b6a7..2b79a9f 100644 --- a/listener/tcp/listener.go +++ b/listener/tcp/listener.go @@ -3,12 +3,13 @@ package tcp import ( "net" - limiter "github.com/go-gost/core/limiter/wrapper" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" ) @@ -49,6 +50,7 @@ func (l *tcpListener) Init(md md.Metadata) (err error) { } ln = metrics.WrapListener(l.options.Service, ln) + ln = admission.WrapListener(l.options.Admission, ln) l.ln = limiter.WrapListener(l.options.RateLimiter, ln) return diff --git a/listener/tls/listener.go b/listener/tls/listener.go index 2c39f3d..8f54b7a 100644 --- a/listener/tls/listener.go +++ b/listener/tls/listener.go @@ -4,12 +4,13 @@ import ( "crypto/tls" "net" - admission "github.com/go-gost/core/admission/wrapper" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" ) @@ -50,6 +51,7 @@ func (l *tlsListener) Init(md md.Metadata) (err error) { } ln = metrics.WrapListener(l.options.Service, ln) ln = admission.WrapListener(l.options.Admission, ln) + ln = limiter.WrapListener(l.options.RateLimiter, ln) l.ln = tls.NewListener(ln, l.options.TLSConfig) diff --git a/listener/tun/listener.go b/listener/tun/listener.go index d47343a..3b2c862 100644 --- a/listener/tun/listener.go +++ b/listener/tun/listener.go @@ -8,9 +8,10 @@ import ( "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" mdata "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" xnet "github.com/go-gost/x/internal/net" + limiter "github.com/go-gost/x/limiter/wrapper" mdx "github.com/go-gost/x/metadata" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" ) @@ -88,6 +89,7 @@ func (l *tunListener) listenLoop() { cancel: cancel, } c = metrics.WrapConn(l.options.Service, c) + c = limiter.WrapConn(l.options.RateLimiter, c) c = withMetadata(mdx.NewMetadata(map[string]any{ "config": l.md.config, }), c) diff --git a/listener/udp/listener.go b/listener/udp/listener.go index d5983f2..323f02f 100644 --- a/listener/udp/listener.go +++ b/listener/udp/listener.go @@ -4,12 +4,13 @@ import ( "net" "github.com/go-gost/core/common/net/udp" - limiter "github.com/go-gost/core/limiter/wrapper" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" ) @@ -55,6 +56,7 @@ func (l *udpListener) Init(md md.Metadata) (err error) { return } conn = metrics.WrapPacketConn(l.options.Service, conn) + conn = admission.WrapPacketConn(l.options.Admission, conn) conn = limiter.WrapPacketConn(l.options.RateLimiter, conn) l.ln = udp.NewListener(conn, &udp.ListenConfig{ diff --git a/listener/ws/listener.go b/listener/ws/listener.go index aa30a90..a6377e2 100644 --- a/listener/ws/listener.go +++ b/listener/ws/listener.go @@ -6,13 +6,14 @@ import ( "net/http" "net/http/httputil" - admission "github.com/go-gost/core/admission/wrapper" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" - metrics "github.com/go-gost/core/metrics/wrapper" + admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" ws_util "github.com/go-gost/x/internal/util/ws" + limiter "github.com/go-gost/x/limiter/wrapper" + metrics "github.com/go-gost/x/metrics/wrapper" "github.com/go-gost/x/registry" "github.com/gorilla/websocket" ) @@ -91,6 +92,7 @@ func (l *wsListener) Init(md md.Metadata) (err error) { } ln = metrics.WrapListener(l.options.Service, ln) ln = admission.WrapListener(l.options.Admission, ln) + ln = limiter.WrapListener(l.options.RateLimiter, ln) if l.tlsEnabled { ln = tls.NewListener(ln, l.options.TLSConfig) diff --git a/metrics/wrapper/conn.go b/metrics/wrapper/conn.go new file mode 100644 index 0000000..2a2ace7 --- /dev/null +++ b/metrics/wrapper/conn.go @@ -0,0 +1,272 @@ +package wrapper + +import ( + "errors" + "io" + "net" + "syscall" + + "github.com/go-gost/core/metrics" + xnet "github.com/go-gost/x/internal/net" + "github.com/go-gost/x/internal/net/udp" +) + +var ( + errUnsupport = errors.New("unsupported operation") +) + +// serverConn is a server side Conn with metrics supported. +type serverConn struct { + net.Conn + service string +} + +func WrapConn(service string, c net.Conn) net.Conn { + if !metrics.IsEnabled() { + return c + } + return &serverConn{ + service: service, + Conn: c, + } +} + +func (c *serverConn) Read(b []byte) (n int, err error) { + n, err = c.Conn.Read(b) + if counter := metrics.GetCounter( + metrics.MetricServiceTransferInputBytesCounter, + metrics.Labels{ + "service": c.service, + }); counter != nil { + counter.Add(float64(n)) + } + return +} + +func (c *serverConn) Write(b []byte) (n int, err error) { + n, err = c.Conn.Write(b) + if counter := metrics.GetCounter( + metrics.MetricServiceTransferOutputBytesCounter, + metrics.Labels{ + "service": c.service, + }); counter != nil { + counter.Add(float64(n)) + } + return +} + +func (c *serverConn) SyscallConn() (rc syscall.RawConn, err error) { + if sc, ok := c.Conn.(syscall.Conn); ok { + rc, err = sc.SyscallConn() + return + } + err = errUnsupport + return +} + +type packetConn struct { + net.PacketConn + service string +} + +func WrapPacketConn(service string, pc net.PacketConn) net.PacketConn { + if !metrics.IsEnabled() { + return pc + } + return &packetConn{ + PacketConn: pc, + service: service, + } +} + +func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) { + n, addr, err = c.PacketConn.ReadFrom(p) + if counter := metrics.GetCounter( + metrics.MetricServiceTransferInputBytesCounter, + metrics.Labels{ + "service": c.service, + }); counter != nil { + counter.Add(float64(n)) + } + return +} + +func (c *packetConn) WriteTo(p []byte, addr net.Addr) (n int, err error) { + n, err = c.PacketConn.WriteTo(p, addr) + if counter := metrics.GetCounter( + metrics.MetricServiceTransferOutputBytesCounter, + metrics.Labels{ + "service": c.service, + }); counter != nil { + counter.Add(float64(n)) + } + return +} + +type udpConn struct { + net.PacketConn + service string +} + +func WrapUDPConn(service string, pc net.PacketConn) udp.Conn { + return &udpConn{ + PacketConn: pc, + service: service, + } +} + +func (c *udpConn) RemoteAddr() net.Addr { + if nc, ok := c.PacketConn.(xnet.RemoteAddr); ok { + return nc.RemoteAddr() + } + return nil +} + +func (c *udpConn) SetReadBuffer(n int) error { + if nc, ok := c.PacketConn.(xnet.SetBuffer); ok { + return nc.SetReadBuffer(n) + } + return errUnsupport +} + +func (c *udpConn) SetWriteBuffer(n int) error { + if nc, ok := c.PacketConn.(xnet.SetBuffer); ok { + return nc.SetWriteBuffer(n) + } + return errUnsupport +} + +func (c *udpConn) Read(b []byte) (n int, err error) { + if nc, ok := c.PacketConn.(io.Reader); ok { + n, err = nc.Read(b) + if counter := metrics.GetCounter( + metrics.MetricServiceTransferInputBytesCounter, + metrics.Labels{ + "service": c.service, + }); counter != nil { + counter.Add(float64(n)) + } + return + } + err = errUnsupport + return +} + +func (c *udpConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) { + n, addr, err = c.PacketConn.ReadFrom(p) + if counter := metrics.GetCounter( + metrics.MetricServiceTransferInputBytesCounter, + metrics.Labels{ + "service": c.service, + }); counter != nil { + counter.Add(float64(n)) + } + return +} + +func (c *udpConn) ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) { + if nc, ok := c.PacketConn.(udp.ReadUDP); ok { + n, addr, err = nc.ReadFromUDP(b) + if counter := metrics.GetCounter( + metrics.MetricServiceTransferInputBytesCounter, + metrics.Labels{ + "service": c.service, + }); counter != nil { + counter.Add(float64(n)) + } + 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 { + n, oobn, flags, addr, err = nc.ReadMsgUDP(b, oob) + if counter := metrics.GetCounter( + metrics.MetricServiceTransferInputBytesCounter, + metrics.Labels{ + "service": c.service, + }); counter != nil { + counter.Add(float64(n)) + } + 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) + if counter := metrics.GetCounter( + metrics.MetricServiceTransferOutputBytesCounter, + metrics.Labels{ + "service": c.service, + }); counter != nil { + counter.Add(float64(n)) + } + return + } + err = errUnsupport + return +} + +func (c *udpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) { + n, err = c.PacketConn.WriteTo(p, addr) + if counter := metrics.GetCounter( + metrics.MetricServiceTransferOutputBytesCounter, + metrics.Labels{ + "service": c.service, + }); counter != nil { + counter.Add(float64(n)) + } + return +} + +func (c *udpConn) WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error) { + if nc, ok := c.PacketConn.(udp.WriteUDP); ok { + n, err = nc.WriteToUDP(b, addr) + if counter := metrics.GetCounter( + metrics.MetricServiceTransferOutputBytesCounter, + metrics.Labels{ + "service": c.service, + }); counter != nil { + counter.Add(float64(n)) + } + return + } + err = errUnsupport + return +} + +func (c *udpConn) WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error) { + if nc, ok := c.PacketConn.(udp.WriteUDP); ok { + n, oobn, err = nc.WriteMsgUDP(b, oob, addr) + if counter := metrics.GetCounter( + metrics.MetricServiceTransferOutputBytesCounter, + metrics.Labels{ + "service": c.service, + }); counter != nil { + counter.Add(float64(n)) + } + return + } + err = errUnsupport + return +} + +func (c *udpConn) SyscallConn() (rc syscall.RawConn, err error) { + if nc, ok := c.PacketConn.(syscall.Conn); ok { + return nc.SyscallConn() + } + err = errUnsupport + return +} + +func (c *udpConn) SetDSCP(n int) error { + if nc, ok := c.PacketConn.(xnet.SetDSCP); ok { + return nc.SetDSCP(n) + } + return nil +} diff --git a/metrics/wrapper/listener.go b/metrics/wrapper/listener.go new file mode 100644 index 0000000..f663cb8 --- /dev/null +++ b/metrics/wrapper/listener.go @@ -0,0 +1,32 @@ +package wrapper + +import ( + "net" + + "github.com/go-gost/core/metrics" +) + +type listener struct { + service string + net.Listener +} + +func WrapListener(service string, ln net.Listener) net.Listener { + if !metrics.IsEnabled() { + return ln + } + + return &listener{ + service: service, + Listener: ln, + } +} + +func (ln *listener) Accept() (net.Conn, error) { + c, err := ln.Listener.Accept() + if err != nil { + return nil, err + } + + return WrapConn(ln.service, c), nil +} diff --git a/registry/limiter.go b/registry/limiter.go index 69a48d7..2ebadf3 100644 --- a/registry/limiter.go +++ b/registry/limiter.go @@ -31,18 +31,18 @@ type rlimiterWrapper struct { r *rlimiterRegistry } -func (w *rlimiterWrapper) Input() limiter.Limiter { +func (w *rlimiterWrapper) In(key string) limiter.Limiter { v := w.r.get(w.name) if v == nil { return nil } - return v.Input() + return v.In(key) } -func (w *rlimiterWrapper) Output() limiter.Limiter { +func (w *rlimiterWrapper) Out(key string) limiter.Limiter { v := w.r.get(w.name) if v == nil { return nil } - return v.Output() + return v.Out(key) }