#!/usr/bin/env node

/**
 * PSK Proxy Exit-Node (Client)
 *
 * Connects to the relay server as a client.
 * Receives OPEN(host,port) to create outbound TCP connections to remote servers,
 * then forwards DATA/CLOSE frames bidirectionally.
 *
 * Also supports UDP relaying for SOCKS5 UDP ASSOCIATE via UDP_* frames.
 */

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

program
  .requiredOption('-H, --relay-host <host>', 'Relay server host to connect to')
  .requiredOption('-P, --relay-port <port>', 'Relay server port for exit connections')
  .requiredOption('--psk-file <path>', 'Path to PSK key file')
  .requiredOption('--identity <identity>', 'PSK identity to use when connecting to relay')
  .option('--connect-timeout <ms>', 'Timeout for outbound TCP connect (ms)', '10000')
  .option('--reconnect-delay <ms>', 'Delay before reconnecting to relay (ms)', '2000')
  .parse();

const options = program.opts();

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

const OUT_CONNECT_TIMEOUT = parseInt(options.connectTimeout, 10) || 10000;
const RECONNECT_DELAY = parseInt(options.reconnectDelay, 10) || 2000;

// 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 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]);

    // Parse as many complete frames as possible
    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 parseOpenPayload(buf) {
  if (buf.length < 4) return null;
  let offset = 0;
  const hostLen = buf.readUInt16BE(offset); offset += 2;
  if (buf.length < 2 + hostLen + 2) return null;
  const host = buf.subarray(offset, offset + hostLen).toString('utf8'); offset += hostLen;
  const port = buf.readUInt16BE(offset); offset += 2;
  return { host, port };
}

function buildOpenResultPayload(success) {
  const b = Buffer.allocUnsafe(1);
  b.writeUInt8(success ? 1 : 0, 0);
  return b;
}

function parseUdpPayload(buf) {
  // [2B hostLen][host][2B port][2B dataLen][data...]
  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 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;
}

// Global state
let relaySocket = null;
const upstreamConns = new Map(); // connectionId -> net.Socket

// UDP association mapping: connectionId -> { udp4?: dgram.Socket, udp6?: dgram.Socket }
const udpAssoc = new Map();

// Framing and fairness controls
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 for upstream->relay direction
const txQueues = new Map(); // connectionId -> { queue: Buffer[], sending: boolean }

function processTxQueue(connectionId, state) {
  if (!relaySocket || relaySocket.destroyed) {
    // Drop queued data if relay is gone
    state.queue = [];
    state.sending = false;
    return;
  }
  if (state.sending) return;
  state.sending = true;

  let bytesThisTick = 0;

  const run = () => {
    if (!relaySocket || relaySocket.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(relaySocket, MSG_TYPES.DATA, connectionId, slice);
        if (!ok) {
          // Backpressure: remember progress and resume on drain
          buf._offset = end;
          relaySocket.once('drain', () => {
            // Reset flag to allow re-entry
            state.sending = false;
            run();
          });
          return;
        }
        offset = end;
        bytesThisTick += slice.length;
        if (bytesThisTick >= BYTES_PER_TICK) {
          // Yield to allow other I/O
          buf._offset = offset;
          bytesThisTick = 0;
          setImmediate(() => {
            state.sending = false;
            run();
          });
          return;
        }
      }

      // Finished this buffer
      delete buf._offset;
      state.queue.shift();
    }

    state.sending = false;
  };

  run();
}

function enqueueToRelay(connectionId, data) {
  let state = txQueues.get(connectionId);
  if (!state) {
    state = { queue: [], sending: false };
    txQueues.set(connectionId, state);
  }
  state.queue.push(data);
  // Attempt to process immediately
  if (!state.sending) {
    state.sending = false; // ensure process can start
    processTxQueue(connectionId, state);
  }
}

function closeAllUpstreams() {
  for (const [, s] of upstreamConns) {
    try { s.destroy(); } catch (_) {}
  }
  upstreamConns.clear();
}

function closeAllUdp() {
  for (const [, obj] of udpAssoc) {
    try { obj.udp4 && obj.udp4.close(); } catch (_) {}
    try { obj.udp6 && obj.udp6.close(); } catch (_) {}
  }
  udpAssoc.clear();
}

function ensureUdpSocketsForAssoc(connectionId) {
  if (udpAssoc.has(connectionId)) return udpAssoc.get(connectionId);
  const onMessage = (msg, rinfo) => {
    // Forward incoming datagrams to client
    writeMessage(relaySocket, MSG_TYPES.UDP_RECV, connectionId, buildUdpPayload(rinfo.address, rinfo.port, msg));
  };
  // Create both IPv4 and IPv6 sockets to support all targets
  const u4 = dgram.createSocket('udp4');
  const u6 = dgram.createSocket('udp6');

  u4.on('message', onMessage);
  u6.on('message', onMessage);

  u4.on('error', (err) => { /* log but keep running */ console.warn(`UDP4 error (conn ${connectionId}): ${err.message}`); });
  u6.on('error', (err) => { /* log but keep running */ console.warn(`UDP6 error (conn ${connectionId}): ${err.message}`); });

  // Bind to ephemeral ports to receive replies
  try { u4.bind(0); } catch (_) {}
  try { u6.bind(0); } catch (_) {}

  const entry = { udp4: u4, udp6: u6 };
  udpAssoc.set(connectionId, entry);
  return entry;
}

