Improve TinySocket connection status handling (#117)

* Create specific TinySocket tests and move wake on lan tests from LGTV * Tentatively handles ECONNRESET issues on macOS by ending the server socket instead of just closing the server * Assert actions that require the client to be connected or disconnected * Make disconnecting a disconnected socket a no-op
WesSouza · Dec 30, 2023 · 06d2fa7 · 06d2fa7
1 parent a9986be
commit 06d2fa7
Show file tree

Hide file tree

Showing 4 changed files with 260 additions and 80 deletions.
diff --git a/packages/lgtv-ip-control/src/classes/LGTV.ts b/packages/lgtv-ip-control/src/classes/LGTV.ts
@@ -42,6 +42,10 @@ export class LGTV {
     return this.encoder.decode(response);
   }
 
+  get connected() {
+    return this.socket.connected;
+  }
+
   async connect(): Promise<void> {
     await this.socket.connect();
   }

diff --git a/packages/lgtv-ip-control/src/classes/TinySocket.ts b/packages/lgtv-ip-control/src/classes/TinySocket.ts
@@ -51,7 +51,8 @@ function assertSettings(settings: SocketSettings) {
 export class TimeoutError extends Error {}
 
 export class TinySocket {
-  private client = new Socket();
+  #client = new Socket();
+  #connected = false;
 
   constructor(
     private host: string,
@@ -65,6 +66,24 @@ export class TinySocket {
     );
   }
 
+  #isConnected() {
+    return (
+      this.#connected && !this.#client.connecting && !this.#client.destroyed
+    );
+  }
+
+  #assertConnected() {
+    assert(this.#isConnected(), 'should be connected');
+  }
+
+  #assertDisconnected() {
+    assert(!this.#isConnected(), 'should not be connected');
+  }
+
+  get connected() {
+    return this.#isConnected();
+  }
+
   wrap<T>(
     method: (
       resolve: (value: T) => void,
@@ -73,16 +92,17 @@ export class TinySocket {
   ): Promise<T> {
     return new Promise((resolve, reject) => {
       const handleTimeout = () => {
-        this.client.end();
+        this.#connected = false;
+        this.#client.end();
         reject(new TimeoutError());
       };
       const cleanup = () => {
-        this.client.removeListener('error', reject);
-        this.client.removeListener('timeout', handleTimeout);
+        this.#client.removeListener('error', reject);
+        this.#client.removeListener('timeout', handleTimeout);
       };
 
-      this.client.once('error', reject);
-      this.client.once('timeout', handleTimeout);
+      this.#client.once('error', reject);
+      this.#client.once('timeout', handleTimeout);
 
       method(
         (value: T) => {
@@ -97,22 +117,35 @@ export class TinySocket {
     });
   }
 
-  connect(): Promise<void> {
-    return this.wrap((resolve) => {
-      this.client.setTimeout(this.settings.networkTimeout);
-      this.client.connect(this.settings.networkPort, this.host, resolve);
+  async connect(): Promise<void> {
+    this.#assertDisconnected();
+
+    await this.wrap<undefined>((resolve) => {
+      this.#client.setTimeout(this.settings.networkTimeout);
+      try {
+        this.#client.connect(this.settings.networkPort, this.host, () =>
+          resolve(undefined),
+        );
+        this.#connected = true;
+      } catch (e) {
+        console.error(e);
+      }
     });
   }
 
   read(): Promise<Buffer> {
+    this.#assertConnected();
+
     return this.wrap((resolve) => {
-      this.client.once('data', resolve);
+      this.#client.once('data', resolve);
     });
   }
 
   write(data: Buffer): Promise<void> {
+    this.#assertConnected();
+
     return this.wrap((resolve, reject) => {
-      this.client.write(data, (error?: Error) => {
+      this.#client.write(data, (error?: Error) => {
         if (error) {
           reject(error);
         } else {
@@ -123,13 +156,20 @@ export class TinySocket {
   }
 
   async sendReceive(data: Buffer): Promise<Buffer> {
+    this.#assertConnected();
+
     await this.write(data);
     return this.read();
   }
 
   disconnect(): Promise<void> {
+    if (!this.#isConnected()) {
+      return Promise.resolve(undefined);
+    }
+
     return this.wrap((resolve) => {
-      this.client.end(resolve);
+      this.#connected = false;
+      this.#client.end(resolve);
     });
   }
 

diff --git a/packages/lgtv-ip-control/test/LGTV.test.ts b/packages/lgtv-ip-control/test/LGTV.test.ts
@@ -1,6 +1,4 @@
-import { createSocket, Socket as DgramSocket, SocketType } from 'dgram';
-import { AddressInfo, Server } from 'net';
-import { promisify } from 'util';
+import { AddressInfo, Server, Socket } from 'net';
 import { afterEach, beforeEach, describe, expect, it } from 'vitest';
 
 import { LGEncoder, LGEncryption } from '../src/classes/LGEncryption.js';
@@ -45,6 +43,7 @@ describe.each([
   ({ address, crypt }) => {
     let mockEncode: LGEncoder;
     let mockServer: Server;
+    let mockServerSocket: Socket;
     let testSettings: typeof DefaultSettings;
     let testTV: LGTV;
 
@@ -74,19 +73,25 @@ describe.each([
           responseTerminator: '\r',
         });
       }
-      mockServer = new Server().listen();
+      mockServer = new Server((socket) => {
+        mockServerSocket = socket;
+      }).listen();
       const port = (<AddressInfo>mockServer.address()).port;
       testSettings = { ...DefaultSettings, networkPort: port };
       testTV = new LGTV(address, MAC, crypt ? CRYPT_KEY : null, testSettings);
     });
 
     afterEach(async () => {
-      await testTV.disconnect();
-      // Graceful shutdown of net.Server doesn't work correctly on macOS.
-      // Work around it by making sure the socket has a chance to close first.
-      await new Promise((resolve) => setImmediate(resolve));
-      await new Promise((resolve) => setImmediate(resolve));
-      await promisify(mockServer.close).bind(mockServer)();
+      if (testTV.connected) {
+        await testTV.disconnect();
+      }
+
+      await new Promise((resolve) => {
+        mockServerSocket.end(() => {
+          resolve(undefined);
+        });
+        mockServer.unref();
+      });
     });
 
     it('connects', async () => {
@@ -325,60 +330,3 @@ describe.each([
     });
   },
 );
-
-describe.each([
-  {
-    ipProto: 'IPv4',
-    address: '127.0.0.1',
-    socketType: 'udp4' as SocketType,
-  },
-  { ipProto: 'IPv6', address: '::1', socketType: 'udp6' as SocketType },
-])('datagram commands $ipProto', ({ address, socketType }) => {
-  let mockSocket: DgramSocket;
-  let testSettings: typeof DefaultSettings;
-  let testTV: LGTV;
-
-  beforeEach(async () => {
-    mockSocket = createSocket(socketType);
-    await promisify(mockSocket.bind).bind(mockSocket)();
-    const port = mockSocket.address().port;
-    testSettings = {
-      ...DefaultSettings,
-      networkWolPort: port,
-      networkWolAddress: address,
-    };
-    testTV = new LGTV(address, MAC, CRYPT_KEY, testSettings);
-  });
-
-  afterEach(async () => {
-    await promisify(mockSocket.close).bind(mockSocket)();
-  });
-
-  it('powers on', async () => {
-    let received = false;
-    let contents: Buffer | null = null;
-    const mocking = new Promise<void>((resolve) => {
-      mockSocket.on('message', (msg) => {
-        received = true;
-        contents = msg;
-        resolve();
-      });
-    });
-    testTV.powerOn();
-    await mocking;
-    expect(received).toBe(true);
-    expect(contents).toStrictEqual(
-      Buffer.from([
-        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xda, 0x0a, 0x0f, 0xe1, 0x60, 0xcb,
-        0xda, 0x0a, 0x0f, 0xe1, 0x60, 0xcb, 0xda, 0x0a, 0x0f, 0xe1, 0x60, 0xcb,
-        0xda, 0x0a, 0x0f, 0xe1, 0x60, 0xcb, 0xda, 0x0a, 0x0f, 0xe1, 0x60, 0xcb,
-        0xda, 0x0a, 0x0f, 0xe1, 0x60, 0xcb, 0xda, 0x0a, 0x0f, 0xe1, 0x60, 0xcb,
-        0xda, 0x0a, 0x0f, 0xe1, 0x60, 0xcb, 0xda, 0x0a, 0x0f, 0xe1, 0x60, 0xcb,
-        0xda, 0x0a, 0x0f, 0xe1, 0x60, 0xcb, 0xda, 0x0a, 0x0f, 0xe1, 0x60, 0xcb,
-        0xda, 0x0a, 0x0f, 0xe1, 0x60, 0xcb, 0xda, 0x0a, 0x0f, 0xe1, 0x60, 0xcb,
-        0xda, 0x0a, 0x0f, 0xe1, 0x60, 0xcb, 0xda, 0x0a, 0x0f, 0xe1, 0x60, 0xcb,
-        0xda, 0x0a, 0x0f, 0xe1, 0x60, 0xcb,
-      ]),
-    );
-  });
-});