#!/usr/bin/env node

/**
 * PSK SOCKS5 Proxy (Client)
 *
 * - Runs a local SOCKS5 proxy (supports:
 *    - CONNECT for arbitrary TCP protocols (HTTP, HTTPS, SSH, etc.)
 *    - UDP ASSOCIATE for UDP-based protocols (e.g., DNS)
 *    - BIND: not supported, returns error)
 * - Maintains a single TLS-PSK tunnel to a remote out-node (proxy-server.js)
 * - Multiplexes many local client connections over one TLS tunnel using frames:
 *
 * Frame: [1 byte type][4 bytes connection id][4 bytes data length][data...]
 * Types:
 *   DATA (2)
 *   CLOSE (3)
 *   OPEN (4)               payload = [2B hostLen][host utf8][2B port]
 *   OPEN_RESULT (5)        payload = [1B status] 1=success,0=failure
 *
 *   UDP_OPEN (6)           payload = empty
 *   UDP_OPEN_RESULT (7)    payload = [1B status]
 *   UDP_SEND (8)           payload = [2B hostLen][host][2B port][2B dataLen][data...]
 *   UDP_RECV (9)           payload = same as UDP_SEND (host:port,data)
 *   UDP_CLOSE (10)         payload = empty
 */

const net = require('net');
const tls = require('tls');
const fs = require('fs');
const dgram = require('dgram');
const { program } = require('commander');

program
  .requiredOption('-H, --server-host <host>', 'Proxy out-node server host')
  .requiredOption('-P, --server-port <port>', 'Proxy out-node server port')
  .requiredOption('--psk-file <path>', 'Path to binary PSK key file')
  .requiredOption('--identity <identity>', 'PSK identity string')
  .requiredOption('--socks-port <port>', 'Local SOCKS5 proxy port to listen on')
  .option('--bind-host <host>', 'Local bind host', '127.0.0.1')
  .option('--connect-timeout <ms>', 'Timeout waiting OPEN/UDP_OPEN result (ms)', '10000')
  .option('--idle-timeout <ms>', 'Idle timeout for local TCP sockets (ms)', '60000')
  .option('--udp-idle-timeout <ms>', 'Idle timeout for UDP associate socket (ms, 0=never)', '60000')
  .parse();

const options = program.opts();

let pskKey;
try {
  pskKey = fs.readFileSync(options.pskFile);
} catch (error) {
  console.error(`Error reading PSK file: ${error.message}`);
  process.exit(1);
}

const OPEN_RESULT_TIMEOUT = parseInt(options.connectTimeout, 10) || 10000;
const IDLE_TIMEOUT = parseInt(options.idleTimeout, 10) || 60000;
const UDP_IDLE_TIMEOUT = parseInt(options.udpIdleTimeout, 10) || 60000;

// Message Types
const MSG_TYPES = {
  DATA: 2,
  CLOSE: 3,
  OPEN: 4,
  OPEN_RESULT: 5,

  UDP_OPEN: 6,
  UDP_OPEN_RESULT: 7,
  UDP_SEND: 8,
  UDP_RECV: 9,
  UDP_CLOSE: 10,
};

function buildOpenPayload(host, port) {
  const hostBuf = Buffer.from(host, 'utf8');
  const buf = Buffer.allocUnsafe(2 + hostBuf.length + 2);
  let off = 0;
  buf.writeUInt16BE(hostBuf.length, off); off += 2;
  hostBuf.copy(buf, off); off += hostBuf.length;
  buf.writeUInt16BE(port, off); off += 2;
  return buf;
}

function buildUdpPayload(host, port, data) {
  const hostBuf = Buffer.from(host, 'utf8');
  const buf = Buffer.allocUnsafe(2 + hostBuf.length + 2 + 2 + data.length);
  let off = 0;
  buf.writeUInt16BE(hostBuf.length, off); off += 2;
  hostBuf.copy(buf, off); off += hostBuf.length;
  buf.writeUInt16BE(port, off); off += 2;
  buf.writeUInt16BE(data.length, off); off += 2;
  data.copy(buf, off);
  return buf;
}

function writeMessage(socket, type, connectionId, data = Buffer.alloc(0)) {
  if (!socket || socket.destroyed) return true;
  const buf = Buffer.allocUnsafe(9 + data.length);
  buf.writeUInt8(type, 0);
  buf.writeUInt32BE(connectionId >>> 0, 1);
  buf.writeUInt32BE(data.length >>> 0, 5);
  if (data.length > 0) data.copy(buf, 9);
  return socket.write(buf);
}

