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

12
handler/tunnel/connect.go
View File

@ -6,9 +6,12 @@ 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/internal/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 {
@ -95,9 +98,16 @@ func (h *tunnelHandler) handleConnect(ctx context.Context, req *relay.Request, c
req.WriteTo(cc)
}
rw := wrapper.WrapReadWriter(h.options.Limiter, conn, tunnelID.String(),
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(conn, cc)
xnet.Transport(rw, cc)
log.WithFields(map[string]any{
"duration": time.Since(t),
}).Debugf("%s >-< %s", conn.RemoteAddr(), cc.RemoteAddr())

8
handler/tunnel/handler.go
View File

@ -15,8 +15,8 @@ import (
"github.com/go-gost/core/recorder"
"github.com/go-gost/core/service"
"github.com/go-gost/relay"
ctxvalue "github.com/go-gost/x/internal/ctx"
xnet "github.com/go-gost/x/internal/net"
auth_util "github.com/go-gost/x/internal/util/auth"
xrecorder "github.com/go-gost/x/recorder"
"github.com/go-gost/x/registry"
xservice "github.com/go-gost/x/service"
@ -169,8 +169,6 @@ func (h *tunnelHandler) Handle(ctx context.Context, conn net.Conn, opts ...handl
}).Infof("%s >< %s", conn.RemoteAddr(), conn.LocalAddr())
}()
ctx = auth_util.ContextWithClientAddr(ctx, auth_util.ClientAddr(conn.RemoteAddr().String()))
if !h.checkRateLimit(conn.RemoteAddr()) {
return ErrRateLimit
}
@ -238,13 +236,13 @@ func (h *tunnelHandler) Handle(ctx context.Context, conn net.Conn, opts ...handl
}
if h.options.Auther != nil {
id, ok := h.options.Auther.Authenticate(ctx, user, pass)
clientID, ok := h.options.Auther.Authenticate(ctx, user, pass)
if !ok {
resp.Status = relay.StatusUnauthorized
resp.WriteTo(conn)
return ErrUnauthorized
}
ctx = auth_util.ContextWithID(ctx, auth_util.ID(id))
ctx = ctxvalue.ContextWithClientID(ctx, ctxvalue.ClientID(clientID))
}
switch req.Cmd & relay.CmdMask {