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

add traffic limiter for proxy handler

Browse Source
This commit is contained in:
ginuerzh
2023-11-18 18:28:09 +08:00
parent 330631fd79
commit 88cc6ff4d5
38 changed files with 633 additions and 200 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
handler/socks/v4/handler.go
View File

@ -8,12 +8,13 @@ import (
"github.com/go-gost/core/chain"
"github.com/go-gost/core/handler"
"github.com/go-gost/core/limiter/traffic"
"github.com/go-gost/core/logger"
md "github.com/go-gost/core/metadata"
"github.com/go-gost/gosocks4"
ctxvalue "github.com/go-gost/x/internal/ctx"
netpkg "github.com/go-gost/x/internal/net"
auth_util "github.com/go-gost/x/internal/util/auth"
sx "github.com/go-gost/x/internal/util/selector"
"github.com/go-gost/x/limiter/traffic/wrapper"
"github.com/go-gost/x/registry"
)
@ -82,8 +83,6 @@ func (h *socks4Handler) Handle(ctx context.Context, conn net.Conn, opts ...handl
conn.SetReadDeadline(time.Now().Add(h.md.readTimeout))
}
ctx = auth_util.ContextWithClientAddr(ctx, auth_util.ClientAddr(conn.RemoteAddr().String()))
req, err := gosocks4.ReadRequest(conn)
if err != nil {
log.Error(err)
@ -100,7 +99,7 @@ func (h *socks4Handler) Handle(ctx context.Context, conn net.Conn, opts ...handl
log.Trace(resp)
return resp.Write(conn)
}
ctx = auth_util.ContextWithID(ctx, auth_util.ID(id))
ctx = ctxvalue.ContextWithClientID(ctx, ctxvalue.ClientID(id))
}
switch req.Cmd {
@ -132,7 +131,7 @@ func (h *socks4Handler) handleConnect(ctx context.Context, conn net.Conn, req *g
switch h.md.hash {
case "host":
ctx = sx.ContextWithHash(ctx, &sx.Hash{Source: addr})
ctx = ctxvalue.ContextWithHash(ctx, &ctxvalue.Hash{Source: addr})
}
cc, err := h.router.Dial(ctx, "tcp", addr)
@ -152,9 +151,16 @@ func (h *socks4Handler) handleConnect(ctx context.Context, conn net.Conn, req *g
return err
}
rw := wrapper.WrapReadWriter(h.options.Limiter, conn, conn.RemoteAddr().String(),
traffic.NetworkOption("tcp"),
traffic.AddrOption(addr),
traffic.ClientOption(string(ctxvalue.ClientIDFromContext(ctx))),
traffic.SrcOption(conn.RemoteAddr().String()),
)
t := time.Now()
log.Infof("%s <-> %s", conn.RemoteAddr(), addr)
netpkg.Transport(conn, cc)
netpkg.Transport(rw, cc)
log.WithFields(map[string]any{
"duration": time.Since(t),
}).Infof("%s >-< %s", conn.RemoteAddr(), addr)