From 5b21bc7cda1279c98bc18133515384ffdc59eddc Mon Sep 17 00:00:00 2001 From: Atsika Date: Thu, 19 Mar 2026 11:04:18 +0100 Subject: [PATCH] Add WASM agent support for JavaScript runtimes Compile the agent to WebAssembly for deployment in JS runtimes (Bun, Node.js, or any runtime implementing the TCPDial/UDPListen interface). - Split network and UDP code with js/native build tags - Extract UDP relay into a platform-agnostic shared layer - Expose runtime-agnostic TCPDial and UDPListen JS hooks - Add azure-sdk-for-go fork as submodule for WASM-compatible mmf fallback - Bump aznet dependency - Document WASM build and JS interface in README --- .gitignore | 3 +- .gitmodules | 3 + Makefile | 9 +- README.md | 36 +++- cmd/agent/main.go | 10 + cmd/proxy/main.go | 20 +- go.mod | 10 +- go.sum | 14 +- pkg/azure-sdk | 1 + pkg/protocol/base.go | 86 +++++--- pkg/protocol/connection.go | 2 +- pkg/protocol/netconn.go | 2 +- pkg/protocol/protocol.go | 32 +++ pkg/proxy/server/server.go | 2 +- pkg/proxy/socks/connect.go | 5 +- pkg/proxy/socks/dial_js.go | 115 +++++++++++ pkg/proxy/socks/dial_native.go | 12 ++ pkg/proxy/socks/udp_associate.go | 327 ++++++++----------------------- pkg/proxy/socks/udp_js.go | 137 +++++++++++++ pkg/proxy/socks/udp_native.go | 41 ++++ pkg/proxy/socks/udp_relay.go | 28 +++ 21 files changed, 593 insertions(+), 302 deletions(-) create mode 100644 .gitmodules create mode 160000 pkg/azure-sdk create mode 100644 pkg/proxy/socks/dial_js.go create mode 100644 pkg/proxy/socks/dial_native.go create mode 100644 pkg/proxy/socks/udp_js.go create mode 100644 pkg/proxy/socks/udp_native.go create mode 100644 pkg/proxy/socks/udp_relay.go diff --git a/.gitignore b/.gitignore index 0da5764..1bc0f7f 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ agent agent.exe +agent.wasm !agent/ proxy proxy.exe @@ -8,8 +9,8 @@ config.json *-config.json .DS_Store .cursor/ +.claude/ CLAUDE.md -AGENTS.md *.txt __blobstorage__/ __azurite*.json \ No newline at end of file diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..826270c --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "pkg/azure-sdk"] + path = pkg/azure-sdk + url = https://github.com/Atsika/azure-sdk-for-go.git diff --git a/Makefile b/Makefile index 2438c17..dd5e215 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: all clean proxy agent +.PHONY: all clean proxy agent wasm ifndef TOKEN CONN_STRING ?= @@ -6,7 +6,7 @@ else CONN_STRING = -X 'main.ConnString=$(TOKEN)' endif -all: clean proxy agent +all: clean proxy agent wasm proxy: go build -ldflags="-s -w" -trimpath -o proxy cmd/proxy/main.go @@ -14,5 +14,8 @@ proxy: agent: go build -ldflags="-s -w $(CONN_STRING)" -trimpath -o agent cmd/agent/main.go +wasm: + GOOS=js GOARCH=wasm go build -ldflags="-s -w $(CONN_STRING)" -trimpath -o agent.wasm cmd/agent/main.go + clean: - rm -f proxy agent \ No newline at end of file + rm -f proxy agent agent.wasm \ No newline at end of file diff --git a/README.md b/README.md index 039daf3..16dce23 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,7 @@ The system consists of two components: - Interactive CLI with auto-completion - Multiple agent management - Local or remote proxy server +- WebAssembly (WASM) agent for deployment in JavaScript runtimes (Bun, Node.js, etc.) ## Prerequisites @@ -120,11 +121,11 @@ The default storage account name is `devstoreaccount1` and the account key is `E ## Installation ```bash -# Clone the repository -git clone https://github.com/quarkslab/proxyblob +# Clone the repository (include submodules) +git clone --recurse-submodules https://github.com/quarkslab/proxyblob cd proxyblob -# Build both components +# Build all components make ``` @@ -132,6 +133,29 @@ This will produce the following binaries: - `proxy` - the proxy server running on your machine - `agent` - the agent running on the target +- `agent.wasm` - the agent compiled for WebAssembly + +The WASM agent requires the JS host runtime to expose two global functions: + +```js +// Open a TCP connection to host:port, call callbacks on events +TCPDial(host, port, onConnect, onData, onClose, onError) + +// Bind a UDP socket on a random port, call callbacks on events +UDPListen(onBind, onData, onError) +``` + +The runner is responsible for implementing these using the runtime's native APIs (e.g. Bun: `Bun.connect` / `Bun.udpSocket`, Node.js: `net` / `dgram`). + +You can also build targets individually: + +```bash +make proxy # proxy server only +make agent # native agent only +make wasm # WASM agent only +make agent TOKEN= # native agent with embedded connection string +make wasm TOKEN= # WASM agent with embedded connection string +``` ## Configuration @@ -401,6 +425,12 @@ If you encounter issues: ## CHANGELOG +**ProxyBlob v2.1 (WASM agent) - 19/03/2026:** + +- WebAssembly agent build target (`agent.wasm`) for deployment in JavaScript runtimes (Bun, Node.js, etc.) +- Build tag separation for JS/native network and UDP code +- Refactored UDP relay into a shared, platform-agnostic layer + **ProxyBlob v2 (aznet boosted) - 16/02/2026:** - Complete architecture rewrite using aznet networking layer diff --git a/cmd/agent/main.go b/cmd/agent/main.go index 14d8b56..3cfda46 100644 --- a/cmd/agent/main.go +++ b/cmd/agent/main.go @@ -138,6 +138,12 @@ func init() { // Use a more human-friendly output for console log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) + + go func() { + for { + time.Sleep(time.Hour) + } + }() } // main is the entry point for the agent process @@ -147,6 +153,10 @@ func main() { flag.StringVar(&ConnString, "c", ConnString, "Connection string") flag.Parse() + if ConnString == "" { + ConnString = os.Getenv("CONNECTION_STRING") + } + if ConnString == "" { os.Exit(ErrNoConnectionString) } diff --git a/cmd/proxy/main.go b/cmd/proxy/main.go index 0e76aaa..69dd682 100644 --- a/cmd/proxy/main.go +++ b/cmd/proxy/main.go @@ -34,8 +34,8 @@ const banner = ` |_| |_| \___/_/\_\\__, |____/|_|\___/|_.__/ |___/ - SOCKS Proxy over More Azure Storage (v2.0 - aznet) - -------------------------------------------------- + SOCKS Proxy over More Azure Storage (v2.1 - wasm) + ------------------------------------------------- ` @@ -619,8 +619,8 @@ func formatRelativeTime(t time.Time) string { func AddCommands(app *grumble.App) { // Parent command for listener management listenerCmd := &grumble.Command{ - Name: "listener", - Help: "manage listeners", + Name: "listener", + Help: "manage listeners", } // Subcommand: listener list @@ -745,8 +745,8 @@ func AddCommands(app *grumble.App) { // Command to create a new agent connection string app.AddCommand(&grumble.Command{ - Name: "new", - Help: "generate a new connection string for an agent", + Name: "new", + Help: "generate a new connection string for an agent", Flags: func(f *grumble.Flags) { f.Duration("d", "duration", 7*24*time.Hour, "duration for the SAS token (default 7 days)") f.String("l", "listener", "", "listener ID to use (defaults to selected listener)") @@ -789,8 +789,8 @@ func AddCommands(app *grumble.App) { // Parent command for agent management agentCmd := &grumble.Command{ - Name: "agent", - Help: "manage agents", + Name: "agent", + Help: "manage agents", } // Subcommand: agent list @@ -839,8 +839,8 @@ func AddCommands(app *grumble.App) { // Subcommand: agent start agentCmd.AddCommand(&grumble.Command{ - Name: "start", - Help: "start SOCKS proxy for the selected agent", + Name: "start", + Help: "start SOCKS proxy for the selected agent", Flags: func(f *grumble.Flags) { f.String("l", "listen", "127.0.0.1:1080", "listen address for SOCKS server") }, diff --git a/go.mod b/go.mod index a6408f7..55b3576 100644 --- a/go.mod +++ b/go.mod @@ -36,7 +36,15 @@ require ( ) require ( - github.com/atsika/aznet v0.0.0-20260216144450-1395eae946cd + github.com/atsika/aznet v0.0.0-20260319000323-1bb45b486a9e github.com/jedib0t/go-pretty v4.3.0+incompatible github.com/rs/zerolog v1.34.0 ) + +replace ( + github.com/Azure/azure-sdk-for-go/sdk/azcore => ./pkg/azure-sdk/sdk/azcore + github.com/Azure/azure-sdk-for-go/sdk/data/aztables => ./pkg/azure-sdk/sdk/data/aztables + github.com/Azure/azure-sdk-for-go/sdk/internal => ./pkg/azure-sdk/sdk/internal + github.com/Azure/azure-sdk-for-go/sdk/storage/azblob => ./pkg/azure-sdk/sdk/storage/azblob + github.com/Azure/azure-sdk-for-go/sdk/storage/azqueue => ./pkg/azure-sdk/sdk/storage/azqueue +) diff --git a/go.sum b/go.sum index f6a20d8..d4853ee 100644 --- a/go.sum +++ b/go.sum @@ -1,26 +1,16 @@ github.com/AlecAivazis/survey/v2 v2.0.5/go.mod h1:WYBhg6f0y/fNYUuesWQc0PKbJcEliGcYHB9sNT3Bg74= -github.com/Azure/azure-sdk-for-go/sdk/azcore v1.21.0 h1:fou+2+WFTib47nS+nz/ozhEBnvU96bKHy6LjRsY4E28= -github.com/Azure/azure-sdk-for-go/sdk/azcore v1.21.0/go.mod h1:t76Ruy8AHvUAC8GfMWJMa0ElSbuIcO03NLpynfbgsPA= github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.13.1 h1:Hk5QBxZQC1jb2Fwj6mpzme37xbCDdNTxU7O9eb5+LB4= github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.13.1/go.mod h1:IYus9qsFobWIc2YVwe/WPjcnyCkPKtnHAqUYeebc8z0= -github.com/Azure/azure-sdk-for-go/sdk/data/aztables v1.4.1 h1:j0hhYS006eJ54vusoap0f2NVZ1YY3QnaAEnLM68f0SQ= -github.com/Azure/azure-sdk-for-go/sdk/data/aztables v1.4.1/go.mod h1:AdtInaXmK8eYmbjezRWgLz+Qs46nc9Up9GWGwteWNfw= -github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.2 h1:9iefClla7iYpfYWdzPCRDozdmndjTm8DXdpCzPajMgA= -github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.2/go.mod h1:XtLgD3ZD34DAaVIIAyG3objl5DynM3CQ/vMcbBNJZGI= github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/storage/armstorage v1.8.1 h1:/Zt+cDPnpC3OVDm/JKLOs7M2DKmLRIIp3XIx9pHHiig= github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/storage/armstorage v1.8.1/go.mod h1:Ng3urmn6dYe8gnbCMoHHVl5APYz2txho3koEkV2o2HA= -github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.6.4 h1:jWQK1GI+LeGGUKBADtcH2rRqPxYB1Ljwms5gFA2LqrM= -github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.6.4/go.mod h1:8mwH4klAm9DUgR2EEHyEEAQlRDvLPyg5fQry3y+cDew= -github.com/Azure/azure-sdk-for-go/sdk/storage/azqueue v1.0.1 h1:qvrrnQ2mIjwY7IVlQuNB0ma43Nr74+9ZTZJ60KlmlV4= -github.com/Azure/azure-sdk-for-go/sdk/storage/azqueue v1.0.1/go.mod h1:FkF/Az07vR3S4sBdjCuisznWfFWOD8u6Ibm/g/oyDAk= github.com/AzureAD/microsoft-authentication-library-for-go v1.6.0 h1:XRzhVemXdgvJqCH0sFfrBUTnUJSBrBf7++ypk+twtRs= github.com/AzureAD/microsoft-authentication-library-for-go v1.6.0/go.mod h1:HKpQxkWaGLJ+D/5H8QRpyQXA1eKjxkFlOMwck5+33Jk= github.com/Netflix/go-expect v0.0.0-20180615182759-c93bf25de8e8/go.mod h1:oX5x61PbNXchhh0oikYAH+4Pcfw5LKv21+Jnpr6r6Pc= github.com/Netflix/go-expect v0.0.0-20190729225929-0e00d9168667/go.mod h1:oX5x61PbNXchhh0oikYAH+4Pcfw5LKv21+Jnpr6r6Pc= github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 h1:DklsrG3dyBCFEj5IhUbnKptjxatkF07cF2ak3yi77so= github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2/go.mod h1:WaHUgvxTVq04UNunO+XhnAqY/wQc+bxr74GqbsZ/Jqw= -github.com/atsika/aznet v0.0.0-20260216144450-1395eae946cd h1:tvxn6W9ky7ROcX5M3SbIHCDdypnQRPiBC/3K/+3+sBo= -github.com/atsika/aznet v0.0.0-20260216144450-1395eae946cd/go.mod h1:0+F6U1Bi8CrghJ51eAY1J7B9pCG5FCJc930Dm1Jqh1I= +github.com/atsika/aznet v0.0.0-20260319000323-1bb45b486a9e h1:a78RRI3y1Gsa1rYJztcPPDvS3Fn95Ty6QTl9y2OPmQk= +github.com/atsika/aznet v0.0.0-20260319000323-1bb45b486a9e/go.mod h1:0+F6U1Bi8CrghJ51eAY1J7B9pCG5FCJc930Dm1Jqh1I= github.com/chzyer/logex v1.1.10 h1:Swpa1K6QvQznwJRcfTfQJmTE72DqScAa40E+fbHEXEE= github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWRnGsAI= github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1 h1:q763qf9huN11kDQavWsoZXJNW3xEE4JJyHa5Q25/sd8= diff --git a/pkg/azure-sdk b/pkg/azure-sdk new file mode 160000 index 0000000..2bbb11c --- /dev/null +++ b/pkg/azure-sdk @@ -0,0 +1 @@ +Subproject commit 2bbb11c8970faaff012386e3859b7bfb96872c13 diff --git a/pkg/protocol/base.go b/pkg/protocol/base.go index 148bd3b..5f67413 100644 --- a/pkg/protocol/base.go +++ b/pkg/protocol/base.go @@ -41,6 +41,9 @@ type BaseHandler struct { // conn handles underlying packet transmission (direct net.Conn) conn net.Conn + // writeCh buffers encoded packets for the writeLoop goroutine + writeCh chan []byte + // Connections maps UUIDs to active Connection objects Connections sync.Map @@ -64,11 +67,14 @@ func NewBaseHandler(parentCtx context.Context, conn net.Conn) *BaseHandler { parentCtx = context.Background() } ctx, cancel := context.WithCancel(parentCtx) - return &BaseHandler{ - conn: conn, - Ctx: ctx, - Cancel: cancel, + h := &BaseHandler{ + conn: conn, + writeCh: make(chan []byte, 1024), + Ctx: ctx, + Cancel: cancel, } + go h.writeLoop() + return h } // ReceiveLoop processes incoming packets until EOF or context cancellation. @@ -127,18 +133,21 @@ func (h *BaseHandler) ReceiveLoop() { continue } - packet := Decode(data) - if packet == nil { - continue - } - - errCode := h.handlePacket(packet) - if errCode != ErrNone { - if h.Ctx.Err() != nil { - continue + // Decode all packets in the message (write coalescing may batch multiple) + for len(data) >= HeaderSize { + packet, rest := DecodeNext(data) + if packet == nil { + break } - // Async close dispatch to avoid blocking ReceiveLoop on aznet writes - go h.SendClose(packet.ConnectionID, errCode) + errCode := h.handlePacket(packet) + if errCode != ErrNone { + if h.Ctx.Err() != nil { + break + } + // Async close dispatch to avoid blocking ReceiveLoop on aznet writes + go h.SendClose(packet.ConnectionID, errCode) + } + data = rest } } } @@ -196,8 +205,8 @@ func (h *BaseHandler) SendClose(connectionID uuid.UUID, errCode byte) byte { return h.sendPacket(CmdClose, connectionID, []byte{errCode}) } -// sendPacket is the internal method that encodes and sends all packet types. -// It handles error checking and context checking. +// sendPacket encodes a packet and submits it to the write coalescing channel. +// The actual aznet write happens asynchronously in writeLoop. func (h *BaseHandler) sendPacket(cmd byte, connectionID uuid.UUID, data []byte) byte { if h.Ctx.Err() != nil { return ErrHandlerStopped @@ -213,15 +222,44 @@ func (h *BaseHandler) sendPacket(cmd byte, connectionID uuid.UUID, data []byte) return ErrInvalidPacket } - _, err := h.conn.Write(encoded) - if err != nil { - if err == io.EOF { - return ErrTransportClosed - } - return ErrTransportError + select { + case h.writeCh <- encoded: + return ErrNone + case <-h.Ctx.Done(): + return ErrHandlerStopped } +} - return ErrNone +// writeLoop coalesces queued packets and writes them in batches to aznet. +// Only this goroutine calls conn.Write, eliminating fmu contention. +func (h *BaseHandler) writeLoop() { + for { + var buf []byte + + // Wait for the first packet + select { + case first := <-h.writeCh: + buf = append(buf, first...) + case <-h.Ctx.Done(): + return + } + + // Drain all queued packets (non-blocking) + for { + select { + case more := <-h.writeCh: + buf = append(buf, more...) + default: + goto flush + } + } + + flush: + if _, err := h.conn.Write(buf); err != nil { + h.Cancel() + return + } + } } func (h *BaseHandler) CloseAllConnections() { diff --git a/pkg/protocol/connection.go b/pkg/protocol/connection.go index 98c65f4..3eba42f 100644 --- a/pkg/protocol/connection.go +++ b/pkg/protocol/connection.go @@ -48,7 +48,7 @@ func NewConnection(id uuid.UUID) *Connection { return &Connection{ ID: id, Closed: make(chan struct{}), - deliverCh: make(chan []byte, 256), + deliverCh: make(chan []byte, 1024), CreatedAt: time.Now(), LastActivity: time.Now(), } diff --git a/pkg/protocol/netconn.go b/pkg/protocol/netconn.go index 8f1894c..aa81739 100644 --- a/pkg/protocol/netconn.go +++ b/pkg/protocol/netconn.go @@ -29,7 +29,7 @@ func NewProtocolConn(ctx context.Context, id uuid.UUID, handler *BaseHandler) *P return &ProtocolConn{ id: id, handler: handler, - readBuffer: make(chan []byte, 100), // Buffered channel for received data + readBuffer: make(chan []byte, 1024), // Buffered channel for received data closed: make(chan struct{}), ctx: ctx, } diff --git a/pkg/protocol/protocol.go b/pkg/protocol/protocol.go index 7a4e097..3a0f90d 100644 --- a/pkg/protocol/protocol.go +++ b/pkg/protocol/protocol.go @@ -112,3 +112,35 @@ func Decode(data []byte) *Packet { return NewPacket(command, id, packetData) } + +// DecodeNext parses the first packet from a buffer that may contain multiple +// concatenated packets (from write coalescing). Returns the decoded packet and +// the remaining unconsumed bytes. Returns nil packet if data is incomplete or invalid. +func DecodeNext(data []byte) (*Packet, []byte) { + if len(data) < HeaderSize { + return nil, data + } + + command := data[0] + if command < CmdNew || command > CmdClose { + return nil, data + } + + var id uuid.UUID + copy(id[:], data[CommandSize:CommandSize+UUIDSize]) + + dataLength := binary.BigEndian.Uint32(data[CommandSize+UUIDSize : HeaderSize]) + + totalSize := HeaderSize + int(dataLength) + if len(data) < totalSize { + return nil, data + } + + var packetData []byte + if dataLength > 0 { + packetData = make([]byte, dataLength) + copy(packetData, data[HeaderSize:totalSize]) + } + + return NewPacket(command, id, packetData), data[totalSize:] +} diff --git a/pkg/proxy/server/server.go b/pkg/proxy/server/server.go index 71e5abc..7b4a8cf 100644 --- a/pkg/proxy/server/server.go +++ b/pkg/proxy/server/server.go @@ -183,7 +183,7 @@ func (s *ProxyServer) handleConnection(clientConn net.Conn) { } // 2. Wait for agent acknowledgment with timeout (ProtocolConn will be created in OnAck) - timeout := time.After(5 * time.Second) + timeout := time.After(30 * time.Second) ticker := time.NewTicker(50 * time.Millisecond) defer ticker.Stop() diff --git a/pkg/proxy/socks/connect.go b/pkg/proxy/socks/connect.go index e5152e3..fa9ca8b 100644 --- a/pkg/proxy/socks/connect.go +++ b/pkg/proxy/socks/connect.go @@ -6,7 +6,6 @@ import ( "fmt" "io" "net" - "time" "proxyblob/pkg/protocol" ) @@ -40,8 +39,8 @@ func (h *SocksHandler) handleConnect(conn *protocol.Connection, cmdData []byte) fmt.Printf("[CONNECT] %s\n", target) - // Establish TCP connection to target - targetConn, err := net.DialTimeout("tcp", target, 10*time.Second) + // Establish TCP connection to target (uses Bun bridge on WASM, native net on other platforms) + targetConn, err := dialTCP(target) if err != nil { // Map network error to appropriate protocol error code errCode = protocol.MapNetError(err) diff --git a/pkg/proxy/socks/dial_js.go b/pkg/proxy/socks/dial_js.go new file mode 100644 index 0000000..c903d9e --- /dev/null +++ b/pkg/proxy/socks/dial_js.go @@ -0,0 +1,115 @@ +//go:build js + +package proxy + +import ( + "fmt" + "io" + "net" + "syscall/js" + "time" +) + +type jsConn struct { + socket js.Value + pr *io.PipeReader + pw *io.PipeWriter + closed chan struct{} +} + +func dialTCP(target string) (net.Conn, error) { + host, port, err := net.SplitHostPort(target) + if err != nil { + return nil, err + } + + pr, pw := io.Pipe() + conn := &jsConn{ + pr: pr, + pw: pw, + closed: make(chan struct{}), + } + + done := make(chan error, 1) + + onConnect := js.FuncOf(func(_ js.Value, args []js.Value) any { + conn.socket = args[0] + select { + case done <- nil: + default: + } + return nil + }) + + onData := js.FuncOf(func(_ js.Value, args []js.Value) any { + jsData := args[1] + data := make([]byte, jsData.Length()) + js.CopyBytesToGo(data, jsData) + go pw.Write(data) // don't block the JS callback + return nil + }) + + onClose := js.FuncOf(func(_ js.Value, args []js.Value) any { + pw.CloseWithError(io.EOF) + select { + case <-conn.closed: + default: + close(conn.closed) + } + return nil + }) + + onError := js.FuncOf(func(_ js.Value, args []js.Value) any { + msg := "connection error" + if len(args) > 0 && args[0].Type() == js.TypeString { + msg = args[0].String() + } + e := fmt.Errorf("%s", msg) + pw.CloseWithError(e) + select { + case done <- e: + default: + } + return nil + }) + + js.Global().Call("TCPDial", host, port, onConnect, onData, onClose, onError) + + if err := <-done; err != nil { + return nil, err + } + return conn, nil +} + +func (c *jsConn) Read(b []byte) (int, error) { return c.pr.Read(b) } + +func (c *jsConn) Write(b []byte) (int, error) { + select { + case <-c.closed: + return 0, net.ErrClosed + default: + } + jsData := js.Global().Get("Uint8Array").New(len(b)) + js.CopyBytesToJS(jsData, b) + c.socket.Call("write", jsData) + return len(b), nil +} + +func (c *jsConn) Close() error { + select { + case <-c.closed: + return nil + default: + close(c.closed) + c.pr.CloseWithError(net.ErrClosed) + c.pw.CloseWithError(net.ErrClosed) + c.socket.Call("end") + return nil + } +} + +func (c *jsConn) LocalAddr() net.Addr { return &net.TCPAddr{} } +func (c *jsConn) RemoteAddr() net.Addr { return &net.TCPAddr{} } +func (c *jsConn) SetDeadline(t time.Time) error { return nil } +func (c *jsConn) SetReadDeadline(t time.Time) error { return nil } +func (c *jsConn) SetWriteDeadline(t time.Time) error { return nil } diff --git a/pkg/proxy/socks/dial_native.go b/pkg/proxy/socks/dial_native.go new file mode 100644 index 0000000..4df25c8 --- /dev/null +++ b/pkg/proxy/socks/dial_native.go @@ -0,0 +1,12 @@ +//go:build !js + +package proxy + +import ( + "net" + "time" +) + +func dialTCP(target string) (net.Conn, error) { + return net.DialTimeout("tcp", target, 10*time.Second) +} diff --git a/pkg/proxy/socks/udp_associate.go b/pkg/proxy/socks/udp_associate.go index c5a3927..ef4f601 100644 --- a/pkg/proxy/socks/udp_associate.go +++ b/pkg/proxy/socks/udp_associate.go @@ -4,310 +4,153 @@ import ( "encoding/binary" "errors" "net" + "os" "time" "proxyblob/pkg/protocol" ) // handleUDPAssociate processes the SOCKS5 UDP ASSOCIATE command. -// It creates a UDP relay that allows clients to send and receive -// UDP datagrams through the SOCKS server. -// -// The process involves: -// 1. Creating a UDP socket for client communication -// 2. Sending the socket address back to the client -// 3. Maintaining the TCP control connection -// 4. Relaying UDP datagrams between client and targets -// -// The command format follows RFC 1928 Section 4. +// It creates a single UDP relay socket, tells the client which port to use, +// then dispatches all relay logic to handleUDPPackets. func (h *SocksHandler) handleUDPAssociate(conn *protocol.Connection) byte { - // Create UDP socket - udpAddr, err := net.ResolveUDPAddr("udp", "0.0.0.0:0") + udpConn, err := listenUDP() if err != nil { - // Send network unreachable error h.SendError(conn, protocol.ErrNetworkUnreachable) return protocol.ErrNetworkUnreachable } - udpConn, err := net.ListenUDP("udp", udpAddr) - if err != nil { - // Map network error to appropriate protocol error code - errCode := protocol.MapNetError(err) - h.SendError(conn, errCode) - return errCode - } + port := udpConn.LocalPort() - // Get the allocated port - localAddr := udpConn.LocalAddr().(*net.UDPAddr) - port := localAddr.Port - - // Create response - // Format: |VER|REP|RSV|ATYP|BND.ADDR|BND.PORT| + // Send success response: VER REP RSV ATYP BND.ADDR(4) BND.PORT(2) response := []byte{ - Version5, // VER - Succeeded, // REP - success - 0, // RSV - reserved, must be 0 - IPv4, // ATYP - IPv4 - 0, 0, 0, 0, // BND.ADDR - 0.0.0.0 (any address) - byte(port >> 8), // BND.PORT - high byte - byte(port & 0xff), // BND.PORT - low byte + Version5, Succeeded, 0, IPv4, + 0, 0, 0, 0, + byte(port >> 8), byte(port & 0xff), } - - errCode := h.SendData(conn.ID, response) - if errCode != protocol.ErrNone { + if errCode := h.SendData(conn.ID, response); errCode != protocol.ErrNone { udpConn.Close() return protocol.ErrPacketSendFailed } - // Store UDP connection and start handling packets (state is already established via ProtocolConn) - conn.Conn = udpConn + go h.handleUDPPackets(conn, udpConn) - go h.handleUDPPackets(conn) - - // Keep the control connection open until it's closed elsewhere + // Hold the control connection open; close the socket when context dies. select { case <-conn.Closed: - // Control connection closed, UDP associate terminated case <-h.Ctx.Done(): udpConn.Close() } - return protocol.ErrNone } -// handleUDPPackets manages the UDP relay for a client. -// It: -// - Receives UDP packets from the client -// - Extracts target addresses from SOCKS headers -// - Forwards packets to targets -// - Receives responses from targets -// - Wraps responses in SOCKS headers -// - Returns them to the client +// handleUDPPackets relays UDP datagrams between the SOCKS client and internet targets. // -// The relay operates until the control connection closes or -// the context is canceled. -func (h *SocksHandler) handleUDPPackets(conn *protocol.Connection) { - udpConn := conn.Conn.(*net.UDPConn) - buffer := make([]byte, 16*1024) // Reduced from 64KB to 16KB +// A single socket handles both directions: +// - Packets from clientAddr → strip SOCKS5 header → forward to target +// - Packets from a known target → wrap in SOCKS5 header → forward to clientAddr +func (h *SocksHandler) handleUDPPackets(conn *protocol.Connection, udpConn UDPRelayConn) { + defer udpConn.Close() + + buf := make([]byte, 16*1024) var clientAddr *net.UDPAddr - // Create a UDP connection for sending to targets - targetConn, err := net.ListenUDP("udp", nil) - if err != nil { - h.SendClose(conn.ID, protocol.ErrNetworkUnreachable) - return - } - - // Ensure connections are properly closed when this handler exits - defer targetConn.Close() - - // Map to track target addresses and their corresponding responses type targetInfo struct { addr *net.UDPAddr lastActive time.Time } targets := make(map[string]*targetInfo) - // Channel for receiving UDP responses - responses := make(chan struct { - data []byte - addr *net.UDPAddr - }, 100) - - // Start a goroutine to handle responses - go func() { - respBuf := make([]byte, 16*1024) // Reduced from 128KB to 16KB - for { - // Check for cancellation - select { - case <-h.Ctx.Done(): - return - case <-conn.Closed: - return - default: - // Continue processing - } - - // Set read deadline to 100ms for faster response (reduced from 300ms) - if err := targetConn.SetReadDeadline(time.Now().Add(100 * time.Millisecond)); err != nil { - if errors.Is(err, net.ErrClosed) { - return - } - continue // Non-fatal, just retry - } - - n, addr, err := targetConn.ReadFromUDP(respBuf) - if err != nil { - if errors.Is(err, net.ErrClosed) { - return - } - continue - } - - if n > 0 { - // Make a copy of the data since respBuf will be reused - data := make([]byte, n) - copy(data, respBuf[:n]) - - // Send response through channel - select { - case responses <- struct { - data []byte - addr *net.UDPAddr - }{data, addr}: - // Successfully sent - default: - // Channel full, log and continue - } - } - } - }() - - // Timeout to clean up inactive targets - ticker := time.NewTicker(30 * time.Second) - defer ticker.Stop() + cleanup := time.NewTicker(30 * time.Second) + defer cleanup.Stop() for { + // Check for shutdown before blocking on I/O. select { case <-h.Ctx.Done(): return case <-conn.Closed: return - case resp := <-responses: - // Find the target that matches this response - var found bool - - for _, target := range targets { - if target.addr.IP.Equal(resp.addr.IP) && target.addr.Port == resp.addr.Port { - // Update last active time - target.lastActive = time.Now() - found = true - break - } - } - - if !found { - // Skip unknown targets - continue - } - - if clientAddr == nil { - // Skip client address not set - continue - } - - // Create response packet with appropriate header - // For simplicity, we'll just use a minimal header with the correct address type - var respHeader []byte - - // Determine address type - var addrType byte - if resp.addr.IP.To4() != nil { - addrType = IPv4 - } else { - addrType = IPv6 - } - - // header: RSV(2) + FRAG(0) + ATYP(1) + ADDR + PORT(2) - respHeader = append(respHeader, 0, 0, 0, addrType) - - if addrType == IPv4 { - respHeader = append(respHeader, resp.addr.IP.To4()...) - } else { - respHeader = append(respHeader, resp.addr.IP.To16()...) - } - - // Add port - portBytes := make([]byte, 2) - binary.BigEndian.PutUint16(portBytes, uint16(resp.addr.Port)) - respHeader = append(respHeader, portBytes...) - - // Combine header and payload - respPacket := append(respHeader, resp.data...) - - _, err = udpConn.WriteToUDP(respPacket, clientAddr) - if err != nil { - if errors.Is(err, net.ErrClosed) { - return - } - } - - case <-ticker.C: - // Clean up inactive targets + case <-cleanup.C: now := time.Now() - for targetKey, targetInfo := range targets { - if now.Sub(targetInfo.lastActive) > 1*time.Minute { - delete(targets, targetKey) + for k, t := range targets { + if now.Sub(t.lastActive) > time.Minute { + delete(targets, k) } } - default: - // Set read deadline to 100ms for faster response (reduced from 300ms) - if err := udpConn.SetReadDeadline(time.Now().Add(100 * time.Millisecond)); err != nil { - if errors.Is(err, net.ErrClosed) { - return - } - continue // Non-fatal + } + + udpConn.SetReadDeadline(time.Now().Add(100 * time.Millisecond)) + n, addr, err := udpConn.ReadFrom(buf) + if err != nil { + if isUDPTimeout(err) { + continue } - - n, remoteAddr, err := udpConn.ReadFromUDP(buffer) - if err != nil { - // Handle various error conditions - if errors.Is(err, net.ErrClosed) { - return - } - - if netErr, ok := err.(net.Error); ok { - if netErr.Timeout() { - // This is expected due to deadline, don't log - continue - } - } - - h.SendClose(conn.ID, protocol.ErrNetworkUnreachable) + if errors.Is(err, net.ErrClosed) { return } + h.SendClose(conn.ID, protocol.ErrNetworkUnreachable) + return + } - // Store client address from first packet - if clientAddr == nil { - clientAddr = remoteAddr - } + if n == 0 { + continue + } - // Only accept packets from original client - if !remoteAddr.IP.Equal(clientAddr.IP) { + if clientAddr == nil { + // First valid SOCKS5 UDP datagram (RSV=0x0000, FRAG=0) sets the client. + if n > 3 && buf[0] == 0 && buf[1] == 0 && buf[2] == 0 { + clientAddr = addr + } else { continue } + } - // Handle UDP packet - if n > 3 { - // Extract target address from UDP header - targetAddr, headerLen, errCode := ExtractUDPHeader(buffer[:n]) - if errCode != protocol.ErrNone { + if addr.IP.Equal(clientAddr.IP) && addr.Port == clientAddr.Port { + // Client → target + if n <= 3 { + continue + } + targetAddr, headerLen, errCode := ExtractUDPHeader(buf[:n]) + if errCode != protocol.ErrNone { + continue + } + targetUDPAddr, err := net.ResolveUDPAddr("udp", targetAddr) + if err != nil { + continue + } + targets[targetAddr] = &targetInfo{addr: targetUDPAddr, lastActive: time.Now()} + udpConn.WriteTo(buf[headerLen:n], targetUDPAddr) + } else { + // Target → client: find the matching target entry and wrap with SOCKS5 header. + for _, t := range targets { + if !t.addr.IP.Equal(addr.IP) || t.addr.Port != addr.Port { continue } + t.lastActive = time.Now() - // Resolve target address - targetUDPAddr, err := net.ResolveUDPAddr("udp", targetAddr) - if err != nil { - continue + var header []byte + if ip4 := addr.IP.To4(); ip4 != nil { + header = append([]byte{0, 0, 0, IPv4}, ip4...) + } else { + header = append([]byte{0, 0, 0, IPv6}, addr.IP.To16()...) } + var portBuf [2]byte + binary.BigEndian.PutUint16(portBuf[:], uint16(addr.Port)) + header = append(header, portBuf[:]...) - // Store target information - targetKey := targetAddr - targets[targetKey] = &targetInfo{ - addr: targetUDPAddr, - lastActive: time.Now(), - } - - // Send payload to target - _, err = targetConn.WriteToUDP(buffer[headerLen:n], targetUDPAddr) - if err != nil { - if errors.Is(err, net.ErrClosed) { - return - } - continue - } + udpConn.WriteTo(append(header, buf[:n]...), clientAddr) + break } } } } + +// isUDPTimeout reports whether err is a read-deadline / timeout error. +func isUDPTimeout(err error) bool { + if netErr, ok := err.(net.Error); ok && netErr.Timeout() { + return true + } + return errors.Is(err, os.ErrDeadlineExceeded) +} diff --git a/pkg/proxy/socks/udp_js.go b/pkg/proxy/socks/udp_js.go new file mode 100644 index 0000000..3babc54 --- /dev/null +++ b/pkg/proxy/socks/udp_js.go @@ -0,0 +1,137 @@ +//go:build js + +package proxy + +import ( + "fmt" + "net" + "os" + "sync" + "syscall/js" + "time" +) + +type jsUDPPacket struct { + data []byte + addr *net.UDPAddr +} + +type jsUDPConn struct { + socket js.Value + port int + packets chan jsUDPPacket + closed chan struct{} + once sync.Once + mu sync.Mutex + dl time.Time // read deadline +} + +func listenUDP() (UDPRelayConn, error) { + c := &jsUDPConn{ + packets: make(chan jsUDPPacket, 256), + closed: make(chan struct{}), + } + + done := make(chan error, 1) + + onBind := js.FuncOf(func(_ js.Value, args []js.Value) any { + c.socket = args[0] + c.port = args[1].Int() + select { + case done <- nil: + default: + } + return nil + }) + + onData := js.FuncOf(func(_ js.Value, args []js.Value) any { + // args: socket, Uint8Array, port, address + jsData := args[1] + port := args[2].Int() + host := args[3].String() + + data := make([]byte, jsData.Length()) + js.CopyBytesToGo(data, jsData) + + addr := &net.UDPAddr{IP: net.ParseIP(host), Port: port} + select { + case c.packets <- jsUDPPacket{data, addr}: + case <-c.closed: + default: + } + return nil + }) + + onError := js.FuncOf(func(_ js.Value, args []js.Value) any { + msg := "udp socket error" + if len(args) > 0 { + msg = args[0].String() + } + select { + case done <- fmt.Errorf("%s", msg): + default: + } + return nil + }) + + js.Global().Call("UDPListen", onBind, onData, onError) + + if err := <-done; err != nil { + return nil, err + } + return c, nil +} + +func (c *jsUDPConn) LocalPort() int { return c.port } + +func (c *jsUDPConn) ReadFrom(b []byte) (int, *net.UDPAddr, error) { + c.mu.Lock() + dl := c.dl + c.mu.Unlock() + + var timer <-chan time.Time + if !dl.IsZero() { + d := time.Until(dl) + if d <= 0 { + return 0, nil, &net.OpError{Op: "read", Net: "udp", Err: os.ErrDeadlineExceeded} + } + timer = time.After(d) + } + + select { + case pkt := <-c.packets: + n := copy(b, pkt.data) + return n, pkt.addr, nil + case <-c.closed: + return 0, nil, net.ErrClosed + case <-timer: + return 0, nil, &net.OpError{Op: "read", Net: "udp", Err: os.ErrDeadlineExceeded} + } +} + +func (c *jsUDPConn) WriteTo(b []byte, addr *net.UDPAddr) error { + select { + case <-c.closed: + return net.ErrClosed + default: + } + jsData := js.Global().Get("Uint8Array").New(len(b)) + js.CopyBytesToJS(jsData, b) + c.socket.Call("send", jsData, addr.Port, addr.IP.String()) + return nil +} + +func (c *jsUDPConn) SetReadDeadline(t time.Time) error { + c.mu.Lock() + c.dl = t + c.mu.Unlock() + return nil +} + +func (c *jsUDPConn) Close() error { + c.once.Do(func() { + close(c.closed) + c.socket.Call("close") + }) + return nil +} diff --git a/pkg/proxy/socks/udp_native.go b/pkg/proxy/socks/udp_native.go new file mode 100644 index 0000000..ec71efd --- /dev/null +++ b/pkg/proxy/socks/udp_native.go @@ -0,0 +1,41 @@ +//go:build !js + +package proxy + +import ( + "net" + "time" +) + +type nativeUDPConn struct { + conn *net.UDPConn +} + +func listenUDP() (UDPRelayConn, error) { + c, err := net.ListenUDP("udp", &net.UDPAddr{}) + if err != nil { + return nil, err + } + return &nativeUDPConn{c}, nil +} + +func (c *nativeUDPConn) LocalPort() int { + return c.conn.LocalAddr().(*net.UDPAddr).Port +} + +func (c *nativeUDPConn) ReadFrom(b []byte) (int, *net.UDPAddr, error) { + return c.conn.ReadFromUDP(b) +} + +func (c *nativeUDPConn) WriteTo(b []byte, addr *net.UDPAddr) error { + _, err := c.conn.WriteToUDP(b, addr) + return err +} + +func (c *nativeUDPConn) SetReadDeadline(t time.Time) error { + return c.conn.SetReadDeadline(t) +} + +func (c *nativeUDPConn) Close() error { + return c.conn.Close() +} diff --git a/pkg/proxy/socks/udp_relay.go b/pkg/proxy/socks/udp_relay.go new file mode 100644 index 0000000..44d3b19 --- /dev/null +++ b/pkg/proxy/socks/udp_relay.go @@ -0,0 +1,28 @@ +package proxy + +import ( + "net" + "time" +) + +// UDPRelayConn is a UDP socket that can send and receive datagrams to/from +// arbitrary addresses. Used by the SOCKS5 UDP ASSOCIATE handler. +// Implementations must be safe for concurrent use. +type UDPRelayConn interface { + // LocalPort returns the port this socket is bound to. + LocalPort() int + + // ReadFrom reads a datagram into b, returning the sender's address. + // Blocks until data arrives, the socket is closed, or the read deadline fires. + ReadFrom(b []byte) (int, *net.UDPAddr, error) + + // WriteTo sends a datagram to addr. + WriteTo(b []byte, addr *net.UDPAddr) error + + // SetReadDeadline sets the deadline for future ReadFrom calls. + // A zero value disables the deadline. + SetReadDeadline(t time.Time) error + + // Close closes the socket. + Close() error +}