function createMessageReader() {
  let buffer = Buffer.alloc(0);
  let expectedLength = 9;
  let currentMessage = null;

  return function onData(data, callback) {
    buffer = Buffer.concat([buffer, data]);

    while (buffer.length >= expectedLength) {
      if (!currentMessage) {
        const type = buffer.readUInt8(0);
        const connectionId = buffer.readUInt32BE(1);
        const dataLength = buffer.readUInt32BE(5);
        currentMessage = { type, connectionId, dataLength };
        expectedLength = 9 + dataLength;

        if (dataLength === 0) {
          callback(currentMessage.type, currentMessage.connectionId, Buffer.alloc(0));
          buffer = buffer.subarray(9);
          currentMessage = null;
          expectedLength = 9;
        }
      } else {
        const messageData = buffer.subarray(9, expectedLength);
        callback(currentMessage.type, currentMessage.connectionId, messageData);
        buffer = buffer.subarray(expectedLength);
        currentMessage = null;
        expectedLength = 9;
      }
    }
  };
}

function pskCallback(/* hint */) {
  return {
    identity: options.identity,
    psk: pskKey,
  };
}

// Global tunnel socket (single multiplexed connection)
let tunnelSocket = null;
let reader = null;

// TCP connection state: id -> { id, localSocket, opened, openTimer }
const connections = new Map();
let nextConnId = 1;

// UDP association state: udpId -> {
//    id, controlSocket, udpSocket, openTimer, opened,
//    clients: Set of "host:port" strings to reply to,
// }
const udpAssocs = new Map();

// Framing and fairness controls for local->tunnel direction
const FRAME_MAX = 16 * 1024; // 16KiB frames to improve interleaving
const BYTES_PER_TICK = 64 * 1024; // Limit per processing burst to yield event loop
// Per-connection TX queues: connectionId -> { queue: Buffer[], sending: boolean }
const txQueues = new Map();

function processTxQueue(connectionId, state) {
  if (!tunnelSocket || tunnelSocket.destroyed) {
    state.queue = [];
    state.sending = false;
    return;
  }
  if (state.sending) return;
  state.sending = true;

  let bytesThisTick = 0;

  const run = () => {
    if (!tunnelSocket || tunnelSocket.destroyed) {
      state.queue = [];
      state.sending = false;
      return;
    }
    while (state.queue.length) {
      const buf = state.queue[0];
      let offset = buf._offset || 0;

      while (offset < buf.length) {
        const end = Math.min(offset + FRAME_MAX, buf.length);
        const slice = buf.subarray(offset, end);
        const ok = writeMessage(tunnelSocket, MSG_TYPES.DATA, connectionId, slice);
        if (!ok) {
          buf._offset = end;
          tunnelSocket.once('drain', () => {
            state.sending = false;
            run();
          });
          return;
        }
        offset = end;
        bytesThisTick += slice.length;
        if (bytesThisTick >= BYTES_PER_TICK) {
          buf._offset = offset;
          bytesThisTick = 0;
          setImmediate(() => {
            state.sending = false;
            run();
          });
          return;
        }
      }

      delete buf._offset;
      state.queue.shift();
    }

    state.sending = false;
  };

  run();
}

function enqueueToTunnel(connectionId, data) {
  let state = txQueues.get(connectionId);
  if (!state) {
    state = { queue: [], sending: false };
    txQueues.set(connectionId, state);
  }
  state.queue.push(data);
  if (!state.sending) {
    state.sending = false;
    processTxQueue(connectionId, state);
  }
}

function genConnId() {
  let id = nextConnId >>> 0;
  do {
    id = (id + 1) >>> 0;
    if (id === 0) id = 1;
  } while (connections.has(id) || udpAssocs.has(id));
  nextConnId = id;
  return id;
}

