Files
x/handler/socks/v4/handler.go
T
2025-08-29 23:36:31 +08:00

385 lines
9.7 KiB
Go

package v4
import (
"bufio"
"bytes"
"context"
"crypto/tls"
"errors"
"net"
"time"
"github.com/go-gost/core/auth"
"github.com/go-gost/core/bypass"
"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"
"github.com/go-gost/core/observer"
"github.com/go-gost/core/observer/stats"
"github.com/go-gost/core/recorder"
"github.com/go-gost/gosocks4"
xctx "github.com/go-gost/x/ctx"
ictx "github.com/go-gost/x/internal/ctx"
xnet "github.com/go-gost/x/internal/net"
"github.com/go-gost/x/internal/util/sniffing"
stats_util "github.com/go-gost/x/internal/util/stats"
tls_util "github.com/go-gost/x/internal/util/tls"
rate_limiter "github.com/go-gost/x/limiter/rate"
cache_limiter "github.com/go-gost/x/limiter/traffic/cache"
traffic_wrapper "github.com/go-gost/x/limiter/traffic/wrapper"
xstats "github.com/go-gost/x/observer/stats"
stats_wrapper "github.com/go-gost/x/observer/stats/wrapper"
xrecorder "github.com/go-gost/x/recorder"
"github.com/go-gost/x/registry"
)
var (
ErrUnknownCmd = errors.New("socks4: unknown command")
ErrUnimplemented = errors.New("socks4: unimplemented")
)
func init() {
registry.HandlerRegistry().Register("socks4", NewHandler)
registry.HandlerRegistry().Register("socks4a", NewHandler)
}
type socks4Handler struct {
md metadata
options handler.Options
stats *stats_util.HandlerStats
limiter traffic.TrafficLimiter
cancel context.CancelFunc
recorder recorder.RecorderObject
certPool tls_util.CertPool
}
func NewHandler(opts ...handler.Option) handler.Handler {
options := handler.Options{}
for _, opt := range opts {
opt(&options)
}
return &socks4Handler{
options: options,
}
}
func (h *socks4Handler) Init(md md.Metadata) (err error) {
if err := h.parseMetadata(md); err != nil {
return err
}
ctx, cancel := context.WithCancel(context.Background())
h.cancel = cancel
if h.options.Observer != nil {
h.stats = stats_util.NewHandlerStats(h.options.Service, h.md.observerResetTraffic)
go h.observeStats(ctx)
}
if h.options.Limiter != nil {
h.limiter = cache_limiter.NewCachedTrafficLimiter(h.options.Limiter,
cache_limiter.RefreshIntervalOption(h.md.limiterRefreshInterval),
cache_limiter.CleanupIntervalOption(h.md.limiterCleanupInterval),
cache_limiter.ScopeOption(limiter.ScopeClient),
)
}
for _, ro := range h.options.Recorders {
if ro.Record == xrecorder.RecorderServiceHandler {
h.recorder = ro
break
}
}
if h.md.certificate != nil && h.md.privateKey != nil {
h.certPool = tls_util.NewMemoryCertPool()
}
return nil
}
func (h *socks4Handler) Handle(ctx context.Context, conn net.Conn, opts ...handler.HandleOption) (err error) {
defer conn.Close()
start := time.Now()
ro := &xrecorder.HandlerRecorderObject{
Network: "tcp",
Service: h.options.Service,
RemoteAddr: conn.RemoteAddr().String(),
LocalAddr: conn.LocalAddr().String(),
SID: xctx.SidFromContext(ctx).String(),
Time: start,
}
if srcAddr := xctx.SrcAddrFromContext(ctx); srcAddr != nil {
ro.ClientAddr = srcAddr.String()
}
log := h.options.Logger.WithFields(map[string]any{
"network": ro.Network,
"remote": conn.RemoteAddr().String(),
"local": conn.LocalAddr().String(),
"client": ro.ClientAddr,
"sid": ro.SID,
})
log.Infof("%s <> %s", conn.RemoteAddr(), conn.LocalAddr())
pStats := xstats.Stats{}
conn = stats_wrapper.WrapConn(conn, &pStats)
defer func() {
if err != nil {
ro.Err = err.Error()
}
ro.InputBytes = pStats.Get(stats.KindInputBytes)
ro.OutputBytes = pStats.Get(stats.KindOutputBytes)
ro.Duration = time.Since(start)
if err := ro.Record(ctx, h.recorder.Recorder); err != nil {
log.Errorf("record: %v", err)
}
log.WithFields(map[string]any{
"duration": time.Since(start),
"inputBytes": ro.InputBytes,
"outputBytes": ro.OutputBytes,
}).Infof("%s >< %s", conn.RemoteAddr(), conn.LocalAddr())
}()
if !h.checkRateLimit(conn.RemoteAddr()) {
return rate_limiter.ErrRateLimit
}
if h.md.readTimeout > 0 {
conn.SetReadDeadline(time.Now().Add(h.md.readTimeout))
}
req, err := gosocks4.ReadRequest(conn)
if err != nil {
log.Error(err)
return err
}
if userid := string(req.Userid); userid != "" {
log = log.WithFields(map[string]any{"user": userid})
ro.ClientID = userid
}
ro.Host = req.Addr.String()
log.Trace(req)
conn.SetReadDeadline(time.Time{})
if h.options.Auther != nil {
clientID, ok := h.options.Auther.Authenticate(ctx, string(req.Userid), "", auth.WithService(h.options.Service))
if !ok {
resp := gosocks4.NewReply(gosocks4.RejectedUserid, nil)
log.Trace(resp)
return resp.Write(conn)
}
ctx = xctx.ContextWithClientID(ctx, xctx.ClientID(clientID))
ro.ClientID = clientID
log = log.WithFields(map[string]any{"clientID": clientID})
}
switch req.Cmd {
case gosocks4.CmdConnect:
return h.handleConnect(ctx, conn, req, ro, log)
case gosocks4.CmdBind:
return h.handleBind(ctx, conn, req)
default:
err = ErrUnknownCmd
log.Error(err)
return err
}
}
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, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error {
addr := req.Addr.String()
log = log.WithFields(map[string]any{
"dst": addr,
"host": addr,
})
log.Debugf("%s >> %s", conn.RemoteAddr(), addr)
{
clientID := xctx.ClientIDFromContext(ctx)
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))
pstats.Add(stats.KindTotalConns, 1)
pstats.Add(stats.KindCurrentConns, 1)
defer pstats.Add(stats.KindCurrentConns, -1)
rw = stats_wrapper.WrapReadWriter(rw, pstats)
}
conn = xnet.NewReadWriteConn(rw, rw, conn)
}
if h.options.Bypass != nil && h.options.Bypass.Contains(ctx, "tcp", addr, bypass.WithService(h.options.Service)) {
resp := gosocks4.NewReply(gosocks4.Rejected, nil)
log.Trace(resp)
log.Debug("bypass: ", addr)
return resp.Write(conn)
}
switch h.md.hash {
case "host":
ctx = xctx.ContextWithHash(ctx, &xctx.Hash{Source: addr})
}
var buf bytes.Buffer
cc, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), "tcp", addr)
ro.Route = buf.String()
if err != nil {
resp := gosocks4.NewReply(gosocks4.Failed, nil)
log.Trace(resp)
resp.Write(conn)
return err
}
defer cc.Close()
log = log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()})
ro.SrcAddr = cc.LocalAddr().String()
ro.DstAddr = cc.RemoteAddr().String()
resp := gosocks4.NewReply(gosocks4.Granted, nil)
log.Trace(resp)
if err := resp.Write(conn); err != nil {
log.Error(err)
return err
}
if h.md.sniffing {
if h.md.sniffingTimeout > 0 {
conn.SetReadDeadline(time.Now().Add(h.md.sniffingTimeout))
}
br := bufio.NewReader(conn)
proto, _ := sniffing.Sniff(ctx, br)
ro.Proto = proto
if h.md.sniffingTimeout > 0 {
conn.SetReadDeadline(time.Time{})
}
dial := func(ctx context.Context, network, address string) (net.Conn, error) {
return cc, nil
}
dialTLS := func(ctx context.Context, network, address string, cfg *tls.Config) (net.Conn, error) {
return cc, nil
}
sniffer := &sniffing.Sniffer{
Websocket: h.md.sniffingWebsocket,
WebsocketSampleRate: h.md.sniffingWebsocketSampleRate,
Recorder: h.recorder.Recorder,
RecorderOptions: h.recorder.Options,
Certificate: h.md.certificate,
PrivateKey: h.md.privateKey,
NegotiatedProtocol: h.md.alpn,
CertPool: h.certPool,
MitmBypass: h.md.mitmBypass,
ReadTimeout: h.md.readTimeout,
}
conn = xnet.NewReadWriteConn(br, conn, conn)
switch proto {
case sniffing.ProtoHTTP:
return sniffer.HandleHTTP(ctx, "tcp", conn,
sniffing.WithService(h.options.Service),
sniffing.WithDial(dial),
sniffing.WithDialTLS(dialTLS),
sniffing.WithRecorderObject(ro),
sniffing.WithLog(log),
)
case sniffing.ProtoTLS:
return sniffer.HandleTLS(ctx, "tcp", conn,
sniffing.WithService(h.options.Service),
sniffing.WithDial(dial),
sniffing.WithDialTLS(dialTLS),
sniffing.WithRecorderObject(ro),
sniffing.WithLog(log),
)
}
}
t := time.Now()
log.Infof("%s <-> %s", conn.RemoteAddr(), ro.Host)
// xnet.Transport(conn, cc)
xnet.Pipe(ctx, conn, cc)
log.WithFields(map[string]any{
"duration": time.Since(t),
}).Infof("%s >-< %s", conn.RemoteAddr(), ro.Host)
return nil
}
func (h *socks4Handler) handleBind(ctx context.Context, conn net.Conn, req *gosocks4.Request) error {
// TODO: bind
return ErrUnimplemented
}
func (h *socks4Handler) checkRateLimit(addr net.Addr) bool {
if h.options.RateLimiter == nil {
return true
}
host, _, _ := net.SplitHostPort(addr.String())
if limiter := h.options.RateLimiter.Limiter(host); limiter != nil {
return limiter.Allow(1)
}
return true
}
func (h *socks4Handler) observeStats(ctx context.Context) {
if h.options.Observer == nil {
return
}
var events []observer.Event
ticker := time.NewTicker(h.md.observerPeriod)
defer ticker.Stop()
for {
select {
case <-ticker.C:
if len(events) > 0 {
if err := h.options.Observer.Observe(ctx, events); err == nil {
events = nil
}
break
}
evs := h.stats.Events()
if err := h.options.Observer.Observe(ctx, evs); err != nil {
events = evs
}
case <-ctx.Done():
return
}
}
}