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

add observer

Browse Source
This commit is contained in:
ginuerzh
2024-01-03 20:55:06 +08:00
parent e1ae379048
commit c959fc2f73
95 changed files with 2371 additions and 890 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -16,6 +16,9 @@ import (
netpkg "github.com/go-gost/x/internal/net"
"github.com/go-gost/x/limiter/traffic/wrapper"
"github.com/go-gost/x/registry"
stats_util "github.com/go-gost/x/internal/util/stats"
"github.com/go-gost/x/stats"
stats_wrapper "github.com/go-gost/x/stats/wrapper"
)
var (
@ -32,6 +35,8 @@ type socks4Handler struct {
router *chain.Router
md metadata
options handler.Options
stats *stats_util.HandlerStats
cancel context.CancelFunc
}
func NewHandler(opts ...handler.Option) handler.Handler {
@ -42,6 +47,7 @@ func NewHandler(opts ...handler.Option) handler.Handler {
return &socks4Handler{
options: options,
stats: stats_util.NewHandlerStats(options.Service),
}
}
@ -55,6 +61,13 @@ func (h *socks4Handler) Init(md md.Metadata) (err error) {
h.router = chain.NewRouter(chain.LoggerRouterOption(h.options.Logger))
}
ctx, cancel := context.WithCancel(context.Background())
h.cancel = cancel
if h.options.Observer != nil {
go h.observeStats(ctx)
}
return nil
}
@ -114,6 +127,13 @@ func (h *socks4Handler) Handle(ctx context.Context, conn net.Conn, opts ...handl
}
}
func (h *socks4Handler) Close() error {
if h.cancel != nil {
h.cancel()
}
return nil
}
func (h *socks4Handler) handleConnect(ctx context.Context, conn net.Conn, req *gosocks4.Request, log logger.Logger) error {
addr := req.Addr.String()
@ -151,12 +171,20 @@ func (h *socks4Handler) handleConnect(ctx context.Context, conn net.Conn, req *g
return err
}
rw := wrapper.WrapReadWriter(h.options.Limiter, conn, conn.RemoteAddr().String(),
clientID := ctxvalue.ClientIDFromContext(ctx)
rw := wrapper.WrapReadWriter(h.options.Limiter, conn,
traffic.NetworkOption("tcp"),
traffic.AddrOption(addr),
traffic.ClientOption(string(ctxvalue.ClientIDFromContext(ctx))),
traffic.ClientOption(string(clientID)),
traffic.SrcOption(conn.RemoteAddr().String()),
)
if h.options.Observer != nil {
pstats := h.stats.Stats(string(clientID))
pstats.Add(stats.KindTotalConns, 1)
pstats.Add(stats.KindCurrentConns, 1)
defer pstats.Add(stats.KindCurrentConns, -1)
rw = stats_wrapper.WrapReadWriter(rw, pstats)
}
t := time.Now()
log.Infof("%s <-> %s", conn.RemoteAddr(), addr)
@ -184,3 +212,21 @@ func (h *socks4Handler) checkRateLimit(addr net.Addr) bool {
return true
}
func (h *socks4Handler) observeStats(ctx context.Context) {
if h.options.Observer == nil {
return
}
ticker := time.NewTicker(5 * time.Second)
defer ticker.Stop()
for {
select {
case <-ticker.C:
h.options.Observer.Observe(ctx, h.stats.Events())
case <-ctx.Done():
return
}
}
}

14
handler/socks/v5/connect.go
View File

@ -12,6 +12,8 @@ import (
ctxvalue "github.com/go-gost/x/internal/ctx"
netpkg "github.com/go-gost/x/internal/net"
"github.com/go-gost/x/limiter/traffic/wrapper"
"github.com/go-gost/x/stats"
stats_wrapper "github.com/go-gost/x/stats/wrapper"
)
func (h *socks5Handler) handleConnect(ctx context.Context, conn net.Conn, network, address string, log logger.Logger) error {
@ -50,12 +52,20 @@ func (h *socks5Handler) handleConnect(ctx context.Context, conn net.Conn, networ
return err
}
rw := wrapper.WrapReadWriter(h.options.Limiter, conn, conn.RemoteAddr().String(),
clientID := ctxvalue.ClientIDFromContext(ctx)
rw := wrapper.WrapReadWriter(h.options.Limiter, conn,
traffic.NetworkOption(network),
traffic.AddrOption(address),
traffic.ClientOption(string(ctxvalue.ClientIDFromContext(ctx))),
traffic.ClientOption(string(clientID)),
traffic.SrcOption(conn.RemoteAddr().String()),
)
if h.options.Observer != nil {
pstats := h.stats.Stats(string(clientID))
pstats.Add(stats.KindTotalConns, 1)
pstats.Add(stats.KindCurrentConns, 1)
defer pstats.Add(stats.KindCurrentConns, -1)
rw = stats_wrapper.WrapReadWriter(rw, pstats)
}
t := time.Now()
log.Infof("%s <-> %s", conn.RemoteAddr(), address)

36
handler/socks/v5/handler.go
View File

@ -12,6 +12,7 @@ import (
"github.com/go-gost/gosocks5"
ctxvalue "github.com/go-gost/x/internal/ctx"
"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 +30,8 @@ type socks5Handler struct {
router *chain.Router
md metadata
options handler.Options
stats *stats_util.HandlerStats
cancel context.CancelFunc
}
func NewHandler(opts ...handler.Option) handler.Handler {
@ -39,6 +42,7 @@ func NewHandler(opts ...handler.Option) handler.Handler {
return &socks5Handler{
options: options,
stats: stats_util.NewHandlerStats(options.Service),
}
}
@ -59,6 +63,13 @@ func (h *socks5Handler) Init(md md.Metadata) (err error) {
noTLS: h.md.noTLS,
}
ctx, cancel := context.WithCancel(context.Background())
h.cancel = cancel
if h.options.Observer != nil {
go h.observeStats(ctx)
}
return
}
@ -125,6 +136,13 @@ func (h *socks5Handler) Handle(ctx context.Context, conn net.Conn, opts ...handl
}
}
func (h *socks5Handler) Close() error {
if h.cancel != nil {
h.cancel()
}
return nil
}
func (h *socks5Handler) checkRateLimit(addr net.Addr) bool {
if h.options.RateLimiter == nil {
return true
@ -136,3 +154,21 @@ func (h *socks5Handler) checkRateLimit(addr net.Addr) bool {
return true
}
func (h *socks5Handler) observeStats(ctx context.Context) {
if h.options.Observer == nil {
return
}
ticker := time.NewTicker(5 * time.Second)
defer ticker.Stop()
for {
select {
case <-ticker.C:
h.options.Observer.Observe(ctx, h.stats.Events())
case <-ctx.Done():
return
}
}
}

15
handler/socks/v5/udp.go
View File

@ -10,8 +10,11 @@ import (
"github.com/go-gost/core/logger"
"github.com/go-gost/gosocks5"
ctxvalue "github.com/go-gost/x/internal/ctx"
"github.com/go-gost/x/internal/net/udp"
"github.com/go-gost/x/internal/util/socks"
"github.com/go-gost/x/stats"
stats_wrapper "github.com/go-gost/x/stats/wrapper"
)
func (h *socks5Handler) handleUDP(ctx context.Context, conn net.Conn, log logger.Logger) error {
@ -67,7 +70,17 @@ func (h *socks5Handler) handleUDP(ctx context.Context, conn net.Conn, log logger
return err
}
r := udp.NewRelay(socks.UDPConn(cc, h.md.udpBufferSize), pc).
var lc net.PacketConn = cc
clientID := ctxvalue.ClientIDFromContext(ctx)
if h.options.Observer != nil {
pstats := h.stats.Stats(string(clientID))
pstats.Add(stats.KindTotalConns, 1)
pstats.Add(stats.KindCurrentConns, 1)
defer pstats.Add(stats.KindCurrentConns, -1)
lc = stats_wrapper.WrapPacketConn(lc, pstats)
}
r := udp.NewRelay(socks.UDPConn(lc, h.md.udpBufferSize), pc).
WithBypass(h.options.Bypass).
WithLogger(log)
r.SetBufferSize(h.md.udpBufferSize)

12
handler/socks/v5/udp_tun.go
View File

@ -7,8 +7,11 @@ import (
"github.com/go-gost/core/logger"
"github.com/go-gost/gosocks5"
ctxvalue "github.com/go-gost/x/internal/ctx"
"github.com/go-gost/x/internal/net/udp"
"github.com/go-gost/x/internal/util/socks"
"github.com/go-gost/x/stats"
stats_wrapper "github.com/go-gost/x/stats/wrapper"
)
func (h *socks5Handler) handleUDPTun(ctx context.Context, conn net.Conn, network, address string, log logger.Logger) error {
@ -56,6 +59,15 @@ func (h *socks5Handler) handleUDPTun(ctx context.Context, conn net.Conn, network
}
log.Debugf("bind on %s OK", pc.LocalAddr())
clientID := ctxvalue.ClientIDFromContext(ctx)
if h.options.Observer != nil {
pstats := h.stats.Stats(string(clientID))
pstats.Add(stats.KindTotalConns, 1)
pstats.Add(stats.KindCurrentConns, 1)
defer pstats.Add(stats.KindCurrentConns, -1)
conn = stats_wrapper.WrapConn(conn, pstats)
}
r := udp.NewRelay(socks.UDPTunServerConn(conn), pc).
WithBypass(h.options.Bypass).
WithLogger(log)