function connectTunnel() {
  const host = options.serverHost;
  const port = parseInt(options.serverPort, 10);
  console.log(`Connecting to proxy out-node ${host}:${port} via TLS-PSK...`);

  const sock = tls.connect(
    {
      host,
      port,
      pskCallback,
      ciphers: 'PSK-AES256-GCM-SHA384:PSK-AES128-GCM-SHA256',
      checkServerIdentity: () => undefined,
    },
    () => {
      console.log('Proxy tunnel connected');
    }
  );

  sock.setNoDelay(true);
  sock.setKeepAlive(true, 30000);

  tunnelSocket = sock;
  reader = createMessageReader();

  sock.on('data', (data) => {
    reader(data, (type, connectionId, payload) => {
      if (type === MSG_TYPES.OPEN_RESULT) {
        const st = connections.get(connectionId);
        if (!st) return;
        const ok = payload.length > 0 ? payload.readUInt8(0) === 1 : false;
        if (st.opened) return;
        clearTimeout(st.openTimer);
        if (!ok) {
          // SOCKS5 CONNECT failure -> reply and close
          try {
            const rep = buildSocksReply(0x05 /* connection refused */, '0.0.0.0', 0);
            st.localSocket.write(rep);
          } catch (_) {}
          destroyLocalTCP(connectionId);
          return;
        }
        st.opened = true;
        // Reply success for CONNECT
        try {
          const rep = buildSocksReply(0x00 /* succeeded */, '0.0.0.0', 0);
          st.localSocket.write(rep);
        } catch (_) {}
        // Switch to streaming mode
        armTcpStreaming(st);
      } else if (type === MSG_TYPES.DATA) {
        const st = connections.get(connectionId);
        if (st && !st.localSocket.destroyed) {
          st.localSocket.write(payload);
        }
      } else if (type === MSG_TYPES.CLOSE) {
        destroyLocalTCP(connectionId);

      } else if (type === MSG_TYPES.UDP_OPEN_RESULT) {
        const as = udpAssocs.get(connectionId);
        if (!as) return;
        const ok = payload.length > 0 ? payload.readUInt8(0) === 1 : false;
        clearTimeout(as.openTimer);
        if (!ok) {
          // Tell SOCKS5 client UDP associate failed
          try {
            const rep = buildSocksReply(0x01 /* general failure */, '0.0.0.0', 0);
            as.controlSocket.write(rep);
          } catch (_) {}
          closeUdpAssoc(connectionId);
          try { as.controlSocket.destroy(); } catch (_) {}
          udpAssocs.delete(connectionId);
          return;
        }
        as.opened = true;
        // Reply with our UDP bind address to client
        const addrInfo = as.udpSocket.address();
        const bndAddr = isIPv6Address(addrInfo.address) ? '::' : addrInfo.address || '0.0.0.0';
        const bndPort = addrInfo.port || 0;
        try {
          const rep = buildSocksReply(0x00 /* succeeded */, bndAddr, bndPort);
          as.controlSocket.write(rep);
        } catch (_) {}
        // Keep control connection open; UDP messages handled on udpSocket
      } else if (type === MSG_TYPES.UDP_RECV) {
        // Payload: [hostLen][host][port][dataLen][data...]
        const parsed = parseUdpPayload(payload);
        if (!parsed) return;
        const as = udpAssocs.get(connectionId);
        if (!as) return;
        // Build SOCKS5 UDP response and send to all known client endpoints
        const { host, port, data } = parsed;
        const udpPacket = buildSocksUdpDatagram(host, port, data);
        for (const ep of as.clients) {
          const [caddr, cportStr] = ep.split(':');
          const cport = parseInt(cportStr, 10);
          try {
            as.udpSocket.send(udpPacket, cport, caddr);
          } catch (_) {}
        }
      } else {
        // ignore unknown message types
      }
    });
  });

  sock.on('close', () => {
    console.log('Proxy tunnel closed. Cleaning up local connections and retrying in 2s...');
    // Close all local TCP connections
    for (const [id] of connections) {
      destroyLocalTCP(id);
    }
    // Close all UDP assocs
    for (const [id] of udpAssocs) {
      closeUdpAssoc(id);
    }
    tunnelSocket = null;
    setTimeout(connectTunnel, 2000);
  });

  sock.on('error', (err) => {
    console.error('Tunnel error:', err.message);
  });
}

function safeWrite(socket, buf) {
  if (!socket || socket.destroyed) return;
  try { socket.write(buf); } catch (_) {}
}

function destroyLocalTCP(connectionId) {
  const st = connections.get(connectionId);
  if (!st) return;
  connections.delete(connectionId);
  clearTimeout(st.openTimer);
  txQueues.delete(connectionId);
  try { st.localSocket.destroy(); } catch (_) {}
}

