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

add scope param for traffic limiter

Browse Source
This commit is contained in:
ginuerzh
2024-08-31 09:30:04 +08:00
parent 9405d1fb48
commit 31c482c649
57 changed files with 1239 additions and 485 deletions
Download Patch File Download Diff File Expand all files Collapse all files
handler/http/handler.go
+19 -7
View File
@@ -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)
handler/http2/handler.go
+20 -7
View File
@@ -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)
handler/relay/bind.go
+1
View File
@@ -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",
})),
handler/relay/connect.go
+12 -7
View File
@@ -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))
handler/relay/entrypoint.go
-4
View File
@@ -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
handler/relay/forward.go
+11 -6
View File
@@ -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))
handler/relay/handler.go
+8 -1
View File
@@ -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
}
handler/socks/v4/handler.go
+19 -7
View File
@@ -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))
handler/socks/v5/connect.go
+12 -7
View File
@@ -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))
handler/socks/v5/handler.go
+8 -1
View File
@@ -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
}
handler/tunnel/bind.go
+9 -1
View File
@@ -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,
handler/tunnel/connect.go
+1 -11
View File
@@ -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())
handler/tunnel/entrypoint.go
-2
View File
@@ -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
handler/tunnel/handler.go
+8 -1
View File
@@ -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
}
handler/tunnel/tunnel.go
+45 -20
View File
@@ -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
}