diff --git a/go.mod b/go.mod index cdb6f135..81f4dd4a 100644 --- a/go.mod +++ b/go.mod @@ -9,10 +9,10 @@ require ( github.com/asaskevich/govalidator v0.0.0-20210307081110-f21760c49a8d github.com/gin-contrib/cors v1.6.0 github.com/gin-gonic/gin v1.9.1 - github.com/go-gost/core v0.1.1 + github.com/go-gost/core v0.1.2 github.com/go-gost/gosocks4 v0.0.1 github.com/go-gost/gosocks5 v0.4.2 - github.com/go-gost/plugin v0.0.0-20240103125338-9c84e29cb81a + github.com/go-gost/plugin v0.1.1 github.com/go-gost/relay v0.5.0 github.com/go-gost/tls-dissector v0.0.2-0.20220408131628-aac992c27451 github.com/go-redis/redis/v8 v8.11.5 diff --git a/go.sum b/go.sum index 763f689b..0bd8cc51 100644 --- a/go.sum +++ b/go.sum @@ -55,12 +55,18 @@ github.com/gin-gonic/gin v1.9.1 h1:4idEAncQnU5cB7BeOkPtxjfCSye0AAm1R0RVIqJ+Jmg= github.com/gin-gonic/gin v1.9.1/go.mod h1:hPrL7YrpYKXt5YId3A/Tnip5kqbEAP+KLuI3SUcPTeU= github.com/go-gost/core v0.1.1 h1:8joR9KJYBvpurNu3i0zqN9orQthVzOjhtT4STumwNF0= github.com/go-gost/core v0.1.1/go.mod h1:WGI43jOka7FAsSAwi/fSMaqxdR+E339ycb4NBGlFr6A= +github.com/go-gost/core v0.1.2 h1:uWGLXEcfqkLYwvpGutXN2eIGXLOUGQX7L1QTe3NUDeA= +github.com/go-gost/core v0.1.2/go.mod h1:WGI43jOka7FAsSAwi/fSMaqxdR+E339ycb4NBGlFr6A= github.com/go-gost/gosocks4 v0.0.1 h1:+k1sec8HlELuQV7rWftIkmy8UijzUt2I6t+iMPlGB2s= github.com/go-gost/gosocks4 v0.0.1/go.mod h1:3B6L47HbU/qugDg4JnoFPHgJXE43Inz8Bah1QaN9qCc= github.com/go-gost/gosocks5 v0.4.2 h1:IianxHTkACPqCwiOAT3MHoMdSUl+SEPSRu1ikawC1Pc= github.com/go-gost/gosocks5 v0.4.2/go.mod h1:1G6I7HP7VFVxveGkoK8mnprnJqSqJjdcASKsdUn4Pp4= github.com/go-gost/plugin v0.0.0-20240103125338-9c84e29cb81a h1:ME7P1Brcg4C640DSPqlvQr7JuvvQfJ8QpmS3yCFlK3A= github.com/go-gost/plugin v0.0.0-20240103125338-9c84e29cb81a/go.mod h1:qXr2Zm9Ex2ATqnWuNUzVZqySPMnuIihvblYZt4MlZLw= +github.com/go-gost/plugin v0.1.0 h1:3C4pm+jyEy8EpGN2DzVUM4eS75T/GVxCq2fU4RuiFJs= +github.com/go-gost/plugin v0.1.0/go.mod h1:oN23l+yGDCIP9G3KnDl/I/0zVGOobZUDCB2Z5yYYXts= +github.com/go-gost/plugin v0.1.1 h1:LNoc/Rqwb3ceGhhxwhpjf1SeeYBGe80MJ9E4lpM3qak= +github.com/go-gost/plugin v0.1.1/go.mod h1:oN23l+yGDCIP9G3KnDl/I/0zVGOobZUDCB2Z5yYYXts= github.com/go-gost/relay v0.5.0 h1:JG1tgy/KWiVXS0ukuVXvbM0kbYuJTWxYpJ5JwzsCf/c= github.com/go-gost/relay v0.5.0/go.mod h1:lcX+23LCQ3khIeASBo+tJ/WbwXFO32/N5YN6ucuYTG8= github.com/go-gost/tls-dissector v0.0.2-0.20220408131628-aac992c27451 h1:xj8gUZGYO3nb5+6Bjw9+tsFkA9sYynrOvDvvC4uDV2I= diff --git a/handler/http/handler.go b/handler/http/handler.go index 9117ef31..4f19849a 100644 --- a/handler/http/handler.go +++ b/handler/http/handler.go @@ -19,13 +19,15 @@ import ( "github.com/asaskevich/govalidator" "github.com/go-gost/core/handler" - traffic "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" md "github.com/go-gost/core/metadata" "github.com/go-gost/core/observer/stats" ctxvalue "github.com/go-gost/x/ctx" xio "github.com/go-gost/x/internal/io" netpkg "github.com/go-gost/x/internal/net" + limiter_util "github.com/go-gost/x/internal/util/limiter" stats_util "github.com/go-gost/x/internal/util/stats" traffic_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" stats_wrapper "github.com/go-gost/x/observer/stats/wrapper" @@ -40,6 +42,7 @@ type httpHandler struct { md metadata options handler.Options stats *stats_util.HandlerStats + limiter traffic.TrafficLimiter cancel context.CancelFunc } @@ -51,7 +54,6 @@ func NewHandler(opts ...handler.Option) handler.Handler { return &httpHandler{ options: options, - stats: stats_util.NewHandlerStats(options.Service), } } @@ -64,9 +66,14 @@ func (h *httpHandler) Init(md md.Metadata) error { h.cancel = cancel if h.options.Observer != nil { + h.stats = stats_util.NewHandlerStats(h.options.Service) go h.observeStats(ctx) } + if limiter := h.options.Limiter; limiter != nil { + h.limiter = limiter_util.NewCachedTrafficLimiter(limiter, 30*time.Second, 60*time.Second) + } + return nil } @@ -222,11 +229,16 @@ func (h *httpHandler) handleRequest(ctx context.Context, conn net.Conn, req *htt } defer cc.Close() - rw := traffic_wrapper.WrapReadWriter(h.options.Limiter, conn, - traffic.NetworkOption(network), - traffic.AddrOption(addr), - traffic.ClientOption(clientID), - traffic.SrcOption(conn.RemoteAddr().String()), + rw := traffic_wrapper.WrapReadWriter( + h.limiter, + conn, + clientID, + limiter.ScopeOption(limiter.ScopeClient), + limiter.ServiceOption(h.options.Service), + limiter.NetworkOption(network), + limiter.AddrOption(addr), + limiter.ClientOption(clientID), + limiter.SrcOption(conn.RemoteAddr().String()), ) if h.options.Observer != nil { pstats := h.stats.Stats(clientID) diff --git a/handler/http2/handler.go b/handler/http2/handler.go index def7ad1b..55500378 100644 --- a/handler/http2/handler.go +++ b/handler/http2/handler.go @@ -19,6 +19,7 @@ import ( "time" "github.com/go-gost/core/handler" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/limiter/traffic" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" @@ -26,8 +27,9 @@ import ( ctxvalue "github.com/go-gost/x/ctx" xio "github.com/go-gost/x/internal/io" netpkg "github.com/go-gost/x/internal/net" + limiter_util "github.com/go-gost/x/internal/util/limiter" stats_util "github.com/go-gost/x/internal/util/stats" - "github.com/go-gost/x/limiter/traffic/wrapper" + traffic_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" stats_wrapper "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" ) @@ -40,6 +42,7 @@ type http2Handler struct { md metadata options handler.Options stats *stats_util.HandlerStats + limiter traffic.TrafficLimiter cancel context.CancelFunc } @@ -51,7 +54,6 @@ func NewHandler(opts ...handler.Option) handler.Handler { return &http2Handler{ options: options, - stats: stats_util.NewHandlerStats(options.Service), } } @@ -64,8 +66,14 @@ func (h *http2Handler) Init(md md.Metadata) error { h.cancel = cancel if h.options.Observer != nil { + h.stats = stats_util.NewHandlerStats(h.options.Service) go h.observeStats(ctx) } + + if limiter := h.options.Limiter; limiter != nil { + h.limiter = limiter_util.NewCachedTrafficLimiter(limiter, 30*time.Second, 60*time.Second) + } + return nil } @@ -216,11 +224,16 @@ func (h *http2Handler) roundTrip(ctx context.Context, w http.ResponseWriter, req return nil } - rw := wrapper.WrapReadWriter(h.options.Limiter, xio.NewReadWriter(req.Body, flushWriter{w}), - traffic.NetworkOption("tcp"), - traffic.AddrOption(addr), - traffic.ClientOption(clientID), - traffic.SrcOption(req.RemoteAddr), + rw := traffic_wrapper.WrapReadWriter( + h.limiter, + xio.NewReadWriter(req.Body, flushWriter{w}), + clientID, + limiter.ScopeOption(limiter.ScopeClient), + limiter.ServiceOption(h.options.Service), + limiter.NetworkOption("tcp"), + limiter.AddrOption(addr), + limiter.ClientOption(clientID), + limiter.SrcOption(req.RemoteAddr), ) if h.options.Observer != nil { pstats := h.stats.Stats(clientID) diff --git a/handler/relay/bind.go b/handler/relay/bind.go index 807655ff..cd2a08d0 100644 --- a/handler/relay/bind.go +++ b/handler/relay/bind.go @@ -92,6 +92,7 @@ func (h *relayHandler) bindTCP(ctx context.Context, conn net.Conn, network, addr epListener := newTCPListener(ln, listener.AddrOption(address), listener.ServiceOption(serviceName), + listener.TrafficLimiterOption(h.options.Limiter), listener.LoggerOption(log.WithFields(map[string]any{ "kind": "listener", })), diff --git a/handler/relay/connect.go b/handler/relay/connect.go index f9567177..34fecf70 100644 --- a/handler/relay/connect.go +++ b/handler/relay/connect.go @@ -8,14 +8,14 @@ import ( "net" "time" - "github.com/go-gost/core/limiter/traffic" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/logger" "github.com/go-gost/core/observer/stats" "github.com/go-gost/relay" ctxvalue "github.com/go-gost/x/ctx" xnet "github.com/go-gost/x/internal/net" serial "github.com/go-gost/x/internal/util/serial" - "github.com/go-gost/x/limiter/traffic/wrapper" + traffic_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" stats_wrapper "github.com/go-gost/x/observer/stats/wrapper" ) @@ -108,11 +108,16 @@ func (h *relayHandler) handleConnect(ctx context.Context, conn net.Conn, network } clientID := ctxvalue.ClientIDFromContext(ctx) - rw := wrapper.WrapReadWriter(h.options.Limiter, conn, - traffic.NetworkOption(network), - traffic.AddrOption(address), - traffic.ClientOption(string(clientID)), - traffic.SrcOption(conn.RemoteAddr().String()), + rw := traffic_wrapper.WrapReadWriter( + h.limiter, + conn, + string(clientID), + limiter.ScopeOption(limiter.ScopeClient), + limiter.ServiceOption(h.options.Service), + limiter.NetworkOption(network), + limiter.AddrOption(address), + limiter.ClientOption(string(clientID)), + limiter.SrcOption(conn.RemoteAddr().String()), ) if h.options.Observer != nil { pstats := h.stats.Stats(string(clientID)) diff --git a/handler/relay/entrypoint.go b/handler/relay/entrypoint.go index 1c3030ab..0814ed24 100644 --- a/handler/relay/entrypoint.go +++ b/handler/relay/entrypoint.go @@ -13,8 +13,6 @@ import ( xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/proxyproto" "github.com/go-gost/x/internal/util/mux" - climiter "github.com/go-gost/x/limiter/conn/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" ) @@ -40,8 +38,6 @@ func (l *tcpListener) Init(md md.Metadata) (err error) { ln = proxyproto.WrapListener(l.options.ProxyProtocol, ln, 10*time.Second) ln = metrics.WrapListener(l.options.Service, ln) ln = admission.WrapListener(l.options.Admission, ln) - ln = limiter.WrapListener(l.options.TrafficLimiter, ln) - ln = climiter.WrapListener(l.options.ConnLimiter, ln) l.ln = ln return diff --git a/handler/relay/forward.go b/handler/relay/forward.go index 9f4f08f0..ba9824ec 100644 --- a/handler/relay/forward.go +++ b/handler/relay/forward.go @@ -7,7 +7,7 @@ import ( "net" "time" - "github.com/go-gost/core/limiter/traffic" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/logger" "github.com/go-gost/core/observer/stats" "github.com/go-gost/relay" @@ -90,11 +90,16 @@ func (h *relayHandler) handleForward(ctx context.Context, conn net.Conn, network } clientID := ctxvalue.ClientIDFromContext(ctx) - rw := wrapper.WrapReadWriter(h.options.Limiter, conn, - traffic.NetworkOption(network), - traffic.AddrOption(target.Addr), - traffic.ClientOption(string(clientID)), - traffic.SrcOption(conn.RemoteAddr().String()), + rw := wrapper.WrapReadWriter( + h.limiter, + conn, + string(clientID), + limiter.ServiceOption(h.options.Service), + limiter.ScopeOption(limiter.ScopeClient), + limiter.NetworkOption(network), + limiter.AddrOption(target.Addr), + limiter.ClientOption(string(clientID)), + limiter.SrcOption(conn.RemoteAddr().String()), ) if h.options.Observer != nil { pstats := h.stats.Stats(string(clientID)) diff --git a/handler/relay/handler.go b/handler/relay/handler.go index c86c02c7..ca9c3c81 100644 --- a/handler/relay/handler.go +++ b/handler/relay/handler.go @@ -9,9 +9,11 @@ import ( "github.com/go-gost/core/handler" "github.com/go-gost/core/hop" + "github.com/go-gost/core/limiter/traffic" md "github.com/go-gost/core/metadata" "github.com/go-gost/relay" ctxvalue "github.com/go-gost/x/ctx" + limiter_util "github.com/go-gost/x/internal/util/limiter" stats_util "github.com/go-gost/x/internal/util/stats" "github.com/go-gost/x/registry" ) @@ -32,6 +34,7 @@ type relayHandler struct { md metadata options handler.Options stats *stats_util.HandlerStats + limiter traffic.TrafficLimiter cancel context.CancelFunc } @@ -43,7 +46,6 @@ func NewHandler(opts ...handler.Option) handler.Handler { return &relayHandler{ options: options, - stats: stats_util.NewHandlerStats(options.Service), } } @@ -56,9 +58,14 @@ func (h *relayHandler) Init(md md.Metadata) (err error) { h.cancel = cancel if h.options.Observer != nil { + h.stats = stats_util.NewHandlerStats(h.options.Service) go h.observeStats(ctx) } + if limiter := h.options.Limiter; limiter != nil { + h.limiter = limiter_util.NewCachedTrafficLimiter(limiter, 30*time.Second, 60*time.Second) + } + return nil } diff --git a/handler/socks/v4/handler.go b/handler/socks/v4/handler.go index 2637b920..d03c7b5f 100644 --- a/handler/socks/v4/handler.go +++ b/handler/socks/v4/handler.go @@ -7,6 +7,7 @@ import ( "time" "github.com/go-gost/core/handler" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/limiter/traffic" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" @@ -14,8 +15,9 @@ import ( "github.com/go-gost/gosocks4" ctxvalue "github.com/go-gost/x/ctx" netpkg "github.com/go-gost/x/internal/net" + limiter_util "github.com/go-gost/x/internal/util/limiter" stats_util "github.com/go-gost/x/internal/util/stats" - "github.com/go-gost/x/limiter/traffic/wrapper" + traffic_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" stats_wrapper "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" ) @@ -34,6 +36,7 @@ type socks4Handler struct { md metadata options handler.Options stats *stats_util.HandlerStats + limiter traffic.TrafficLimiter cancel context.CancelFunc } @@ -45,7 +48,6 @@ func NewHandler(opts ...handler.Option) handler.Handler { return &socks4Handler{ options: options, - stats: stats_util.NewHandlerStats(options.Service), } } @@ -58,9 +60,14 @@ func (h *socks4Handler) Init(md md.Metadata) (err error) { h.cancel = cancel if h.options.Observer != nil { + h.stats = stats_util.NewHandlerStats(h.options.Service) go h.observeStats(ctx) } + if limiter := h.options.Limiter; limiter != nil { + h.limiter = limiter_util.NewCachedTrafficLimiter(limiter, 30*time.Second, 60*time.Second) + } + return nil } @@ -165,11 +172,16 @@ func (h *socks4Handler) handleConnect(ctx context.Context, conn net.Conn, req *g } clientID := ctxvalue.ClientIDFromContext(ctx) - rw := wrapper.WrapReadWriter(h.options.Limiter, conn, - traffic.NetworkOption("tcp"), - traffic.AddrOption(addr), - traffic.ClientOption(string(clientID)), - traffic.SrcOption(conn.RemoteAddr().String()), + rw := traffic_wrapper.WrapReadWriter( + h.limiter, + conn, + string(clientID), + limiter.ScopeOption(limiter.ScopeClient), + limiter.ServiceOption(h.options.Service), + limiter.NetworkOption("tcp"), + limiter.AddrOption(addr), + limiter.ClientOption(string(clientID)), + limiter.SrcOption(conn.RemoteAddr().String()), ) if h.options.Observer != nil { pstats := h.stats.Stats(string(clientID)) diff --git a/handler/socks/v5/connect.go b/handler/socks/v5/connect.go index 6c2a266d..b532439a 100644 --- a/handler/socks/v5/connect.go +++ b/handler/socks/v5/connect.go @@ -6,13 +6,13 @@ import ( "net" "time" - "github.com/go-gost/core/limiter/traffic" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/logger" "github.com/go-gost/core/observer/stats" "github.com/go-gost/gosocks5" ctxvalue "github.com/go-gost/x/ctx" netpkg "github.com/go-gost/x/internal/net" - "github.com/go-gost/x/limiter/traffic/wrapper" + traffic_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" stats_wrapper "github.com/go-gost/x/observer/stats/wrapper" ) @@ -53,11 +53,16 @@ func (h *socks5Handler) handleConnect(ctx context.Context, conn net.Conn, networ } clientID := ctxvalue.ClientIDFromContext(ctx) - rw := wrapper.WrapReadWriter(h.options.Limiter, conn, - traffic.NetworkOption(network), - traffic.AddrOption(address), - traffic.ClientOption(string(clientID)), - traffic.SrcOption(conn.RemoteAddr().String()), + rw := traffic_wrapper.WrapReadWriter( + h.limiter, + conn, + string(clientID), + limiter.ServiceOption(h.options.Service), + limiter.ScopeOption(limiter.ScopeClient), + limiter.NetworkOption(network), + limiter.AddrOption(address), + limiter.ClientOption(string(clientID)), + limiter.SrcOption(conn.RemoteAddr().String()), ) if h.options.Observer != nil { pstats := h.stats.Stats(string(clientID)) diff --git a/handler/socks/v5/handler.go b/handler/socks/v5/handler.go index f24358c5..3374f5fb 100644 --- a/handler/socks/v5/handler.go +++ b/handler/socks/v5/handler.go @@ -7,9 +7,11 @@ import ( "time" "github.com/go-gost/core/handler" + "github.com/go-gost/core/limiter/traffic" md "github.com/go-gost/core/metadata" "github.com/go-gost/gosocks5" ctxvalue "github.com/go-gost/x/ctx" + limiter_util "github.com/go-gost/x/internal/util/limiter" "github.com/go-gost/x/internal/util/socks" stats_util "github.com/go-gost/x/internal/util/stats" "github.com/go-gost/x/registry" @@ -29,6 +31,7 @@ type socks5Handler struct { md metadata options handler.Options stats *stats_util.HandlerStats + limiter traffic.TrafficLimiter cancel context.CancelFunc } @@ -40,7 +43,6 @@ func NewHandler(opts ...handler.Option) handler.Handler { return &socks5Handler{ options: options, - stats: stats_util.NewHandlerStats(options.Service), } } @@ -60,9 +62,14 @@ func (h *socks5Handler) Init(md md.Metadata) (err error) { h.cancel = cancel if h.options.Observer != nil { + h.stats = stats_util.NewHandlerStats(h.options.Service) go h.observeStats(ctx) } + if limiter := h.options.Limiter; limiter != nil { + h.limiter = limiter_util.NewCachedTrafficLimiter(limiter, 30*time.Second, 60*time.Second) + } + return } diff --git a/handler/tunnel/bind.go b/handler/tunnel/bind.go index a5bb319f..5d76c3ef 100644 --- a/handler/tunnel/bind.go +++ b/handler/tunnel/bind.go @@ -65,7 +65,15 @@ func (h *tunnelHandler) handleBind(ctx context.Context, conn net.Conn, network, if h.stats != nil { stats = h.stats.Stats(tunnelID.String()) } - h.pool.Add(tunnelID, NewConnector(connectorID, tunnelID, h.id, session, h.md.sd, stats), h.md.tunnelTTL) + + c := NewConnector(connectorID, tunnelID, h.id, session, &ConnectorOptions{ + service: h.options.Service, + sd: h.md.sd, + stats: stats, + limiter: h.limiter, + }) + + h.pool.Add(tunnelID, c, h.md.tunnelTTL) if h.md.ingress != nil { h.md.ingress.SetRule(ctx, &ingress.Rule{ Hostname: endpoint, diff --git a/handler/tunnel/connect.go b/handler/tunnel/connect.go index eb798c64..21a22b47 100644 --- a/handler/tunnel/connect.go +++ b/handler/tunnel/connect.go @@ -6,12 +6,9 @@ import ( "net" "time" - "github.com/go-gost/core/limiter/traffic" "github.com/go-gost/core/logger" "github.com/go-gost/relay" - ctxvalue "github.com/go-gost/x/ctx" xnet "github.com/go-gost/x/internal/net" - "github.com/go-gost/x/limiter/traffic/wrapper" ) func (h *tunnelHandler) handleConnect(ctx context.Context, req *relay.Request, conn net.Conn, network, srcAddr string, dstAddr string, tunnelID relay.TunnelID, log logger.Logger) error { @@ -100,16 +97,9 @@ func (h *tunnelHandler) handleConnect(ctx context.Context, req *relay.Request, c req.WriteTo(cc) } - rw := wrapper.WrapReadWriter(h.options.Limiter, conn, - traffic.NetworkOption(network), - traffic.AddrOption(dstAddr), - traffic.ClientOption(string(ctxvalue.ClientIDFromContext(ctx))), - traffic.SrcOption(conn.RemoteAddr().String()), - ) - t := time.Now() log.Debugf("%s <-> %s", conn.RemoteAddr(), cc.RemoteAddr()) - xnet.Transport(rw, cc) + xnet.Transport(conn, cc) log.WithFields(map[string]any{ "duration": time.Since(t), }).Debugf("%s >-< %s", conn.RemoteAddr(), cc.RemoteAddr()) diff --git a/handler/tunnel/entrypoint.go b/handler/tunnel/entrypoint.go index 17aec069..90859ee8 100644 --- a/handler/tunnel/entrypoint.go +++ b/handler/tunnel/entrypoint.go @@ -24,7 +24,6 @@ import ( xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/proxyproto" climiter "github.com/go-gost/x/limiter/conn/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" ) @@ -374,7 +373,6 @@ func (l *tcpListener) Init(md md.Metadata) (err error) { ln = proxyproto.WrapListener(l.options.ProxyProtocol, ln, 10*time.Second) ln = metrics.WrapListener(l.options.Service, ln) ln = admission.WrapListener(l.options.Admission, ln) - ln = limiter.WrapListener(l.options.TrafficLimiter, ln) ln = climiter.WrapListener(l.options.ConnLimiter, ln) l.ln = ln diff --git a/handler/tunnel/handler.go b/handler/tunnel/handler.go index aebc3dc5..18dad0e3 100644 --- a/handler/tunnel/handler.go +++ b/handler/tunnel/handler.go @@ -9,6 +9,7 @@ import ( "time" "github.com/go-gost/core/handler" + "github.com/go-gost/core/limiter/traffic" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" @@ -17,6 +18,7 @@ import ( "github.com/go-gost/relay" ctxvalue "github.com/go-gost/x/ctx" xnet "github.com/go-gost/x/internal/net" + limiter_util "github.com/go-gost/x/internal/util/limiter" stats_util "github.com/go-gost/x/internal/util/stats" xrecorder "github.com/go-gost/x/recorder" "github.com/go-gost/x/registry" @@ -47,6 +49,7 @@ type tunnelHandler struct { md metadata log logger.Logger stats *stats_util.HandlerStats + limiter traffic.TrafficLimiter cancel context.CancelFunc } @@ -58,7 +61,6 @@ func NewHandler(opts ...handler.Option) handler.Handler { return &tunnelHandler{ options: options, - stats: stats_util.NewHandlerStats(options.Service), } } @@ -105,9 +107,14 @@ func (h *tunnelHandler) Init(md md.Metadata) (err error) { h.cancel = cancel if h.options.Observer != nil { + h.stats = stats_util.NewHandlerStats(h.options.Service) go h.observeStats(ctx) } + if limiter := h.options.Limiter; limiter != nil { + h.limiter = limiter_util.NewCachedTrafficLimiter(limiter, 30*time.Second, 60*time.Second) + } + return nil } diff --git a/handler/tunnel/tunnel.go b/handler/tunnel/tunnel.go index fe6bc01d..0c31fffd 100644 --- a/handler/tunnel/tunnel.go +++ b/handler/tunnel/tunnel.go @@ -6,12 +6,15 @@ import ( "sync" "time" + "github.com/go-gost/core/limiter" + "github.com/go-gost/core/limiter/traffic" "github.com/go-gost/core/logger" "github.com/go-gost/core/sd" "github.com/go-gost/relay" "github.com/go-gost/x/internal/util/mux" "github.com/go-gost/core/observer/stats" + traffic_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" stats_wrapper "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/selector" "github.com/google/uuid" @@ -21,25 +24,34 @@ const ( MaxWeight uint8 = 0xff ) -type Connector struct { - id relay.ConnectorID - tid relay.TunnelID - node string - sd sd.SD - t time.Time - s *mux.Session - stats *stats.Stats +type ConnectorOptions struct { + service string + sd sd.SD + stats *stats.Stats + limiter traffic.TrafficLimiter } -func NewConnector(id relay.ConnectorID, tid relay.TunnelID, node string, s *mux.Session, sd sd.SD, stats *stats.Stats) *Connector { +type Connector struct { + id relay.ConnectorID + tid relay.TunnelID + node string + s *mux.Session + t time.Time + opts *ConnectorOptions +} + +func NewConnector(id relay.ConnectorID, tid relay.TunnelID, node string, s *mux.Session, opts *ConnectorOptions) *Connector { + if opts == nil { + opts = &ConnectorOptions{} + } + c := &Connector{ - id: id, - tid: tid, - node: node, - sd: sd, - stats: stats, - t: time.Now(), - s: s, + id: id, + tid: tid, + node: node, + s: s, + t: time.Now(), + opts: opts, } go c.accept() return c @@ -51,8 +63,8 @@ func (c *Connector) accept() { if err != nil { logger.Default().Errorf("connector %s: %v", c.id, err) c.s.Close() - if c.sd != nil { - c.sd.Deregister(context.Background(), &sd.Service{ + if c.opts.sd != nil { + c.opts.sd.Deregister(context.Background(), &sd.Service{ ID: c.id.String(), Name: c.tid.String(), Node: c.node, @@ -78,9 +90,22 @@ func (c *Connector) GetConn() (net.Conn, error) { return nil, err } - if c.stats != nil { - conn = stats_wrapper.WrapConn(conn, c.stats) + conn = stats_wrapper.WrapConn(conn, c.opts.stats) + + network := "tcp" + if c.id.IsUDP() { + network = "udp" } + conn = traffic_wrapper.WrapConn( + conn, + c.opts.limiter, + c.tid.String(), + limiter.ScopeOption(limiter.ScopeClient), + limiter.ServiceOption(c.opts.service), + limiter.ClientOption(c.tid.String()), + limiter.NetworkOption(network), + limiter.SrcOption(conn.RemoteAddr().String()), + ) return conn, nil } diff --git a/internal/util/limiter/cache.go b/internal/util/limiter/cache.go new file mode 100644 index 00000000..b2174e58 --- /dev/null +++ b/internal/util/limiter/cache.go @@ -0,0 +1,64 @@ +package limiter + +import ( + "sync" + "time" +) + +type Item struct { + v interface{} + expiration int64 +} + +func NewItem(v interface{}, d time.Duration) *Item { + var expiration int64 + if d > 0 { + expiration = time.Now().Add(d).UnixNano() + } + + return &Item{ + v: v, + expiration: expiration, + } +} + +func (p *Item) Expired() bool { + if p == nil { + return true + } + return p.expiration > 0 && time.Now().UnixNano() > p.expiration +} + +func (p *Item) Value() interface{} { + if p == nil { + return nil + } + return p.v +} + +type Cache struct { + items map[string]*Item + cleanupInterval time.Duration + mu sync.RWMutex +} + +func NewCache(cleanupInterval time.Duration) *Cache { + return &Cache{ + cleanupInterval: cleanupInterval, + items: make(map[string]*Item), + } +} + +func (c *Cache) Set(key string, item *Item) { + c.mu.Lock() + defer c.mu.Unlock() + + c.items[key] = item +} + +func (c *Cache) Get(key string) *Item { + c.mu.RLock() + defer c.mu.RUnlock() + + return c.items[key] +} diff --git a/internal/util/limiter/limiter.go b/internal/util/limiter/limiter.go new file mode 100644 index 00000000..65c47748 --- /dev/null +++ b/internal/util/limiter/limiter.go @@ -0,0 +1,92 @@ +package limiter + +import ( + "context" + "time" + + "github.com/go-gost/core/limiter" + "github.com/go-gost/core/limiter/traffic" +) + +type cachedTrafficLimiter struct { + inLimits *Cache + outLimits *Cache + limiter traffic.TrafficLimiter + ttl time.Duration +} + +func NewCachedTrafficLimiter(limiter traffic.TrafficLimiter, ttl time.Duration, cleanupInterval time.Duration) traffic.TrafficLimiter { + if limiter == nil { + return nil + } + + lim := &cachedTrafficLimiter{ + inLimits: NewCache(cleanupInterval), + outLimits: NewCache(cleanupInterval), + limiter: limiter, + ttl: ttl, + } + return lim +} + +func (p *cachedTrafficLimiter) In(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter { + if p.limiter == nil { + return nil + } + + item := p.inLimits.Get(key) + lim, _ := item.Value().(traffic.Limiter) + if !item.Expired() { + return lim + } + + limNew := p.limiter.In(ctx, key, opts...) + if limNew == nil { + limNew = lim + } + if item == nil || !p.equal(lim, limNew) { + p.inLimits.Set(key, NewItem(limNew, p.ttl)) + return limNew + } + + p.inLimits.Set(key, NewItem(lim, p.ttl)) + + return lim +} + +func (p *cachedTrafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter { + if p.limiter == nil { + return nil + } + + item := p.outLimits.Get(key) + lim, _ := item.Value().(traffic.Limiter) + if !item.Expired() { + return lim + } + + limNew := p.limiter.Out(ctx, key, opts...) + if limNew == nil { + limNew = lim + } + if item == nil || !p.equal(lim, limNew) { + p.outLimits.Set(key, NewItem(limNew, p.ttl)) + return limNew + } + + p.outLimits.Set(key, NewItem(lim, p.ttl)) + + return lim +} + +func (p *cachedTrafficLimiter) equal(lim1, lim2 traffic.Limiter) bool { + if lim1 == lim2 { + return true + } + + if lim1 == nil || lim2 == nil { + return false + } + + return lim1.Limit() == lim2.Limit() +} diff --git a/limiter/traffic/plugin/grpc.go b/limiter/traffic/plugin/grpc.go index c05a5e40..4a888acf 100644 --- a/limiter/traffic/plugin/grpc.go +++ b/limiter/traffic/plugin/grpc.go @@ -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, diff --git a/limiter/traffic/plugin/http.go b/limiter/traffic/plugin/http.go index 7b85ffd1..28429407 100644 --- a/limiter/traffic/plugin/http.go +++ b/limiter/traffic/plugin/http.go @@ -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)) diff --git a/limiter/traffic/traffic.go b/limiter/traffic/traffic.go index 4349daa4..2c68cfe0 100644 --- a/limiter/traffic/traffic.go +++ b/limiter/traffic/traffic.go @@ -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) } diff --git a/limiter/traffic/wrapper/conn.go b/limiter/traffic/wrapper/conn.go index 67a84b20..07979661 100644 --- a/limiter/traffic/wrapper/conn.go +++ b/limiter/traffic/wrapper/conn.go @@ -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 } diff --git a/limiter/traffic/wrapper/io.go b/limiter/traffic/wrapper/io.go index 86923d2a..4dbb5ce9 100644 --- a/limiter/traffic/wrapper/io.go +++ b/limiter/traffic/wrapper/io.go @@ -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) } diff --git a/limiter/traffic/wrapper/listener.go b/limiter/traffic/wrapper/listener.go index 0e4ec5ba..9b579ce7 100644 --- a/limiter/traffic/wrapper/listener.go +++ b/limiter/traffic/wrapper/listener.go @@ -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 } diff --git a/listener/dns/listener.go b/listener/dns/listener.go index a6276796..d5f1d115 100644 --- a/listener/dns/listener.go +++ b/listener/dns/listener.go @@ -10,13 +10,16 @@ import ( "net" "net/http" "strings" + "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_util "github.com/go-gost/x/internal/util/limiter" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -76,6 +79,13 @@ func (l *dnsListener) Init(md md.Metadata) (err error) { if err != nil { return } + + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) + l.server = &dnsServer{ server: &dns.Server{ Net: "tcp", @@ -111,6 +121,12 @@ func (l *dnsListener) Init(md md.Metadata) (err error) { } ln = tls.NewListener(ln, l.options.TLSConfig) + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) + l.server = &dnsServer{ server: &dns.Server{ Net: "tcp-tls", @@ -147,6 +163,12 @@ func (l *dnsListener) Init(md md.Metadata) (err error) { } ln = tls.NewListener(ln, l.options.TLSConfig) + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) + l.server = &dohServer{ addr: l.options.Addr, tlsConfig: l.options.TLSConfig, @@ -181,6 +203,15 @@ func (l *dnsListener) Init(md md.Metadata) (err error) { return } + pc = limiter_wrapper.WrapPacketConn( + pc, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + "", + limiter.ScopeOption(limiter.ScopeService), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(network), + ) + l.server = &dnsServer{ server: &dns.Server{ Net: "udp", @@ -218,7 +249,15 @@ func (l *dnsListener) Accept() (conn net.Conn, err error) { conn = metrics.WrapConn(l.options.Service, conn) conn = stats.WrapConn(conn, l.options.Stats) conn = admission.WrapConn(l.options.Admission, conn) - conn = limiter.WrapConn(l.options.TrafficLimiter, conn) + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) case err, ok = <-l.errChan: if !ok { err = listener.ErrClosed diff --git a/listener/dtls/listener.go b/listener/dtls/listener.go index 228cf4d6..eaa285f3 100644 --- a/listener/dtls/listener.go +++ b/listener/dtls/listener.go @@ -6,6 +6,7 @@ import ( "net" "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" @@ -13,8 +14,9 @@ import ( xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/proxyproto" xdtls "github.com/go-gost/x/internal/util/dtls" + limiter_util "github.com/go-gost/x/internal/util/limiter" climiter "github.com/go-gost/x/limiter/conn/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -82,7 +84,11 @@ func (l *dtlsListener) Init(md md.Metadata) (err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = stats.WrapListener(ln, l.options.Stats) ln = admission.WrapListener(l.options.Admission, ln) - ln = limiter.WrapListener(l.options.TrafficLimiter, ln) + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) ln = climiter.WrapListener(l.options.ConnLimiter, ln) l.ln = ln @@ -96,6 +102,15 @@ func (l *dtlsListener) Accept() (conn net.Conn, err error) { return } conn = xdtls.Conn(c, l.md.bufferSize) + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) return } diff --git a/listener/ftcp/listener.go b/listener/ftcp/listener.go index 6e60e2c1..6b691cf6 100644 --- a/listener/ftcp/listener.go +++ b/listener/ftcp/listener.go @@ -2,14 +2,17 @@ package ftcp import ( "net" + "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/udp" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_util "github.com/go-gost/x/internal/util/limiter" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -55,7 +58,14 @@ func (l *ftcpListener) Init(md md.Metadata) (err error) { conn = metrics.WrapPacketConn(l.options.Service, conn) conn = stats.WrapPacketConn(conn, l.options.Stats) conn = admission.WrapPacketConn(l.options.Admission, conn) - conn = limiter.WrapPacketConn(l.options.TrafficLimiter, conn) + conn = limiter_wrapper.WrapPacketConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + "", + limiter.ScopeOption(limiter.ScopeService), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + ) l.ln = udp.NewListener( conn, @@ -71,7 +81,22 @@ func (l *ftcpListener) Init(md md.Metadata) (err error) { } func (l *ftcpListener) Accept() (conn net.Conn, err error) { - return l.ln.Accept() + conn, err = l.ln.Accept() + if err != nil { + return + } + + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) + + return } func (l *ftcpListener) Addr() net.Addr { diff --git a/listener/grpc/listener.go b/listener/grpc/listener.go index 1bf16d63..7d3f80f6 100644 --- a/listener/grpc/listener.go +++ b/listener/grpc/listener.go @@ -5,6 +5,7 @@ import ( "net" "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" @@ -12,8 +13,9 @@ import ( xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/proxyproto" pb "github.com/go-gost/x/internal/util/grpc/proto" + limiter_util "github.com/go-gost/x/internal/util/limiter" climiter "github.com/go-gost/x/limiter/conn/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -69,7 +71,11 @@ func (l *grpcListener) Init(md md.Metadata) (err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = stats.WrapListener(ln, l.options.Stats) ln = admission.WrapListener(l.options.Admission, ln) - ln = limiter.WrapListener(l.options.TrafficLimiter, ln) + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) ln = climiter.WrapListener(l.options.ConnLimiter, ln) var opts []grpc.ServerOption @@ -116,6 +122,15 @@ func (l *grpcListener) Accept() (conn net.Conn, err error) { var ok bool select { case conn = <-l.cqueue: + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) case err, ok = <-l.errChan: if !ok { err = listener.ErrClosed diff --git a/listener/http2/h2/listener.go b/listener/http2/h2/listener.go index aa2ad57d..22c90e9c 100644 --- a/listener/http2/h2/listener.go +++ b/listener/http2/h2/listener.go @@ -9,14 +9,16 @@ import ( "net/http/httputil" "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/proxyproto" + limiter_util "github.com/go-gost/x/internal/util/limiter" climiter "github.com/go-gost/x/limiter/conn/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -90,7 +92,11 @@ func (l *h2Listener) Init(md md.Metadata) (err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = stats.WrapListener(ln, l.options.Stats) ln = admission.WrapListener(l.options.Admission, ln) - ln = limiter.WrapListener(l.options.TrafficLimiter, ln) + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) ln = climiter.WrapListener(l.options.ConnLimiter, ln) if l.h2c { @@ -122,6 +128,15 @@ func (l *h2Listener) Accept() (conn net.Conn, err error) { var ok bool select { case conn = <-l.cqueue: + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) case err, ok = <-l.errChan: if !ok { err = listener.ErrClosed diff --git a/listener/http2/listener.go b/listener/http2/listener.go index 682a3677..538ae4ac 100644 --- a/listener/http2/listener.go +++ b/listener/http2/listener.go @@ -7,14 +7,16 @@ import ( "net/http" "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/proxyproto" + limiter_util "github.com/go-gost/x/internal/util/limiter" climiter "github.com/go-gost/x/limiter/conn/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" mdx "github.com/go-gost/x/metadata" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" @@ -79,7 +81,11 @@ func (l *http2Listener) Init(md md.Metadata) (err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = stats.WrapListener(ln, l.options.Stats) ln = admission.WrapListener(l.options.Admission, ln) - ln = limiter.WrapListener(l.options.TrafficLimiter, ln) + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) ln = climiter.WrapListener(l.options.ConnLimiter, ln) ln = tls.NewListener( @@ -103,6 +109,15 @@ func (l *http2Listener) Accept() (conn net.Conn, err error) { var ok bool select { case conn = <-l.cqueue: + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) case err, ok = <-l.errChan: if !ok { err = listener.ErrClosed diff --git a/listener/http3/h3/listener.go b/listener/http3/h3/listener.go index 0ebef93c..6c03a904 100644 --- a/listener/http3/h3/listener.go +++ b/listener/http3/h3/listener.go @@ -2,14 +2,17 @@ package h3 import ( "net" + "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" + limiter_util "github.com/go-gost/x/internal/util/limiter" pht_util "github.com/go-gost/x/internal/util/pht" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -88,7 +91,15 @@ func (l *http3Listener) Accept() (conn net.Conn, err error) { conn = metrics.WrapConn(l.options.Service, conn) conn = stats.WrapConn(conn, l.options.Stats) conn = admission.WrapConn(l.options.Admission, conn) - conn = limiter.WrapConn(l.options.TrafficLimiter, conn) + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) return conn, nil } diff --git a/listener/http3/wt/listener.go b/listener/http3/wt/listener.go index 8a0d0a92..0c3e778e 100644 --- a/listener/http3/wt/listener.go +++ b/listener/http3/wt/listener.go @@ -4,14 +4,17 @@ import ( "net" "net/http" "net/http/httputil" + "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" + limiter_util "github.com/go-gost/x/internal/util/limiter" wt_util "github.com/go-gost/x/internal/util/wt" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -65,11 +68,24 @@ func (l *wtListener) Init(md md.Metadata) (err error) { } l.addr = laddr - pc, err := net.ListenUDP(network, laddr) + var pc net.PacketConn + pc, err = net.ListenUDP(network, laddr) if err != nil { return } + pc = metrics.WrapPacketConn(l.options.Service, pc) + pc = stats.WrapPacketConn(pc, l.options.Stats) + pc = admission.WrapPacketConn(l.options.Admission, pc) + pc = limiter_wrapper.WrapPacketConn( + pc, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + "", + limiter.ScopeOption(limiter.ScopeService), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(network), + ) + mux := http.NewServeMux() mux.Handle(l.md.path, http.HandlerFunc(l.upgrade)) @@ -112,10 +128,15 @@ func (l *wtListener) Accept() (conn net.Conn, err error) { var ok bool select { case conn = <-l.cqueue: - conn = metrics.WrapConn(l.options.Service, conn) - conn = stats.WrapConn(conn, l.options.Stats) - conn = admission.WrapConn(l.options.Admission, conn) - conn = limiter.WrapConn(l.options.TrafficLimiter, conn) + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) case err, ok = <-l.errChan: if !ok { err = listener.ErrClosed diff --git a/listener/icmp/listener.go b/listener/icmp/listener.go index adc84976..cfbc526b 100644 --- a/listener/icmp/listener.go +++ b/listener/icmp/listener.go @@ -3,13 +3,16 @@ package quic import ( "context" "net" + "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" 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/traffic/wrapper" + limiter_util "github.com/go-gost/x/internal/util/limiter" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -78,7 +81,14 @@ func (l *icmpListener) Init(md md.Metadata) (err error) { conn = metrics.WrapPacketConn(l.options.Service, conn) conn = stats.WrapPacketConn(conn, l.options.Stats) conn = admission.WrapPacketConn(l.options.Admission, conn) - conn = limiter.WrapPacketConn(l.options.TrafficLimiter, conn) + conn = limiter_wrapper.WrapPacketConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + "", + limiter.ScopeOption(limiter.ScopeService), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + ) config := &quic.Config{ KeepAlivePeriod: l.md.keepAlivePeriod, @@ -111,6 +121,15 @@ func (l *icmpListener) Accept() (conn net.Conn, err error) { var ok bool select { case conn = <-l.cqueue: + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) case err, ok = <-l.errChan: if !ok { err = listener.ErrClosed diff --git a/listener/kcp/listener.go b/listener/kcp/listener.go index be2901af..ea4e3fc5 100644 --- a/listener/kcp/listener.go +++ b/listener/kcp/listener.go @@ -4,13 +4,15 @@ import ( "net" "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" 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/traffic/wrapper" + limiter_util "github.com/go-gost/x/internal/util/limiter" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -78,7 +80,14 @@ func (l *kcpListener) Init(md md.Metadata) (err error) { conn = metrics.WrapUDPConn(l.options.Service, conn) conn = stats.WrapUDPConn(conn, l.options.Stats) conn = admission.WrapUDPConn(l.options.Admission, conn) - conn = limiter.WrapUDPConn(l.options.TrafficLimiter, conn) + conn = limiter_wrapper.WrapUDPConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + "", + limiter.ScopeOption(limiter.ScopeService), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + ) ln, err := kcp.ServeConn( kcp_util.BlockCrypt(config.Key, config.Crypt, kcp_util.DefaultSalt), @@ -113,6 +122,15 @@ func (l *kcpListener) Accept() (conn net.Conn, err error) { var ok bool select { case conn = <-l.cqueue: + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) case err, ok = <-l.errChan: if !ok { err = listener.ErrClosed diff --git a/listener/mtcp/listener.go b/listener/mtcp/listener.go index 3d8d7df1..9f867cae 100644 --- a/listener/mtcp/listener.go +++ b/listener/mtcp/listener.go @@ -5,15 +5,17 @@ import ( "net" "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/proxyproto" + limiter_util "github.com/go-gost/x/internal/util/limiter" "github.com/go-gost/x/internal/util/mux" climiter "github.com/go-gost/x/limiter/conn/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -69,7 +71,11 @@ func (l *mtcpListener) Init(md md.Metadata) (err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = stats.WrapListener(ln, l.options.Stats) ln = admission.WrapListener(l.options.Admission, ln) - ln = limiter.WrapListener(l.options.TrafficLimiter, ln) + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) ln = climiter.WrapListener(l.options.ConnLimiter, ln) l.ln = ln @@ -93,6 +99,15 @@ func (l *mtcpListener) Accept() (conn net.Conn, err error) { var ok bool select { case conn = <-l.cqueue: + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) case err, ok = <-l.errChan: if !ok { err = listener.ErrClosed diff --git a/listener/mtls/listener.go b/listener/mtls/listener.go index 5d67a4bd..b43b11f0 100644 --- a/listener/mtls/listener.go +++ b/listener/mtls/listener.go @@ -6,15 +6,17 @@ import ( "net" "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/proxyproto" + limiter_util "github.com/go-gost/x/internal/util/limiter" "github.com/go-gost/x/internal/util/mux" climiter "github.com/go-gost/x/limiter/conn/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -68,7 +70,11 @@ func (l *mtlsListener) Init(md md.Metadata) (err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = stats.WrapListener(ln, l.options.Stats) ln = admission.WrapListener(l.options.Admission, ln) - ln = limiter.WrapListener(l.options.TrafficLimiter, ln) + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) ln = climiter.WrapListener(l.options.ConnLimiter, ln) l.Listener = tls.NewListener(ln, l.options.TLSConfig) @@ -84,6 +90,15 @@ func (l *mtlsListener) Accept() (conn net.Conn, err error) { var ok bool select { case conn = <-l.cqueue: + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) case err, ok = <-l.errChan: if !ok { err = listener.ErrClosed diff --git a/listener/mws/listener.go b/listener/mws/listener.go index c50d80ce..90b4d21e 100644 --- a/listener/mws/listener.go +++ b/listener/mws/listener.go @@ -8,16 +8,18 @@ import ( "net/http/httputil" "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/proxyproto" + limiter_util "github.com/go-gost/x/internal/util/limiter" "github.com/go-gost/x/internal/util/mux" ws_util "github.com/go-gost/x/internal/util/ws" climiter "github.com/go-gost/x/limiter/conn/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -110,7 +112,11 @@ func (l *mwsListener) Init(md md.Metadata) (err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = stats.WrapListener(ln, l.options.Stats) ln = admission.WrapListener(l.options.Admission, ln) - ln = limiter.WrapListener(l.options.TrafficLimiter, ln) + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) ln = climiter.WrapListener(l.options.ConnLimiter, ln) if l.tlsEnabled { @@ -134,6 +140,15 @@ func (l *mwsListener) Accept() (conn net.Conn, err error) { var ok bool select { case conn = <-l.cqueue: + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) case err, ok = <-l.errChan: if !ok { err = listener.ErrClosed diff --git a/listener/obfs/http/listener.go b/listener/obfs/http/listener.go index ccedafc3..8c6710d0 100644 --- a/listener/obfs/http/listener.go +++ b/listener/obfs/http/listener.go @@ -5,14 +5,16 @@ import ( "net" "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/proxyproto" + limiter_util "github.com/go-gost/x/internal/util/limiter" climiter "github.com/go-gost/x/limiter/conn/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -63,7 +65,11 @@ func (l *obfsListener) Init(md md.Metadata) (err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = stats.WrapListener(ln, l.options.Stats) ln = admission.WrapListener(l.options.Admission, ln) - ln = limiter.WrapListener(l.options.TrafficLimiter, ln) + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) ln = climiter.WrapListener(l.options.ConnLimiter, ln) l.Listener = ln @@ -71,13 +77,23 @@ func (l *obfsListener) Init(md md.Metadata) (err error) { } func (l *obfsListener) Accept() (net.Conn, error) { - c, err := l.Listener.Accept() + conn, err := l.Listener.Accept() if err != nil { return nil, err } + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) + return &obfsHTTPConn{ - Conn: c, + Conn: conn, header: l.md.header, logger: l.logger, }, nil diff --git a/listener/obfs/tls/listener.go b/listener/obfs/tls/listener.go index a7a6ccd4..84cd1b0d 100644 --- a/listener/obfs/tls/listener.go +++ b/listener/obfs/tls/listener.go @@ -5,14 +5,16 @@ import ( "net" "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/proxyproto" + limiter_util "github.com/go-gost/x/internal/util/limiter" climiter "github.com/go-gost/x/limiter/conn/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -63,7 +65,11 @@ func (l *obfsListener) Init(md md.Metadata) (err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = stats.WrapListener(ln, l.options.Stats) ln = admission.WrapListener(l.options.Admission, ln) - ln = limiter.WrapListener(l.options.TrafficLimiter, ln) + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) ln = climiter.WrapListener(l.options.ConnLimiter, ln) l.Listener = ln @@ -71,12 +77,22 @@ func (l *obfsListener) Init(md md.Metadata) (err error) { } func (l *obfsListener) Accept() (net.Conn, error) { - c, err := l.Listener.Accept() + conn, err := l.Listener.Accept() if err != nil { return nil, err } + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) + return &obfsTLSConn{ - Conn: c, + Conn: conn, }, nil } diff --git a/listener/pht/listener.go b/listener/pht/listener.go index 1a8ba6f2..cb753089 100644 --- a/listener/pht/listener.go +++ b/listener/pht/listener.go @@ -4,14 +4,17 @@ package pht import ( "net" + "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" + limiter_util "github.com/go-gost/x/internal/util/limiter" pht_util "github.com/go-gost/x/internal/util/pht" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -95,7 +98,15 @@ func (l *phtListener) Accept() (conn net.Conn, err error) { conn = metrics.WrapConn(l.options.Service, conn) conn = stats.WrapConn(conn, l.options.Stats) conn = admission.WrapConn(l.options.Admission, conn) - conn = limiter.WrapConn(l.options.TrafficLimiter, conn) + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) return } diff --git a/listener/quic/listener.go b/listener/quic/listener.go index dff8bfa0..77ec522c 100644 --- a/listener/quic/listener.go +++ b/listener/quic/listener.go @@ -4,14 +4,17 @@ import ( "context" "net" "strings" + "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" + limiter_util "github.com/go-gost/x/internal/util/limiter" quic_util "github.com/go-gost/x/internal/util/quic" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -70,6 +73,18 @@ func (l *quicListener) Init(md md.Metadata) (err error) { conn = quic_util.CipherPacketConn(conn, l.md.cipherKey) } + conn = metrics.WrapPacketConn(l.options.Service, conn) + conn = stats.WrapPacketConn(conn, l.options.Stats) + conn = admission.WrapPacketConn(l.options.Admission, conn) + conn = limiter_wrapper.WrapPacketConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + "", + limiter.ScopeOption(limiter.ScopeService), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + ) + config := &quic.Config{ KeepAlivePeriod: l.md.keepAlivePeriod, HandshakeIdleTimeout: l.md.handshakeTimeout, @@ -102,10 +117,15 @@ func (l *quicListener) Accept() (conn net.Conn, err error) { var ok bool select { case conn = <-l.cqueue: - conn = metrics.WrapConn(l.options.Service, conn) - conn = stats.WrapConn(conn, l.options.Stats) - conn = admission.WrapConn(l.options.Admission, conn) - conn = limiter.WrapConn(l.options.TrafficLimiter, conn) + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) 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 4907698d..725f5f22 100644 --- a/listener/redirect/tcp/listener.go +++ b/listener/redirect/tcp/listener.go @@ -5,14 +5,16 @@ import ( "net" "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/proxyproto" + limiter_util "github.com/go-gost/x/internal/util/limiter" climiter "github.com/go-gost/x/limiter/conn/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -68,14 +70,33 @@ func (l *redirectListener) Init(md md.Metadata) (err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = stats.WrapListener(ln, l.options.Stats) ln = admission.WrapListener(l.options.Admission, ln) - ln = limiter.WrapListener(l.options.TrafficLimiter, ln) + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) ln = climiter.WrapListener(l.options.ConnLimiter, ln) l.ln = ln return } func (l *redirectListener) Accept() (conn net.Conn, err error) { - return l.ln.Accept() + conn, err = l.ln.Accept() + if err != nil { + return + } + + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) + + return } func (l *redirectListener) Addr() net.Addr { diff --git a/listener/redirect/udp/listener.go b/listener/redirect/udp/listener.go index ff0698a2..02ef3044 100644 --- a/listener/redirect/udp/listener.go +++ b/listener/redirect/udp/listener.go @@ -2,12 +2,15 @@ package udp import ( "net" + "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_util "github.com/go-gost/x/internal/util/limiter" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -57,7 +60,15 @@ func (l *redirectListener) Accept() (conn net.Conn, err error) { conn = metrics.WrapConn(l.options.Service, conn) conn = stats.WrapConn(conn, l.options.Stats) conn = admission.WrapConn(l.options.Admission, conn) - conn = limiter.WrapConn(l.options.TrafficLimiter, conn) + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) return } diff --git a/listener/rtcp/listener.go b/listener/rtcp/listener.go index 564c2676..9a4351f9 100644 --- a/listener/rtcp/listener.go +++ b/listener/rtcp/listener.go @@ -4,15 +4,18 @@ import ( "context" "net" "sync" + "time" "github.com/go-gost/core/chain" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" + limiter_util "github.com/go-gost/x/internal/util/limiter" climiter "github.com/go-gost/x/limiter/conn/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -81,7 +84,11 @@ func (l *rtcpListener) Accept() (conn net.Conn, err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = stats.WrapListener(ln, l.options.Stats) ln = admission.WrapListener(l.options.Admission, ln) - ln = limiter.WrapListener(l.options.TrafficLimiter, ln) + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) ln = climiter.WrapListener(l.options.ConnLimiter, ln) l.setListener(ln) } @@ -99,6 +106,17 @@ func (l *rtcpListener) Accept() (conn net.Conn, err error) { l.setListener(nil) return nil, listener.NewAcceptError(err) } + + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) + return } diff --git a/listener/rudp/listener.go b/listener/rudp/listener.go index b98bb643..97972ef4 100644 --- a/listener/rudp/listener.go +++ b/listener/rudp/listener.go @@ -4,14 +4,18 @@ import ( "context" "net" "sync" + "time" "github.com/go-gost/core/chain" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_util "github.com/go-gost/x/internal/util/limiter" + climiter "github.com/go-gost/x/limiter/conn/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -81,6 +85,14 @@ func (l *rudpListener) Accept() (conn net.Conn, err error) { if err != nil { return nil, listener.NewAcceptError(err) } + + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) + ln = climiter.WrapListener(l.options.ConnLimiter, ln) + l.setListener(ln) } @@ -102,7 +114,14 @@ func (l *rudpListener) Accept() (conn net.Conn, err error) { uc := metrics.WrapUDPConn(l.options.Service, pc) uc = stats.WrapUDPConn(uc, l.options.Stats) uc = admission.WrapUDPConn(l.options.Admission, uc) - conn = limiter.WrapUDPConn(l.options.TrafficLimiter, uc) + conn = limiter_wrapper.WrapUDPConn( + uc, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + "", + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + ) } return diff --git a/listener/serial/listener.go b/listener/serial/listener.go index 862dd3c9..26a21b9b 100644 --- a/listener/serial/listener.go +++ b/listener/serial/listener.go @@ -5,11 +5,13 @@ import ( "net" "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" + limiter_util "github.com/go-gost/x/internal/util/limiter" serial "github.com/go-gost/x/internal/util/serial" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -94,12 +96,19 @@ func (l *serialListener) listenLoop() { return err } - c := serial.NewConn(port, l.addr, cancel) - c = metrics.WrapConn(l.options.Service, c) - c = stats.WrapConn(c, l.options.Stats) - c = limiter.WrapConn(l.options.TrafficLimiter, c) + conn := serial.NewConn(port, l.addr, cancel) + conn = metrics.WrapConn(l.options.Service, conn) + conn = stats.WrapConn(conn, l.options.Stats) + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + "", + limiter.ScopeOption(limiter.ScopeService), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + ) - l.cqueue <- c + l.cqueue <- conn return nil }() diff --git a/listener/ssh/listener.go b/listener/ssh/listener.go index 2701b095..c3bc6f46 100644 --- a/listener/ssh/listener.go +++ b/listener/ssh/listener.go @@ -6,15 +6,17 @@ import ( "net" "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/proxyproto" + limiter_util "github.com/go-gost/x/internal/util/limiter" ssh_util "github.com/go-gost/x/internal/util/ssh" climiter "github.com/go-gost/x/limiter/conn/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -70,7 +72,11 @@ func (l *sshListener) Init(md md.Metadata) (err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = stats.WrapListener(ln, l.options.Stats) ln = admission.WrapListener(l.options.Admission, ln) - ln = limiter.WrapListener(l.options.TrafficLimiter, ln) + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) ln = climiter.WrapListener(l.options.ConnLimiter, ln) l.Listener = ln @@ -96,6 +102,15 @@ func (l *sshListener) Accept() (conn net.Conn, err error) { var ok bool select { case conn = <-l.cqueue: + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) case err, ok = <-l.errChan: if !ok { err = listener.ErrClosed diff --git a/listener/sshd/listener.go b/listener/sshd/listener.go index 6c49d5a8..2badfe25 100644 --- a/listener/sshd/listener.go +++ b/listener/sshd/listener.go @@ -7,16 +7,18 @@ import ( "strconv" "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/proxyproto" + limiter_util "github.com/go-gost/x/internal/util/limiter" ssh_util "github.com/go-gost/x/internal/util/ssh" sshd_util "github.com/go-gost/x/internal/util/sshd" climiter "github.com/go-gost/x/limiter/conn/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -78,7 +80,11 @@ func (l *sshdListener) Init(md md.Metadata) (err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = stats.WrapListener(ln, l.options.Stats) ln = admission.WrapListener(l.options.Admission, ln) - ln = limiter.WrapListener(l.options.TrafficLimiter, ln) + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) ln = climiter.WrapListener(l.options.ConnLimiter, ln) l.Listener = ln @@ -104,6 +110,15 @@ func (l *sshdListener) Accept() (conn net.Conn, err error) { var ok bool select { case conn = <-l.cqueue: + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) case err, ok = <-l.errChan: if !ok { err = listener.ErrClosed diff --git a/listener/tap/listener.go b/listener/tap/listener.go index c01cfbb3..108bb442 100644 --- a/listener/tap/listener.go +++ b/listener/tap/listener.go @@ -5,11 +5,13 @@ import ( "net" "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" mdata "github.com/go-gost/core/metadata" xnet "github.com/go-gost/x/internal/net" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_util "github.com/go-gost/x/internal/util/limiter" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" mdx "github.com/go-gost/x/metadata" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" @@ -92,7 +94,14 @@ func (l *tapListener) listenLoop() { } c = metrics.WrapConn(l.options.Service, c) c = stats.WrapConn(c, l.options.Stats) - c = limiter.WrapConn(l.options.TrafficLimiter, c) + c = limiter_wrapper.WrapConn( + c, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + c.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeService), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(c.LocalAddr().Network()), + ) 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 9e42ef44..e2a2ab01 100644 --- a/listener/tcp/listener.go +++ b/listener/tcp/listener.go @@ -5,14 +5,16 @@ import ( "net" "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/proxyproto" + limiter_util "github.com/go-gost/x/internal/util/limiter" climiter "github.com/go-gost/x/limiter/conn/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -66,7 +68,11 @@ func (l *tcpListener) Init(md md.Metadata) (err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = stats.WrapListener(ln, l.options.Stats) ln = admission.WrapListener(l.options.Admission, ln) - ln = limiter.WrapListener(l.options.TrafficLimiter, ln) + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) ln = climiter.WrapListener(l.options.ConnLimiter, ln) l.ln = ln @@ -74,7 +80,22 @@ func (l *tcpListener) Init(md md.Metadata) (err error) { } func (l *tcpListener) Accept() (conn net.Conn, err error) { - return l.ln.Accept() + conn, err = l.ln.Accept() + if err != nil { + return + } + + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) + + return } func (l *tcpListener) Addr() net.Addr { diff --git a/listener/tls/listener.go b/listener/tls/listener.go index 2e36fb76..ef5662f6 100644 --- a/listener/tls/listener.go +++ b/listener/tls/listener.go @@ -6,14 +6,16 @@ import ( "net" "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/proxyproto" + limiter_util "github.com/go-gost/x/internal/util/limiter" climiter "github.com/go-gost/x/limiter/conn/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -64,7 +66,11 @@ func (l *tlsListener) Init(md md.Metadata) (err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = stats.WrapListener(ln, l.options.Stats) ln = admission.WrapListener(l.options.Admission, ln) - ln = limiter.WrapListener(l.options.TrafficLimiter, ln) + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) ln = climiter.WrapListener(l.options.ConnLimiter, ln) l.ln = tls.NewListener(ln, l.options.TLSConfig) @@ -73,7 +79,22 @@ func (l *tlsListener) Init(md md.Metadata) (err error) { } func (l *tlsListener) Accept() (conn net.Conn, err error) { - return l.ln.Accept() + conn, err = l.ln.Accept() + if err != nil { + return + } + + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) + + return } func (l *tlsListener) Addr() net.Addr { diff --git a/listener/tun/listener.go b/listener/tun/listener.go index bf4f96d6..7784f7e7 100644 --- a/listener/tun/listener.go +++ b/listener/tun/listener.go @@ -5,12 +5,14 @@ import ( "net" "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" mdata "github.com/go-gost/core/metadata" "github.com/go-gost/core/router" xnet "github.com/go-gost/x/internal/net" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_util "github.com/go-gost/x/internal/util/limiter" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" mdx "github.com/go-gost/x/metadata" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" @@ -93,7 +95,14 @@ func (l *tunListener) listenLoop() { } c = metrics.WrapConn(l.options.Service, c) c = stats.WrapConn(c, l.options.Stats) - c = limiter.WrapConn(l.options.TrafficLimiter, c) + c = limiter_wrapper.WrapConn( + c, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + c.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeService), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(c.LocalAddr().Network()), + ) 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 e69a3b6b..d2fda937 100644 --- a/listener/udp/listener.go +++ b/listener/udp/listener.go @@ -2,14 +2,17 @@ package udp import ( "net" + "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/udp" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_util "github.com/go-gost/x/internal/util/limiter" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -59,7 +62,14 @@ func (l *udpListener) Init(md md.Metadata) (err error) { conn = metrics.WrapPacketConn(l.options.Service, conn) conn = stats.WrapPacketConn(conn, l.options.Stats) conn = admission.WrapPacketConn(l.options.Admission, conn) - conn = limiter.WrapPacketConn(l.options.TrafficLimiter, conn) + conn = limiter_wrapper.WrapPacketConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + "", + limiter.ScopeOption(limiter.ScopeService), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + ) l.ln = udp.NewListener(conn, &udp.ListenConfig{ Backlog: l.md.backlog, diff --git a/listener/unix/listener.go b/listener/unix/listener.go index 11f718ae..c5919cf7 100644 --- a/listener/unix/listener.go +++ b/listener/unix/listener.go @@ -2,13 +2,16 @@ package unix import ( "net" + "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" + limiter_util "github.com/go-gost/x/internal/util/limiter" climiter "github.com/go-gost/x/limiter/conn/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -49,7 +52,11 @@ func (l *unixListener) Init(md md.Metadata) (err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = stats.WrapListener(ln, l.options.Stats) ln = admission.WrapListener(l.options.Admission, ln) - ln = limiter.WrapListener(l.options.TrafficLimiter, ln) + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) ln = climiter.WrapListener(l.options.ConnLimiter, ln) l.ln = ln @@ -57,7 +64,22 @@ func (l *unixListener) Init(md md.Metadata) (err error) { } func (l *unixListener) Accept() (conn net.Conn, err error) { - return l.ln.Accept() + conn, err = l.ln.Accept() + if err != nil { + return + } + + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) + + return } func (l *unixListener) Addr() net.Addr { diff --git a/listener/ws/listener.go b/listener/ws/listener.go index 2a42a255..f5e299b5 100644 --- a/listener/ws/listener.go +++ b/listener/ws/listener.go @@ -8,15 +8,17 @@ import ( "net/http/httputil" "time" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/listener" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" admission "github.com/go-gost/x/admission/wrapper" xnet "github.com/go-gost/x/internal/net" "github.com/go-gost/x/internal/net/proxyproto" + limiter_util "github.com/go-gost/x/internal/util/limiter" ws_util "github.com/go-gost/x/internal/util/ws" climiter "github.com/go-gost/x/limiter/conn/wrapper" - limiter "github.com/go-gost/x/limiter/traffic/wrapper" + limiter_wrapper "github.com/go-gost/x/limiter/traffic/wrapper" metrics "github.com/go-gost/x/metrics/wrapper" stats "github.com/go-gost/x/observer/stats/wrapper" "github.com/go-gost/x/registry" @@ -105,7 +107,11 @@ func (l *wsListener) Init(md md.Metadata) (err error) { ln = metrics.WrapListener(l.options.Service, ln) ln = stats.WrapListener(ln, l.options.Stats) ln = admission.WrapListener(l.options.Admission, ln) - ln = limiter.WrapListener(l.options.TrafficLimiter, ln) + ln = limiter_wrapper.WrapListener( + l.options.Service, + ln, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + ) ln = climiter.WrapListener(l.options.ConnLimiter, ln) if l.tlsEnabled { @@ -129,6 +135,15 @@ func (l *wsListener) Accept() (conn net.Conn, err error) { var ok bool select { case conn = <-l.cqueue: + conn = limiter_wrapper.WrapConn( + conn, + limiter_util.NewCachedTrafficLimiter(l.options.TrafficLimiter, 30*time.Second, 60*time.Second), + conn.RemoteAddr().String(), + limiter.ScopeOption(limiter.ScopeConn), + limiter.ServiceOption(l.options.Service), + limiter.NetworkOption(conn.LocalAddr().Network()), + limiter.SrcOption(conn.RemoteAddr().String()), + ) case err, ok = <-l.errChan: if !ok { err = listener.ErrClosed diff --git a/registry/limiter.go b/registry/limiter.go index fedc718d..00fedc07 100644 --- a/registry/limiter.go +++ b/registry/limiter.go @@ -3,6 +3,7 @@ package registry import ( "context" + "github.com/go-gost/core/limiter" "github.com/go-gost/core/limiter/conn" "github.com/go-gost/core/limiter/rate" "github.com/go-gost/core/limiter/traffic" @@ -32,7 +33,7 @@ type trafficLimiterWrapper struct { r *trafficLimiterRegistry } -func (w *trafficLimiterWrapper) In(ctx context.Context, key string, opts ...traffic.Option) traffic.Limiter { +func (w *trafficLimiterWrapper) In(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter { v := w.r.get(w.name) if v == nil { return nil @@ -40,7 +41,7 @@ func (w *trafficLimiterWrapper) In(ctx context.Context, key string, opts ...traf return v.In(ctx, key, opts...) } -func (w *trafficLimiterWrapper) Out(ctx context.Context, key string, opts ...traffic.Option) traffic.Limiter { +func (w *trafficLimiterWrapper) Out(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter { v := w.r.get(w.name) if v == nil { return nil