function startOpenTcp(connectionId, host, port, st) {
  if (!tunnelSocket || tunnelSocket.destroyed) {
    // Reply SOCKS failure
    try {
      const rep = buildSocksReply(0x01 /* general failure */, '0.0.0.0', 0);
      st.localSocket.write(rep);
    } catch (_) {}
    destroyLocalTCP(connectionId);
    return;
  }

  st.openTimer = setTimeout(() => {
    if (!st.opened) {
      try {
        const rep = buildSocksReply(0x06 /* TTL expired (timeout) */, '0.0.0.0', 0);
        st.localSocket.write(rep);
      } catch (_) {}
      writeMessage(tunnelSocket, MSG_TYPES.CLOSE, connectionId);
      destroyLocalTCP(connectionId);
    }
  }, OPEN_RESULT_TIMEOUT);

  writeMessage(tunnelSocket, MSG_TYPES.OPEN, connectionId, buildOpenPayload(host, port));
}

function armTcpStreaming(st) {
  // After CONNECT success, forward subsequent data with framing/backpressure
  const forward = (chunk) => {
    if (!tunnelSocket || tunnelSocket.destroyed) return;
    enqueueToTunnel(st.id, chunk);
  };
  st.localSocket.on('data', forward);
  st.localSocket.once('close', () => {
    writeMessage(tunnelSocket, MSG_TYPES.CLOSE, st.id);
    connections.delete(st.id);
    clearTimeout(st.openTimer);
    txQueues.delete(st.id);
  });
  st.localSocket.once('error', () => {
    writeMessage(tunnelSocket, MSG_TYPES.CLOSE, st.id);
    connections.delete(st.id);
    clearTimeout(st.openTimer);
    txQueues.delete(st.id);
  });
}

/* ============== SOCKS5 Helpers ============== */

function isIPv4Address(addr) {
  return net.isIPv4(addr);
}
function isIPv6Address(addr) {
  return net.isIPv6(addr);
}

function buildSocksReply(rep, bndAddr, bndPort) {
  // VER=5, REP, RSV=0, ATYP, BND.ADDR, BND.PORT
  const addrBufInfo = encodeSocksAddr(bndAddr);
  const buf = Buffer.allocUnsafe(4 + addrBufInfo.addr.length + 2);
  let off = 0;
  buf.writeUInt8(5, off++);       // VER
  buf.writeUInt8(rep, off++);     // REP
  buf.writeUInt8(0, off++);       // RSV
  buf.writeUInt8(addrBufInfo.atyp, off++); // ATYP
  addrBufInfo.addr.copy(buf, off); off += addrBufInfo.addr.length;
  buf.writeUInt16BE(bndPort >>> 0, off); off += 2;
  return buf;
}

function encodeSocksAddr(addr) {
  if (isIPv4Address(addr)) {
    return { atyp: 0x01, addr: Buffer.from(addr.split('.').map((x) => parseInt(x, 10))) };
  }
  if (isIPv6Address(addr)) {
    return { atyp: 0x04, addr: ipToBuffer(addr) };
  }
  const nameBuf = Buffer.from(addr, 'utf8');
  const len = Math.min(nameBuf.length, 255);
  const nb = nameBuf.subarray(0, len);
  return { atyp: 0x03, addr: Buffer.concat([Buffer.from([nb.length]), nb]) };
}

function ipToBuffer(ip) {
  // Convert IPv6 or IPv4-mapped IPv6 to buffer
  if (isIPv4Address(ip)) {
    return Buffer.from(ip.split('.').map((x) => parseInt(x, 10)));
  }
  // For IPv6 literals, use URL parser trick
  // Node doesn't provide direct conversion; implement simple parser:
  const sections = ip.split('::');
  let head = [];
  let tail = [];
  if (sections.length === 1) {
    head = ip.split(':');
  } else if (sections.length === 2) {
    head = sections[0] ? sections[0].split(':') : [];
    tail = sections[1] ? sections[1].split(':') : [];
  } else {
    // invalid
    return Buffer.alloc(16);
  }
  // Expand head and tail into 8 blocks
  const total = head.length + tail.length;
  const fill = 8 - total;
  const parts = [...head, ...Array(fill).fill('0'), ...tail].map((p) => parseInt(p || '0', 16));
  const buf = Buffer.alloc(16);
  for (let i = 0; i < 8; i++) {
    buf.writeUInt16BE(parts[i] || 0, i * 2);
  }
  return buf;
}

