diodejs 0.4.2 → 0.4.4

package/bindPort.js

@@ -8,6 +8,8 @@ const DiodeRPC = require('./rpc');
 const nativeCrypto = require('./nativeCrypto');
 const logger = require('./logger');
 
+const BIND_PORT_LISTENER_STATE = Symbol.for('diodejs.bindPort.listenerState');
+
 // Custom Duplex stream to handle the Diode connection
 class DiodeSocket extends Duplex {
   constructor(ref, rpc) {
@@ -129,6 +131,104 @@ class BindPort {
     return this.connection;
   }
 
+  _connectionKey(connection) {
+    if (!connection) return '';
+    if (connection._managerHostKey) return connection._managerHostKey;
+    if (connection.host && connection.port) return `${connection.host}:${connection.port}`;
+    return '';
+  }
+
+  _isUsableConnection(connection) {
+    return connection && connection.socket && !connection.socket.destroyed;
+  }
+
+  _clearDeviceRelayCache(deviceIdHex) {
+    const cache = this.connection && this.connection.deviceRelayCache;
+    if (!cache || typeof cache.delete !== 'function') {
+      return;
+    }
+    const normalized = this._stripHexPrefix(deviceIdHex || '').toLowerCase();
+    cache.delete(normalized);
+    cache.delete(`0x${normalized}`);
+  }
+
+  async _getApiRelayCandidates(deviceId, deviceIdHex) {
+    const candidates = [];
+    const seen = new Set();
+    const push = (connection) => {
+      if (!this._isUsableConnection(connection)) {
+        return;
+      }
+      const key = this._connectionKey(connection) || String(candidates.length);
+      if (seen.has(key)) {
+        return;
+      }
+      seen.add(key);
+      candidates.push(connection);
+    };
+
+    try {
+      push(await this._resolveConnectionForDevice(deviceId));
+    } catch (error) {
+      logger.warn(() => `Error resolving relay for device ${deviceIdHex}: ${error}`);
+    }
+
+    if (typeof (this.connection && this.connection.getNearestConnection) === 'function') {
+      push(this.connection.getNearestConnection());
+    }
+
+    if (typeof (this.connection && this.connection.getConnections) === 'function') {
+      for (const connection of this.connection.getConnections()) {
+        push(connection);
+      }
+    } else {
+      push(this.connection);
+    }
+
+    return candidates;
+  }
+
+  async _openApiPortWithRelayFallback(deviceId, deviceIdHex, formattedTargetPort, flags = 'rw') {
+    let candidates = await this._getApiRelayCandidates(deviceId, deviceIdHex);
+    let clearedCache = false;
+    let lastError = null;
+
+    for (let index = 0; index < candidates.length; index += 1) {
+      const connection = candidates[index];
+      const rpc = this._getRpcFor(connection);
+      const relayKey = this._connectionKey(connection) || 'unknown relay';
+
+      try {
+        const ref = await rpc.portOpen(deviceId, formattedTargetPort, flags);
+        if (ref) {
+          if (index > 0) {
+            logger.info(() => `Port ${formattedTargetPort} opened via fallback relay ${relayKey}`);
+          }
+          return { connection, rpc, ref };
+        }
+        lastError = new Error(`portopen returned no ref via ${relayKey}`);
+      } catch (error) {
+        lastError = error;
+      }
+
+      logger.warn(() => `Port ${formattedTargetPort} did not open via ${relayKey}: ${lastError}`);
+      if (!clearedCache) {
+        clearedCache = true;
+        this._clearDeviceRelayCache(deviceIdHex);
+        const refreshed = await this._getApiRelayCandidates(deviceId, deviceIdHex);
+        for (const candidate of refreshed) {
+          const key = this._connectionKey(candidate) || String(candidates.length);
+          const exists = candidates.some((existing) => (this._connectionKey(existing) || '') === key);
+          if (!exists) {
+            candidates.push(candidate);
+          }
+        }
+      }
+    }
+
+    throw lastError || new Error('No relay connection available');
+  }
+
   async _openTlsHandshakeChannel(connection, rpc, ref) {
     const diodeSocket = new DiodeSocket(ref, rpc);
     const certPem = connection.getDeviceCertificate();
@@ -230,8 +330,26 @@ class BindPort {
   }
   
   _setupMessageListener() {
-    // Listen for data events from the device
-    this.connection.on('unsolicited', (message, sourceConnection) => {
+    const rootConnection = this.connection;
+    if (!rootConnection || typeof rootConnection.on !== 'function') {
+      return;
+    }
+
+    let state = rootConnection[BIND_PORT_LISTENER_STATE];
+    if (state) {
+      state.instances.add(this);
+      return;
+    }
+
+    state = {
+      instances: new Set([this]),
+    };
+    rootConnection[BIND_PORT_LISTENER_STATE] = state;
+
+    // Listen for data events from the device. This is shared per connection
+    // manager so multiple BindPort instances do not duplicate payloads into
+    // the same client socket.
+    state.onUnsolicited = (message, sourceConnection) => {
       const connection = sourceConnection || this.connection;
       if (!connection || typeof connection.getClientSocket !== 'function') {
         logger.warn(() => 'Received unsolicited message without a valid connection context');
@@ -276,7 +394,11 @@ class BindPort {
           if (clientSocket.diodeSocket) {
             clientSocket.diodeSocket._destroy(null, () => {});
           }
-          clientSocket.end();
+          if (typeof clientSocket.end === 'function') {
+            clientSocket.end();
+          } else if (typeof clientSocket.close === 'function') {
+            clientSocket.close();
+          }
           connection.deleteClientSocket(dataRef);
           logger.info(() => `Port closed for ref: ${dataRef.toString('hex')}`);
         }
@@ -285,19 +407,26 @@ class BindPort {
           logger.warn(() => `Unknown unsolicited message type: ${messageType}`);
         }
       }
-    });
+    };
+    rootConnection.on('unsolicited', state.onUnsolicited);
     
     // Handle device disconnect
-    this.connection.on('end', () => {
+    state.onEnd = () => {
       logger.info(() => 'Disconnected from Diode.io server');
-      this.closeAllServers();
-    });
+      for (const instance of state.instances) {
+        instance.closeAllServers();
+      }
+    };
+    rootConnection.on('end', state.onEnd);
 
     // Handle connection errors
-    this.connection.on('error', (err) => {
+    state.onError = (err) => {
       logger.error(() => `Connection error: ${err}`);
-      this.closeAllServers();
-    });
+      for (const instance of state.instances) {
+        instance.closeAllServers();
+      }
+    };
+    rootConnection.on('error', state.onError);
   }
 
   addPort(localPort, targetPort, deviceIdHex, protocol = 'tls', transport = undefined) {
@@ -471,39 +600,23 @@ class BindPort {
         let entry = server.clientRefs && server.clientRefs[clientKey];
         
         if (!entry) {
-          let connection;
           try {
-            connection = await this._resolveConnectionForDevice(deviceId);
-          } catch (error) {
-            logger.error(() => `Error resolving relay for device ${deviceIdHex}: ${error}`);
-            return;
-          }
-          if (!connection) {
-            logger.error(() => `No relay connection available for device ${deviceIdHex}`);
-            return;
-          }
-          const rpc = this._getRpcFor(connection);
-          try {
-            const ref = await rpc.portOpen(deviceId, formattedTargetPort, 'rw');
-            if (!ref) {
-              logger.error(() => `Error opening port ${formattedTargetPort} on deviceId: ${deviceIdHex}`);
-              return;
-            } else {
-              logger.info(() => `Port ${formattedTargetPort} opened on device with ref: ${ref.toString('hex')} for udp client ${clientKey}`);
-              if (!server.clientRefs) server.clientRefs = {};
-              entry = { ref, connection };
-              server.clientRefs[clientKey] = entry;
-              
-              // Store the client info
-              connection.addClientSocket(ref, {
-                address: rinfo.address,
-                port: rinfo.port,
-                protocol: 'udp',
-                write: (data) => {
-                  server.send(data, rinfo.port, rinfo.address);
-                }
-              });
-            }
+            const opened = await this._openApiPortWithRelayFallback(deviceId, deviceIdHex, formattedTargetPort, 'rw');
+            const { ref, connection } = opened;
+            logger.info(() => `Port ${formattedTargetPort} opened on device with ref: ${ref.toString('hex')} for udp client ${clientKey}`);
+            if (!server.clientRefs) server.clientRefs = {};
+            entry = { ref, connection };
+            server.clientRefs[clientKey] = entry;
+            
+            // Store the client info
+            connection.addClientSocket(ref, {
+              address: rinfo.address,
+              port: rinfo.port,
+              protocol: 'udp',
+              write: (data) => {
+                server.send(data, rinfo.port, rinfo.address);
+              }
+            });
           } catch (error) {
             logger.error(() => `Error opening port ${formattedTargetPort} on device: ${error}`);
             return;
@@ -539,24 +652,69 @@ class BindPort {
       const server = net.createServer(async (clientSocket) => {
         logger.info(() => `Client connected to local server on port ${localPort}`);
         clientSocket.setNoDelay(true);
-
-        let connection;
-        try {
-          connection = await this._resolveConnectionForDevice(deviceId);
-        } catch (error) {
-          logger.error(() => `Error resolving relay for device ${deviceIdHex}: ${error}`);
-          clientSocket.destroy();
-          return;
-        }
-        if (!connection) {
-          logger.error(() => `No relay connection available for device ${deviceIdHex}`);
-          clientSocket.destroy();
-          return;
+        if (useNative) {
+          clientSocket.pause();
         }
-        const rpc = this._getRpcFor(connection);
+
+        let connection = null;
+        let rpc = null;
 
         let ref;
+        let clientClosed = false;
+        let remoteCleanupStarted = false;
+        let tlsSocketWrapper = null;
+
+        const closeRemoteRef = async () => {
+          if (remoteCleanupStarted) {
+            return;
+          }
+
+          if (!ref || !connection || !rpc) {
+            return;
+          }
+          remoteCleanupStarted = true;
+
+          const storedSocket = typeof connection.getClientSocket === 'function'
+            ? connection.getClientSocket(ref)
+            : null;
+          if (storedSocket && storedSocket !== clientSocket && typeof storedSocket.end === 'function') {
+            try { storedSocket.end(); } catch {}
+          }
+          if (typeof connection.deleteClientSocket === 'function') {
+            try { connection.deleteClientSocket(ref); } catch {}
+          }
+
+          try {
+            await rpc.portClose(ref);
+            logger.info(() => `Port closed on device for ref: ${ref.toString('hex')}`);
+          } catch (error) {
+            logger.error(() => `Error closing port on device: ${error}`);
+          }
+        };
+
+        clientSocket.once('close', () => {
+          clientClosed = true;
+          closeRemoteRef();
+        });
+        clientSocket.once('error', (err) => {
+          logger.error(() => `Client socket error: ${err}`);
+        });
+
         if (useNative) {
+          try {
+            connection = await this._resolveConnectionForDevice(deviceId);
+          } catch (error) {
+            logger.error(() => `Error resolving relay for device ${deviceIdHex}: ${error}`);
+            clientSocket.destroy();
+            return;
+          }
+          if (!connection) {
+            logger.error(() => `No relay connection available for device ${deviceIdHex}`);
+            clientSocket.destroy();
+            return;
+          }
+          rpc = this._getRpcFor(connection);
+
           // Open a new native relay port on the device for this client
           let physicalPort;
           try {
@@ -647,26 +805,30 @@ class BindPort {
           });
           clientSocket.on('end', cleanup);
           relaySocket.on('end', cleanup);
+          clientSocket.resume();
 
           return;
         }
 
         // Legacy API relay
         try {
-          ref = await rpc.portOpen(deviceId, formattedTargetPort, 'rw');
-          if (!ref) {
-            logger.error(() => `Error opening port ${formattedTargetPort} on deviceId: ${deviceIdHex}`);
-            clientSocket.destroy();
-            return;
-          } else {
-            logger.info(() => `Port ${formattedTargetPort} opened on device with ref: ${ref.toString('hex')} for client`);
-          }
+          const opened = await this._openApiPortWithRelayFallback(deviceId, deviceIdHex, formattedTargetPort, 'rw');
+          connection = opened.connection;
+          rpc = opened.rpc;
+          ref = opened.ref;
+          logger.info(() => `Port ${formattedTargetPort} opened on device with ref: ${ref.toString('hex')} for client`);
         } catch (error) {
           logger.error(() => `Error opening port ${formattedTargetPort} on device: ${error}`);
           clientSocket.destroy();
           return;
         }
 
+        if (clientClosed || clientSocket.destroyed) {
+          logger.warn(() => `Local client disconnected before port ${formattedTargetPort} opened; closing ref ${ref.toString('hex')}`);
+          await closeRemoteRef();
+          return;
+        }
+
         if (protocol === 'tls') {
           // For tls protocol, create a proper tls connection
           try {
@@ -713,16 +875,27 @@ class BindPort {
               diodeSocket,
               tlsSocket,
               end: () => {
-                tlsSocket.end();
-                diodeSocket._destroy(null, () => {});
+                try { tlsSocket.end(); } catch {}
+                try { diodeSocket._destroy(null, () => {}); } catch {}
               }
             };
+            tlsSocketWrapper = socketWrapper;
             
             // Store the socket wrapper
             connection.addClientSocket(ref, socketWrapper);
+            tlsSocket.once('close', () => {
+              if (!clientSocket.destroyed) {
+                clientSocket.end();
+              }
+              closeRemoteRef();
+            });
             
           } catch (error) {
             logger.error(() => `Error setting up tls connection: ${error}`);
+            if (tlsSocketWrapper && typeof tlsSocketWrapper.end === 'function') {
+              tlsSocketWrapper.end();
+            }
+            await closeRemoteRef();
             clientSocket.destroy();
             return;
           }
@@ -742,24 +915,9 @@ class BindPort {
         }
 
         // Handle client socket closure (common for all protocols)
-        clientSocket.on('end', async () => {
+        clientSocket.once('end', () => {
           logger.info(() => 'Client disconnected');
-          if (ref && connection.hasClientSocket(ref)) {
-            try {
-              await rpc.portClose(ref);
-              logger.info(() => `Port closed on device for ref: ${ref.toString('hex')}`);
-              connection.deleteClientSocket(ref);
-            } catch (error) {
-              logger.error(() => `Error closing port on device: ${error}`);
-            }
-          } else {
-            logger.warn(() => 'Ref is invalid or no longer in clientSockets.');
-          }
-        });
-
-        // Handle client socket errors
-        clientSocket.on('error', (err) => {
-          logger.error(() => `Client socket error: ${err}`);
+          closeRemoteRef();
         });
       });
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "diodejs",
-  "version": "0.4.2",
+  "version": "0.4.4",
   "description": "A JavaScript client for interacting with the Diode network. It provides functionalities to bind and publish ports, send RPC commands, and handle responses.",
   "main": "index.js",
   "scripts": {

package/publishPort.js

@@ -399,6 +399,7 @@ class PublishPort extends EventEmitter {
         if (session.localSocket && typeof session.localSocket.resume === 'function') {
           session.localSocket.resume();
         }
+        this._flushNativeTCPRelayPending(session);
 
         if (session.protocol === 'udp' && session.relaySocket && session.session) {
           const probe = nativeCrypto.createUdpPacket(session.session, Buffer.alloc(0));
@@ -502,6 +503,7 @@ class PublishPort extends EventEmitter {
       relaySocket: null,
       localSocket: null,
       timer: null,
+      pendingRelayChunks: [],
     };
     session.timer = setTimeout(() => {
       if (!session.ready) {
@@ -546,6 +548,35 @@ class PublishPort extends EventEmitter {
     this.nativeSessions.delete(session.physicalPort);
   }
 
+  _consumeNativeTCPRelayData(session, data) {
+    if (!session || !session.session || !session.localSocket) {
+      return true;
+    }
+    try {
+      const messages = nativeCrypto.consumeTcpFrames(session.session, data);
+      for (const msg of messages) {
+        session.localSocket.write(msg);
+      }
+      return true;
+    } catch (error) {
+      logger.error(() => `TCP decrypt error (${session.deviceId}): ${error}`);
+      return false;
+    }
+  }
+
+  _flushNativeTCPRelayPending(session) {
+    if (!session || session.protocol !== 'tcp' || !Array.isArray(session.pendingRelayChunks)) {
+      return;
+    }
+    const pending = session.pendingRelayChunks.splice(0, session.pendingRelayChunks.length);
+    for (const chunk of pending) {
+      if (!this._consumeNativeTCPRelayData(session, chunk)) {
+        this._cleanupNativeSession(session);
+        return;
+      }
+    }
+  }
+
   handleNativeTCPRelay(sessionId, physicalPortRef, session, connection) {
     const rpc = this._getRpcFor(connection);
     const { physicalPort, port, host, deviceId } = session;
@@ -611,14 +642,11 @@ class PublishPort extends EventEmitter {
     localSocket.on('end', cleanup);
 
     relaySocket.on('data', (data) => {
-      if (!session.ready || !session.session) return;
-      try {
-        const messages = nativeCrypto.consumeTcpFrames(session.session, data);
-        for (const msg of messages) {
-          localSocket.write(msg);
-        }
-      } catch (error) {
-        logger.error(() => `TCP decrypt error (${deviceId}): ${error}`);
+      if (!session.ready || !session.session) {
+        session.pendingRelayChunks.push(Buffer.from(data));
+        return;
+      }
+      if (!this._consumeNativeTCPRelayData(session, data)) {
         cleanup();
       }
     });

package/test/bindPort.test.js

@@ -0,0 +1,242 @@
+const test = require('node:test');
+const assert = require('node:assert/strict');
+const EventEmitter = require('events');
+const net = require('net');
+const { once } = require('events');
+
+const BindPort = require('../bindPort');
+
+function makeRef(hex) {
+  return Buffer.from(hex.padStart(8, '0'), 'hex');
+}
+
+function makeRelay(hostKey, result, calls) {
+  const clientSockets = new Map();
+  const portCloseCalls = [];
+  return {
+    _managerHostKey: hostKey,
+    socket: { destroyed: false },
+    clientSockets,
+    portCloseCalls,
+    RPC: {
+      portOpen: async (_deviceId, port, flags) => {
+        calls.push({ hostKey, port, flags });
+        if (result instanceof Error) {
+          throw result;
+        }
+        return typeof result === 'function' ? result() : result;
+      },
+      portClose: async (ref) => {
+        portCloseCalls.push(ref);
+      },
+    },
+    addClientSocket(ref, socket) {
+      clientSockets.set(ref.toString('hex'), socket);
+    },
+    getClientSocket(ref) {
+      return clientSockets.get(ref.toString('hex'));
+    },
+    deleteClientSocket(ref) {
+      return clientSockets.delete(ref.toString('hex'));
+    },
+    hasClientSocket(ref) {
+      return clientSockets.has(ref.toString('hex'));
+    },
+  };
+}
+
+function deferred() {
+  let resolve;
+  let reject;
+  const promise = new Promise((res, rej) => {
+    resolve = res;
+    reject = rej;
+  });
+  return { promise, resolve, reject };
+}
+
+async function waitFor(predicate, timeoutMs = 1000) {
+  const startedAt = Date.now();
+  while (Date.now() - startedAt < timeoutMs) {
+    if (predicate()) {
+      return;
+    }
+    await new Promise((resolve) => setTimeout(resolve, 10));
+  }
+  assert.equal(predicate(), true);
+}
+
+class FakeManager extends EventEmitter {
+  constructor({ relays, resolvedRelay, nearestRelay, resolveError = null }) {
+    super();
+    this.relays = relays;
+    this.resolvedRelay = resolvedRelay;
+    this.nearestRelay = nearestRelay;
+    this.resolveError = resolveError;
+    this.deviceRelayCache = new Map();
+  }
+
+  async getConnectionForDevice() {
+    if (this.resolveError) {
+      throw this.resolveError;
+    }
+    return this.resolvedRelay;
+  }
+
+  getNearestConnection() {
+    return this.nearestRelay;
+  }
+
+  getConnections() {
+    return this.relays.slice();
+  }
+}
+
+test('API portopen falls back to another relay and clears stale device cache', async () => {
+  const calls = [];
+  const bad = makeRelay('bad.relay:41046', undefined, calls);
+  const goodRef = makeRef('01020304');
+  const good = makeRelay('good.relay:41046', goodRef, calls);
+  const manager = new FakeManager({
+    relays: [bad, good],
+    resolvedRelay: bad,
+    nearestRelay: bad,
+  });
+  manager.deviceRelayCache.set('8a72468957504d50247a260deb0218d504dd091b', {
+    hostKey: 'bad.relay:41046',
+  });
+
+  const bind = new BindPort(manager, {});
+  const result = await bind._openApiPortWithRelayFallback(
+    Buffer.from('8a72468957504d50247a260deb0218d504dd091b', 'hex'),
+    '8a72468957504d50247a260deb0218d504dd091b',
+    'tls:8088',
+    'rw'
+  );
+
+  assert.equal(result.connection, good);
+  assert.equal(result.ref, goodRef);
+  assert.deepEqual(calls.map((entry) => entry.hostKey), [
+    'bad.relay:41046',
+    'good.relay:41046',
+  ]);
+  assert.equal(manager.deviceRelayCache.has('8a72468957504d50247a260deb0218d504dd091b'), false);
+});
+
+test('API portopen tries connected relays when device relay lookup fails', async () => {
+  const calls = [];
+  const ref = makeRef('05060708');
+  const fallback = makeRelay('fallback.relay:41046', ref, calls);
+  const manager = new FakeManager({
+    relays: [fallback],
+    resolvedRelay: null,
+    nearestRelay: null,
+    resolveError: new Error('ticket lookup failed'),
+  });
+
+  const bind = new BindPort(manager, {});
+  const result = await bind._openApiPortWithRelayFallback(
+    Buffer.from('8a72468957504d50247a260deb0218d504dd091b', 'hex'),
+    '8a72468957504d50247a260deb0218d504dd091b',
+    'tls:8088',
+    'rw'
+  );
+
+  assert.equal(result.connection, fallback);
+  assert.equal(result.ref, ref);
+  assert.deepEqual(calls, [{
+    hostKey: 'fallback.relay:41046',
+    port: 'tls:8088',
+    flags: 'rw',
+  }]);
+});
+
+test('multiple BindPort instances share one manager listener set', () => {
+  const manager = new FakeManager({
+    relays: [],
+    resolvedRelay: null,
+    nearestRelay: null,
+  });
+
+  new BindPort(manager, {});
+  new BindPort(manager, {});
+
+  assert.equal(manager.listenerCount('unsolicited'), 1);
+  assert.equal(manager.listenerCount('end'), 1);
+  assert.equal(manager.listenerCount('error'), 1);
+});
+
+test('API bind closes ref when local TCP client disconnects before portopen returns', async () => {
+  const calls = [];
+  const opened = deferred();
+  const ref = makeRef('0a0b0c0d');
+  const relay = makeRelay('relay.example:41046', () => opened.promise, calls);
+  const manager = new FakeManager({
+    relays: [relay],
+    resolvedRelay: relay,
+    nearestRelay: relay,
+  });
+
+  const bind = new BindPort(manager, {
+    0: {
+      targetPort: 8088,
+      deviceIdHex: '8a72468957504d50247a260deb0218d504dd091b',
+      protocol: 'tcp',
+    }
+  });
+  bind.addPort(0, 8088, '8a72468957504d50247a260deb0218d504dd091b', 'tcp');
+  const server = bind.servers.get(0);
+
+  try {
+    await once(server, 'listening');
+    const client = net.connect(server.address().port, '127.0.0.1');
+    await once(client, 'connect');
+    client.destroy();
+    await once(client, 'close');
+
+    opened.resolve(ref);
+    await waitFor(() => relay.portCloseCalls.length === 1);
+
+    assert.equal(relay.portCloseCalls[0].toString('hex'), ref.toString('hex'));
+    assert.equal(relay.hasClientSocket(ref), false);
+  } finally {
+    server.close();
+  }
+});
+
+test('API bind closes ref when local TCP client disconnects after portopen returns', async () => {
+  const calls = [];
+  const ref = makeRef('0e0f1011');
+  const relay = makeRelay('relay.example:41046', ref, calls);
+  const manager = new FakeManager({
+    relays: [relay],
+    resolvedRelay: relay,
+    nearestRelay: relay,
+  });
+
+  const bind = new BindPort(manager, {
+    0: {
+      targetPort: 8088,
+      deviceIdHex: '8a72468957504d50247a260deb0218d504dd091b',
+      protocol: 'tcp',
+    }
+  });
+  bind.addPort(0, 8088, '8a72468957504d50247a260deb0218d504dd091b', 'tcp');
+  const server = bind.servers.get(0);
+
+  try {
+    await once(server, 'listening');
+    const client = net.connect(server.address().port, '127.0.0.1');
+    await once(client, 'connect');
+    await waitFor(() => relay.hasClientSocket(ref));
+
+    client.end();
+    await once(client, 'close');
+    await waitFor(() => relay.portCloseCalls.length === 1);
+
+    assert.equal(relay.portCloseCalls[0].toString('hex'), ref.toString('hex'));
+    assert.equal(relay.hasClientSocket(ref), false);
+  } finally {
+    server.close();
+  }
+});

package/test/publishPort.test.js

@@ -6,6 +6,7 @@ const tls = require('tls');
 const dgram = require('dgram');
 
 const PublishPort = require('../publishPort');
+const nativeCrypto = require('../nativeCrypto');
 
 class FakeStreamSocket extends EventEmitter {
   constructor() {
@@ -13,11 +14,14 @@ class FakeStreamSocket extends EventEmitter {
     this.destroyed = false;
     this.remoteAddress = undefined;
     this.remotePort = undefined;
+    this.writes = [];
   }
 
   setNoDelay() {}
   pause() {}
-  write() {}
+  write(data) {
+    this.writes.push(data);
+  }
   end() {}
   destroy() {
     this.destroyed = true;
@@ -307,6 +311,63 @@ test('native TCP publish connects local socket to configured host', () => {
   assert.deepEqual(connectCalls[1], { port: 8089, host: '10.0.0.8' });
 });
 
+test('native TCP publish buffers relay data until handshake is ready', () => {
+  const connection = new FakeConnection();
+  const publishPort = new PublishPort(connection, {
+    8089: { mode: 'public', host: '10.0.0.8' },
+  });
+  const originalConnect = net.connect;
+  const originalConsumeTcpFrames = nativeCrypto.consumeTcpFrames;
+  const sockets = [];
+
+  net.connect = (options, callback) => {
+    const socket = new FakeStreamSocket();
+    socket.options = options;
+    sockets.push(socket);
+    process.nextTick(() => {
+      if (typeof callback === 'function') {
+        callback();
+      }
+      socket.emit('connect');
+    });
+    return socket;
+  };
+  nativeCrypto.consumeTcpFrames = () => [Buffer.from('plain-http')];
+
+  try {
+    const session = {
+      physicalPort: 41000,
+      port: 8089,
+      host: '10.0.0.8',
+      protocol: 'tcp',
+      deviceId: '0x' + '11'.repeat(20),
+      ready: false,
+      session: null,
+      pendingRelayChunks: [],
+    };
+    publishPort.handleNativeTCPRelay(makeSessionId('07'), 41000, session, connection);
+
+    const relaySocket = sockets[0];
+    const localSocket = sockets[1];
+    relaySocket.emit('data', Buffer.from('early-encrypted'));
+
+    assert.equal(session.pendingRelayChunks.length, 1);
+    assert.equal(localSocket.writes.length, 0);
+
+    session.ready = true;
+    session.session = {};
+    publishPort._flushNativeTCPRelayPending(session);
+
+    assert.equal(session.pendingRelayChunks.length, 0);
+    assert.equal(localSocket.writes.length, 1);
+    assert.equal(localSocket.writes[0].toString(), 'plain-http');
+  } finally {
+    nativeCrypto.consumeTcpFrames = originalConsumeTcpFrames;
+    net.connect = originalConnect;
+    publishPort.stopListening();
+  }
+});
+
 test('native UDP publish connects local socket to configured host', () => {
   const connection = new FakeConnection();
   const publishPort = new PublishPort(connection, {