function connectToRelay() {
  console.log(`Connecting to relay server ${options.relayHost}:${options.relayPort} via TLS-PSK...`);

  const pskCb = () => ({
    identity: options.identity,
    psk: Buffer.from(pskKey, 'hex'),
  });

  const sock = tls.connect(
    {
      host: options.relayHost,
      port: parseInt(options.relayPort, 10),
      pskCallback: pskCb,
      ciphers: 'PSK-AES256-GCM-SHA384:PSK-AES128-GCM-SHA256',
      checkServerIdentity: () => undefined,
    },
    () => {
      console.log('Connected to relay server');
    }
  );

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

  relaySocket = sock;
  const reader = createMessageReader();

  sock.on('data', (data) => {
    reader(data, (type, connectionId, payload) => {
      if (type === MSG_TYPES.OPEN) {
        const spec = parseOpenPayload(payload);
        if (!spec) {
          console.warn(`Invalid OPEN payload for connection ${connectionId}`);
          writeMessage(relaySocket, MSG_TYPES.OPEN_RESULT, connectionId, buildOpenResultPayload(false));
          return;
        }

        const { host, port } = spec;
        // Create outbound TCP connection
        const upstream = net.createConnection({ host, port });
        upstream.setNoDelay(true);
        upstream.setKeepAlive(true, 30000);

        let connected = false;
        const connectTimer = setTimeout(() => {
          if (!connected) {
            upstream.destroy(new Error('Connect timeout'));
          }
        }, OUT_CONNECT_TIMEOUT);

        upstream.once('connect', () => {
          connected = true;
          clearTimeout(connectTimer);
          upstreamConns.set(connectionId, upstream);
          // Notify open success
          writeMessage(relaySocket, MSG_TYPES.OPEN_RESULT, connectionId, buildOpenResultPayload(true));
        });

        upstream.on('data', (chunk) => {
          // Queue data with framing and backpressure-aware sending
          enqueueToRelay(connectionId, chunk);
        });

        const cleanup = () => {
          clearTimeout(connectTimer);
          if (upstreamConns.get(connectionId) === upstream) {
            upstreamConns.delete(connectionId);
          }
          // Drop any pending TX data for this connection
          txQueues.delete(connectionId);
          writeMessage(relaySocket, MSG_TYPES.CLOSE, connectionId);
        };

        upstream.on('error', (err) => {
          if (!connected) {
            clearTimeout(connectTimer);
            writeMessage(relaySocket, MSG_TYPES.OPEN_RESULT, connectionId, buildOpenResultPayload(false));
          } else {
            // Upstream error after established
            cleanup();
          }
        });

        upstream.on('close', () => {
          cleanup();
        });
      } else if (type === MSG_TYPES.DATA) {
        const upstream = upstreamConns.get(connectionId);
        if (upstream && !upstream.destroyed) {
          upstream.write(payload);
        }
      } else if (type === MSG_TYPES.CLOSE) {
        const upstream = upstreamConns.get(connectionId);
        if (upstream) {
          upstream.destroy();
          upstreamConns.delete(connectionId);
        }

      // UDP handling
      } else if (type === MSG_TYPES.UDP_OPEN) {
        try {
          ensureUdpSocketsForAssoc(connectionId);
          writeMessage(relaySocket, MSG_TYPES.UDP_OPEN_RESULT, connectionId, buildOpenResultPayload(true));
        } catch (e) {
          console.warn(`Failed to create UDP association for ${connectionId}: ${e.message}`);
          writeMessage(relaySocket, MSG_TYPES.UDP_OPEN_RESULT, connectionId, buildOpenResultPayload(false));
        }
      } else if (type === MSG_TYPES.UDP_SEND) {
        const parsed = parseUdpPayload(payload);
        if (!parsed) {
          console.warn(`Invalid UDP_SEND payload for ${connectionId}`);
          return;
        }
        const { host, port, data } = parsed;
        const entry = udpAssoc.get(connectionId) || ensureUdpSocketsForAssoc(connectionId);
        // Choose v6 socket if IPv6 literal detected
        const isV6 = host.includes(':');
        const sock = isV6 ? entry.udp6 : entry.udp4;
        try {
          sock.send(data, port, host);
        } catch (e) {
          console.warn(`UDP send failed (conn ${connectionId}) to ${host}:${port} - ${e.message}`);
        }
      } else if (type === MSG_TYPES.UDP_CLOSE) {
        const entry = udpAssoc.get(connectionId);
        if (entry) {
          try { entry.udp4 && entry.udp4.close(); } catch (_) {}
          try { entry.udp6 && entry.udp6.close(); } catch (_) {}
          udpAssoc.delete(connectionId);
        }
      } else {
        // ignore unknown types
      }
    });
  });

  sock.on('close', () => {
    console.log(`Disconnected from relay server. Retrying in ${RECONNECT_DELAY}ms...`);
    relaySocket = null;
    closeAllUpstreams();
    closeAllUdp();
    setTimeout(connectToRelay, RECONNECT_DELAY);
  });

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

// Start connection to the relay server
connectToRelay();

process.on('SIGINT', () => {
  console.log('Shutting down...');
  if (relaySocket) try { relaySocket.destroy(); } catch (_) {}
  closeAllUpstreams();
  closeAllUdp();
  process.exit(0);
});