function parseSocksRequest(buf) {
  // VER CMD RSV ATYP DST.ADDR DST.PORT
  if (buf.length < 4) return null;
  const ver = buf.readUInt8(0);
  if (ver !== 5) return { error: 'VERSION' };
  const cmd = buf.readUInt8(1);
  const rsv = buf.readUInt8(2);
  const atyp = buf.readUInt8(3);
  let off = 4;
  let host = null;
  if (atyp === 0x01) {
    if (buf.length < off + 4 + 2) return null;
    host = Array.from(buf.subarray(off, off + 4)).join('.');
    off += 4;
  } else if (atyp === 0x03) {
    if (buf.length < off + 1) return null;
    const len = buf.readUInt8(off); off += 1;
    if (buf.length < off + len + 2) return null;
    host = buf.subarray(off, off + len).toString('utf8');
    off += len;
  } else if (atyp === 0x04) {
    if (buf.length < off + 16 + 2) return null;
    const b = buf.subarray(off, off + 16);
    host = [...Array(8).keys()].map((i) => b.readUInt16BE(i * 2).toString(16)).join(':');
    off += 16;
  } else {
    return { error: 'ATYP' };
  }
  const port = buf.readUInt16BE(off); off += 2;
  return { cmd, host, port, size: off };
}

function buildSocksMethodsResponse(method) {
  const buf = Buffer.allocUnsafe(2);
  buf.writeUInt8(5, 0);
  buf.writeUInt8(method, 1);
  return buf;
}

function parseUdpPayload(buf) {
  if (buf.length < 2) return null;
  let off = 0;
  const hostLen = buf.readUInt16BE(off); off += 2;
  if (buf.length < 2 + hostLen + 2 + 2) return null;
  const host = buf.subarray(off, off + hostLen).toString('utf8'); off += hostLen;
  const port = buf.readUInt16BE(off); off += 2;
  const dataLen = buf.readUInt16BE(off); off += 2;
  if (buf.length < 2 + hostLen + 2 + 2 + dataLen) return null;
  const data = buf.subarray(off, off + dataLen);
  return { host, port, data };
}

function buildSocksUdpDatagram(host, port, data) {
  const addrInfo = encodeSocksAddr(host);
  const header = Buffer.allocUnsafe(2 + 1 + 1);
  header.writeUInt8(0x00, 0);
  header.writeUInt8(0x00, 1);
  header.writeUInt8(0x00, 2); // FRAG=0
  header.writeUInt8(addrInfo.atyp, 3);
  const portBuf = Buffer.allocUnsafe(2);
  portBuf.writeUInt16BE(port >>> 0, 0);
  return Buffer.concat([header, addrInfo.addr, portBuf, data]);
}

function parseSocksUdpRequest(buf) {
  // UDP request: RSV(2)=0x0000, FRAG(1), ATYP, DST.ADDR, DST.PORT, DATA
  if (buf.length < 4) return null;
  const rsv0 = buf.readUInt8(0), rsv1 = buf.readUInt8(1);
  if (rsv0 !== 0x00 || rsv1 !== 0x00) return { error: 'RSV' };
  const frag = buf.readUInt8(2);
  if (frag !== 0x00) return { error: 'FRAG' }; // we don't support fragmentation
  const atyp = buf.readUInt8(3);
  let off = 4;
  let host = null;
  if (atyp === 0x01) { // IPv4
    if (buf.length < off + 4 + 2) return null;
    host = Array.from(buf.subarray(off, off + 4)).join('.');
    off += 4;
  } else if (atyp === 0x03) { // domain
    if (buf.length < off + 1) return null;
    const len = buf.readUInt8(off); off += 1;
    if (buf.length < off + len + 2) return null;
    host = buf.subarray(off, off + len).toString('utf8'); off += len;
  } else if (atyp === 0x04) { // IPv6
    if (buf.length < off + 16 + 2) return null;
    const b = buf.subarray(off, off + 16);
    host = [...Array(8).keys()].map((i) => b.readUInt16BE(i * 2).toString(16)).join(':');
    off += 16;
  } else {
    return { error: 'ATYP' };
  }
  const port = buf.readUInt16BE(off); off += 2;
  const data = buf.subarray(off);
  return { host, port, data };
}

