exit node as client

proxy-server.js

@@ -3,9 +3,11 @@
 /**
  * PSK Proxy Relay-Node (Server)
  *
- * Listens for a TLS-PSK tunnel connection from the proxy client.
- * Establishes a single TLS-PSK tunnel connection to an exit node.
- * Relays frames between the client and the exit node.
+ * Exposes two ports:
+ * 1. Client port: Listens for TLS-PSK tunnel connections from proxy clients
+ * 2. Exit port: Listens for TLS-PSK tunnel connections from exit nodes
+ * 
+ * Relays frames between connected clients and exits.
  */
 
 const tls = require('tls');
@@ -13,12 +15,12 @@ const fs = require('fs');
 const { program } = require('commander');
 
 program
-  .requiredOption('-P, --tunnel-port <port>', 'Port for proxy client TLS-PSK tunnel connections')
+  .requiredOption('-C, --client-port <port>', 'Port for proxy client TLS-PSK tunnel connections')
+  .requiredOption('-E, --exit-port <port>', 'Port for exit node TLS-PSK tunnel connections')
   .requiredOption('-H, --host <host>', 'Host to bind to (e.g., 0.0.0.0)')
   .requiredOption('--psk-file <path>', 'Path to PSK key file for client and exit connections')
-  .requiredOption('--exit-host <host>', 'Exit node host')
-  .requiredOption('--exit-port <port>', 'Exit node port')
-  .requiredOption('--exit-identity <identity>', 'PSK identity for the exit node')
+  .requiredOption('--client-identity <identity>', 'Expected PSK identity for client connections')
+  .requiredOption('--exit-identity <identity>', 'Expected PSK identity for exit connections')
   .parse();
 
 const options = program.opts();
@@ -80,60 +82,31 @@ function writeMessage(socket, type, connectionId, data = Buffer.alloc(0)) {
   socket.write(buf);
 }
 
-function connectToExitNode() {
-  console.log(`Connecting to exit node ${options.exitHost}:${options.exitPort}...`);
-
-  const pskCb = () => ({
-    identity: options.exitIdentity,
-    psk: Buffer.from(pskKey, 'hex'),
-  });
-
-  const sock = tls.connect(
-    {
-      host: options.exitHost,
-      port: parseInt(options.exitPort, 10),
-      pskCallback: pskCb,
-      ciphers: 'PSK-AES256-GCM-SHA384:PSK-AES128-GCM-SHA256',
-      checkServerIdentity: () => undefined,
-    },
-    () => {
-      console.log('Connected to exit node');
-    }
-  );
-
-  sock.setNoDelay(true);
-  sock.setKeepAlive(true, 30000);
-
-  exitSocket = sock;
-  exitReader = createMessageReader((type, connId, data) => {
-    // Forward messages from exit to client
-    if (clientSocket && !clientSocket.destroyed) {
-      writeMessage(clientSocket, type, connId, data);
-    }
-  });
-
-  sock.on('data', (data) => exitReader(data));
-
-  sock.on('close', () => {
-    console.log('Disconnected from exit node. Retrying in 2s...');
-    exitSocket = null;
-    // If client is still here, it will be disconnected by the client server logic
-    if (clientSocket) {
-        try { clientSocket.destroy(); } catch (_) {}
-    }
-    setTimeout(connectToExitNode, 2000);
-  });
-
-  sock.on('error', (err) => {
-    console.error('Exit node connection error:', err.message);
-  });
-}
-
+// Client PSK callback
 const clientPskCallback = (socket, identity) => {
   console.log(`Client identity: ${identity}`);
+  
+  if (identity !== options.clientIdentity) {
+    console.warn(`Client PSK identity mismatch. Expected '${options.clientIdentity}', got '${identity}'.`);
+    return null;
+  }
+  
   return Buffer.from(pskKey, 'hex');
 };
 
+// Exit PSK callback
+const exitPskCallback = (socket, identity) => {
+  console.log(`Exit identity: ${identity}`);
+  
+  if (identity !== options.exitIdentity) {
+    console.warn(`Exit PSK identity mismatch. Expected '${options.exitIdentity}', got '${identity}'.`);
+    return null;
+  }
+  
+  return Buffer.from(pskKey, 'hex');
+};
+
+// Create server for client connections
 const clientServer = tls.createServer(
   {
     pskCallback: clientPskCallback,
@@ -173,16 +146,60 @@ const clientServer = tls.createServer(
   }
 );
 
-clientServer.listen(parseInt(options.tunnelPort, 10), options.host, () => {
-  console.log(`PSK Proxy Relay-Node listening for clients on ${options.host}:${options.tunnelPort}`);
+// Create server for exit connections
+const exitServer = tls.createServer(
+  {
+    pskCallback: exitPskCallback,
+    ciphers: 'PSK-AES256-GCM-SHA384:PSK-AES128-GCM-SHA256',
+  },
+  (socket) => {
+    if (exitSocket) {
+      console.log('Rejecting new exit connection, already have one.');
+      socket.destroy();
+      return;
+    }
+    console.log('Exit node connected');
+    exitSocket = socket;
+
+    socket.setNoDelay(true);
+    socket.setKeepAlive(true, 30000);
+
+    exitReader = createMessageReader((type, connId, data) => {
+      // Forward messages from exit to client
+      if (clientSocket && !clientSocket.destroyed) {
+        writeMessage(clientSocket, type, connId, data);
+      }
+    });
+
+    socket.on('data', (data) => exitReader(data));
+
+    socket.on('close', () => {
+      console.log('Exit node disconnected');
+      exitSocket = null;
+      // When exit disconnects, we can optionally close the client connection
+      // or keep it alive. For simplicity, we keep it.
+    });
+
+    socket.on('error', (err) => {
+      console.error('Exit socket error:', err.message);
+    });
+  }
+);
+
+// Start listening for client connections
+clientServer.listen(parseInt(options.clientPort, 10), options.host, () => {
+  console.log(`PSK Proxy Relay-Node listening for clients on ${options.host}:${options.clientPort}`);
 });
 
-// Start connection to the exit node
-connectToExitNode();
+// Start listening for exit connections
+exitServer.listen(parseInt(options.exitPort, 10), options.host, () => {
+  console.log(`PSK Proxy Relay-Node listening for exit nodes on ${options.host}:${options.exitPort}`);
+});
 
 process.on('SIGINT', () => {
   console.log('Shutting down...');
   try { clientServer.close(); } catch (_) {}
+  try { exitServer.close(); } catch (_) {}
   if (clientSocket) try { clientSocket.destroy(); } catch (_) {}
   if (exitSocket) try { exitSocket.destroy(); } catch (_) {}
   process.exit(0);