#!/usr/bin/env node

/**
 * PSK Proxy Out-Node (Server)
 *
 * Listens for a single TLS-PSK tunnel connection from the proxy client.
 * Receives OPEN(host,port) to create outbound TCP connections to remote servers,
 * then forwards DATA/CLOSE frames bidirectionally.
 *
 * Protocol (over a single TLS socket):
 *   Header: [1 byte type][4 bytes connection id][4 bytes data length][data...]
 *
 * Message Types:
 *   DATA (2):            Carry stream data
 *   CLOSE (3):           Close a stream
 *   OPEN (4):            Open stream to host:port, payload = [2B hostLen][host][2B port]
 *   OPEN_RESULT (5):     Result for OPEN, payload = [1B status] (1 = success, 0 = failure)
 */

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

program
  .requiredOption('--tunnel-port <port>', 'Port for proxy client TLS-PSK tunnel connections')
  .requiredOption('--host <host>', 'Host to bind to (e.g., 0.0.0.0)')
  .requiredOption('--psk-file <path>', 'Path to PSK key file')
  .option('--connect-timeout <ms>', 'Timeout for outbound TCP connect (ms)', '10000')
  .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;

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

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

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 pskCallback(socket, identity) {
  console.log(`Tunnel client identity: ${identity}`);
  return Buffer.from(pskKey, 'hex');
}

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

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

const server = tls.createServer({
  pskCallback,
  ciphers: 'PSK-AES256-GCM-SHA384:PSK-AES128-GCM-SHA256',
}, (socket) => {
  console.log('Proxy tunnel client connected');
  tunnelSocket = socket;

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

  const reader = createMessageReader();

  socket.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(tunnelSocket, 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(tunnelSocket, MSG_TYPES.OPEN_RESULT, connectionId, buildOpenResultPayload(true));
        });

        upstream.on('data', (chunk) => {
          writeMessage(tunnelSocket, MSG_TYPES.DATA, connectionId, chunk);
        });

        const cleanup = () => {
          clearTimeout(connectTimer);
          if (upstreamConns.get(connectionId) === upstream) {
            upstreamConns.delete(connectionId);
          }
          writeMessage(tunnelSocket, MSG_TYPES.CLOSE, connectionId);
        };

        upstream.on('error', (err) => {
          if (!connected) {
            clearTimeout(connectTimer);
            writeMessage(tunnelSocket, 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);
        }
      } else {
        // ignore unknown types
      }
    });
  });

  socket.on('close', () => {
    console.log('Proxy tunnel disconnected');
    tunnelSocket = null;
    closeAllUpstreams();
  });

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

server.listen(parseInt(options.tunnelPort, 10), options.host, () => {
  console.log(`PSK Proxy Out-Node listening on ${options.host}:${options.tunnelPort}`);
});

process.on('SIGINT', () => {
  console.log('Shutting down...');
  try { server.close(); } catch (_) {}
  closeAllUpstreams();
  process.exit(0);
});