/* ============== Local SOCKS5 Server and Flow ============== */

function handleSocksConnection(localSocket) {
  localSocket.setNoDelay(true);
  localSocket.setKeepAlive(true, 30000);
  if (IDLE_TIMEOUT > 0) {
    localSocket.setTimeout(IDLE_TIMEOUT, () => {
      try { localSocket.destroy(); } catch (_) {}
    });
  }

  let stage = 'NEGOTIATION';
  let buf = Buffer.alloc(0);
  let tcpConnId = null; // for CONNECT
  let udpAssocId = null; // for UDP ASSOCIATE

  function onData(chunk) {
    buf = Buffer.concat([buf, chunk]);

    if (stage === 'NEGOTIATION') {
      if (buf.length < 2) return;
      const ver = buf.readUInt8(0);
      const nmethods = buf.readUInt8(1);
      if (ver !== 5) {
        localSocket.end();
        return;
      }
      if (buf.length < 2 + nmethods) return;
      const methods = Array.from(buf.subarray(2, 2 + nmethods));
      buf = buf.subarray(2 + nmethods);

      // We support only NO AUTH (0x00)
      const method = methods.includes(0x00) ? 0x00 : 0xFF;
      localSocket.write(buildSocksMethodsResponse(method));
      if (method === 0xFF) {
        localSocket.end();
        return;
      }
      stage = 'REQUEST';
    }

    while (stage === 'REQUEST') {
      const parsed = parseSocksRequest(buf);
      if (parsed === null) return; // need more
      if (parsed.error) {
        // invalid
        try { localSocket.write(buildSocksReply(0x01, '0.0.0.0', 0)); } catch (_) {}
        localSocket.end();
        return;
      }
      const { cmd, host, port, size } = parsed;
      buf = buf.subarray(size);

      if (cmd === 0x01 /* CONNECT */) {
        // Allocate TCP connection id and OPEN through tunnel
        tcpConnId = genConnId();
        const state = {
          id: tcpConnId,
          localSocket,
          opened: false,
          openTimer: null,
        };
        connections.set(tcpConnId, state);
        startOpenTcp(tcpConnId, host, port, state);
        // Wait for OPEN_RESULT to send reply, then go streaming
        stage = 'CONNECT_WAIT';

      } else if (cmd === 0x03 /* UDP ASSOCIATE */) {
        // Create UDP association
        udpAssocId = genConnId();
        const as = {
          id: udpAssocId,
          controlSocket: localSocket,
          udpSocket: null,
          openTimer: null,
          opened: false,
          clients: new Set(),
        };
        udpAssocs.set(udpAssocId, as);

        // Create local UDP socket for client to send datagrams to
        const udpSock = dgram.createSocket('udp4');
        as.udpSocket = udpSock;

        udpSock.on('message', (msg, rinfo) => {
          // Parse SOCKS5 UDP request
          const req = parseSocksUdpRequest(msg);
          if (!req || req.error) {
            // ignore invalid/unhandled fragmentation
            return;
          }
          // Track client endpoint so we can send replies back
          as.clients.add(`${rinfo.address}:${rinfo.port}`);
          // Send over tunnel as UDP_SEND
          if (!tunnelSocket || tunnelSocket.destroyed) return;
          writeMessage(tunnelSocket, MSG_TYPES.UDP_SEND, udpAssocId, buildUdpPayload(req.host, req.port, req.data));
        });

        udpSock.on('error', () => {
          // Close association
          try { udpSock.close(); } catch (_) {}
          writeMessage(tunnelSocket, MSG_TYPES.UDP_CLOSE, udpAssocId);
          udpAssocs.delete(udpAssocId);
          try { localSocket.write(buildSocksReply(0x01, '0.0.0.0', 0)); } catch (_) {}
          try { localSocket.destroy(); } catch (_) {}
        });

        udpSock.bind(0, options.bindHost, () => {
          // After binding, request UDP_OPEN on tunnel
          if (!tunnelSocket || tunnelSocket.destroyed) {
            try { localSocket.write(buildSocksReply(0x01, '0.0.0.0', 0)); } catch (_) {}
            try { localSocket.destroy(); } catch (_) {}
            return;
          }

          as.openTimer = setTimeout(() => {
            if (!as.opened) {
              try { localSocket.write(buildSocksReply(0x06, '0.0.0.0', 0)); } catch (_) {}
              writeMessage(tunnelSocket, MSG_TYPES.UDP_CLOSE, udpAssocId);
              closeUdpAssoc(udpAssocId);
              try { localSocket.destroy(); } catch (_) {}
            }
          }, OPEN_RESULT_TIMEOUT);

          writeMessage(tunnelSocket, MSG_TYPES.UDP_OPEN, udpAssocId);
        });

        if (UDP_IDLE_TIMEOUT > 0) {
          udpSock.on('listening', () => {
            udpSock.setRecvBufferSize?.(1 << 20);
          });
          let idleTimer = setTimeout(() => {
            try { udpSock.close(); } catch (_) {}
            writeMessage(tunnelSocket, MSG_TYPES.UDP_CLOSE, udpAssocId);
            udpAssocs.delete(udpAssocId);
            try { localSocket.destroy(); } catch (_) {}
          }, UDP_IDLE_TIMEOUT);
          // Reset idle timer on activity
          udpSock.on('message', () => {
            clearTimeout(idleTimer);
            idleTimer = setTimeout(() => {
              try { udpSock.close(); } catch (_) {}
              writeMessage(tunnelSocket, MSG_TYPES.UDP_CLOSE, udpAssocId);
              udpAssocs.delete(udpAssocId);
              try { localSocket.destroy(); } catch (_) {}
            }, UDP_IDLE_TIMEOUT);
          });
        }

        // Keep stage as REQUEST to allow multiple requests? SOCKS typically one per TCP session.
        // We keep control connection open until client closes.
        stage = 'ASSOCIATED';

      } else if (cmd === 0x02 /* BIND */) {
        // Not supported
        try { localSocket.write(buildSocksReply(0x07 /* Command not supported */, '0.0.0.0', 0)); } catch (_) {}
        localSocket.end();
        return;
      } else {
        // Unknown
        try { localSocket.write(buildSocksReply(0x07, '0.0.0.0', 0)); } catch (_) {}
        localSocket.end();
        return;
      }
    }
  }

  localSocket.on('data', onData);

  localSocket.on('close', () => {
    // For TCP CONNECT stream
    if (tcpConnId !== null) {
      writeMessage(tunnelSocket, MSG_TYPES.CLOSE, tcpConnId);
      connections.delete(tcpConnId);
    }
    // For UDP associate
    if (udpAssocId !== null) {
      writeMessage(tunnelSocket, MSG_TYPES.UDP_CLOSE, udpAssocId);
      closeUdpAssoc(udpAssocId);
    }
  });

  localSocket.on('error', () => {
    if (tcpConnId !== null) {
      writeMessage(tunnelSocket, MSG_TYPES.CLOSE, tcpConnId);
      connections.delete(tcpConnId);
    }
    if (udpAssocId !== null) {
      writeMessage(tunnelSocket, MSG_TYPES.UDP_CLOSE, udpAssocId);
      closeUdpAssoc(udpAssocId);
    }
  });
}

function closeUdpAssoc(id) {
  const as = udpAssocs.get(id);
  if (!as) return;
  try { as.udpSocket && as.udpSocket.close(); } catch (_) {}
  udpAssocs.delete(id);
}

/* ============== Server bootstrap ============== */

// Start local SOCKS5 server
const socksServer = net.createServer((socket) => {
  handleSocksConnection(socket);
});

socksServer.on('error', (err) => {
  console.error('Local SOCKS5 server error:', err.message);
  process.exit(1);
});

socksServer.listen(parseInt(options.socksPort, 10), options.bindHost, () => {
  console.log(`Local SOCKS5 proxy listening on ${options.bindHost}:${options.socksPort}`);
});

// Connect tunnel and maintain it
connectTunnel();

// Graceful shutdown
process.on('SIGINT', () => {
  console.log('Shutting down...');
  try { socksServer.close(); } catch (_) {}
  try { tunnelSocket && tunnelSocket.destroy(); } catch (_) {}
  for (const [id] of connections) {
    destroyLocalTCP(id);
  }
  for (const [id] of udpAssocs) {
    writeMessage(tunnelSocket, MSG_TYPES.UDP_CLOSE, id);
    closeUdpAssoc(id);
  }
  process.exit(0);
});