fix socks issue

This commit is contained in:
KingKDot
2026-05-18 10:10:19 -05:00
parent 95f7db5dd5
commit b8121ade9a
2 changed files with 60 additions and 6 deletions
+2 -1
View File
@@ -5,6 +5,7 @@ import (
"fmt" "fmt"
"log" "log"
"net" "net"
"strconv"
"sync" "sync"
"time" "time"
@@ -63,7 +64,7 @@ func HandleProxyConnect(ctx context.Context, env *runtime.Env, connID string, pa
}) })
} }
target := fmt.Sprintf("%s:%d", host, port) target := net.JoinHostPort(host, strconv.Itoa(port))
goSafe("tunnel-dial-"+connID, nil, func() { goSafe("tunnel-dial-"+connID, nil, func() {
@@ -38,6 +38,8 @@ type TunnelConnection = {
connected: boolean; connected: boolean;
/** buffer data received from SOCKS client before agent confirms */ /** buffer data received from SOCKS client before agent confirms */
pendingData: Buffer[]; pendingData: Buffer[];
/** data destined for the SOCKS client that the kernel buffer rejected; flushed on drain */
writeQueue: Buffer[];
}; };
type ProxySocketData = { type ProxySocketData = {
@@ -110,6 +112,7 @@ export function startProxy(
socket, socket,
connected: false, connected: false,
pendingData: [], pendingData: [],
writeQueue: [],
}); });
logger.debug( logger.debug(
`[socks5] new connection ${connectionId} on port ${port}`, `[socks5] new connection ${connectionId} on port ${port}`,
@@ -149,6 +152,14 @@ export function startProxy(
logger.debug(`[socks5] connection ${connId} closed`); logger.debug(`[socks5] connection ${connId} closed`);
}, },
drain(socket) {
const entry = activeProxies.get(socket.data.proxyPort);
if (!entry) return;
const tunnel = entry.connections.get(socket.data.connectionId);
if (!tunnel) return;
flushWriteQueue(tunnel);
},
error(socket, err) { error(socket, err) {
logger.error( logger.error(
`[socks5] socket error conn=${socket.data?.connectionId}`, `[socks5] socket error conn=${socket.data?.connectionId}`,
@@ -228,16 +239,50 @@ export function handleProxyTunnelData(
if (entry.clientId !== clientId) continue; if (entry.clientId !== clientId) continue;
const tunnel = entry.connections.get(connectionId); const tunnel = entry.connections.get(connectionId);
if (tunnel) { if (tunnel) {
try { writeToTunnelSocket(tunnel, data);
tunnel.socket.write(data);
} catch {
entry.connections.delete(connectionId);
}
return; return;
} }
} }
} }
function writeToTunnelSocket(
tunnel: TunnelConnection,
data: Buffer | Uint8Array,
): void {
const buf = Buffer.isBuffer(data) ? data : Buffer.from(data);
if (tunnel.writeQueue.length > 0) {
tunnel.writeQueue.push(buf);
return;
}
let written: number;
try {
written = tunnel.socket.write(buf);
} catch {
return;
}
if (written < buf.length) {
tunnel.writeQueue.push(buf.subarray(Math.max(written, 0)));
}
}
function flushWriteQueue(tunnel: TunnelConnection): void {
while (tunnel.writeQueue.length > 0) {
const next = tunnel.writeQueue[0];
let written: number;
try {
written = tunnel.socket.write(next);
} catch {
tunnel.writeQueue.length = 0;
return;
}
if (written < next.length) {
tunnel.writeQueue[0] = next.subarray(Math.max(written, 0));
return;
}
tunnel.writeQueue.shift();
}
}
/** Called when the agent closes its side of a tunnel */ /** Called when the agent closes its side of a tunnel */
export function handleProxyTunnelClose( export function handleProxyTunnelClose(
clientId: string, clientId: string,
@@ -327,6 +372,14 @@ function handleSocksData(
} }
if (buf.length < 2 + nmethods) return; if (buf.length < 2 + nmethods) return;
const methods = buf.subarray(2, 2 + nmethods);
if (!methods.includes(0x00)) {
socket.write(Buffer.from([0x05, 0xff]));
socket.end();
entry.connections.delete(connectionId);
return;
}
socket.write(Buffer.from([0x05, 0x00])); socket.write(Buffer.from([0x05, 0x00]));
socket.data.phase = "request"; socket.data.phase = "request";
socket.data.buffer = buf.subarray(2 + nmethods); socket.data.buffer = buf.subarray(2 + nmethods);