From a88c7348ceca020125b3b10fe3fb202d2339e0dc Mon Sep 17 00:00:00 2001
From: flakey5 <73616808+flakey5@users.noreply.github.com>
Date: Mon, 6 Nov 2023 03:14:48 +0000
Subject: [PATCH] src: add .opened promise (#14)

---
 src/index.ts       | 45 ++++++++++++++++++++++++++++++++++++++++-----
 src/types.ts       |  5 +++++
 test/index.test.ts | 10 +++++++---
 3 files changed, 52 insertions(+), 8 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index f9c1afe..b32619c 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -2,7 +2,7 @@ import net from 'node:net';
 import tls from 'node:tls';
 import { Duplex } from 'node:stream';
 import type { ReadableStream, WritableStream } from 'node:stream/web';
-import type { SocketAddress, SocketOptions } from './types';
+import type { SocketAddress, SocketInfo, SocketOptions } from './types';
 import { isSocketAddress } from './is-socket-address';
 
 export function connect(
@@ -23,11 +23,15 @@ export function connect(
 export class Socket {
   readable: ReadableStream<unknown>;
   writable: WritableStream<unknown>;
+  opened: Promise<SocketInfo>;
   closed: Promise<void>;
 
   private socket: net.Socket | tls.TLSSocket;
   private allowHalfOpen: boolean;
   private secureTransport: SocketOptions['secureTransport'];
+  private openedIsResolved: boolean;
+  private openedResolve!: (info: SocketInfo) => void;
+  private openedReject!: (reason?: unknown) => void;
   private closedResolve!: () => void;
   private closedReject!: (reason?: unknown) => void;
   private startTlsCalled = false;
@@ -39,6 +43,19 @@ export class Socket {
     this.secureTransport = options?.secureTransport ?? 'off';
     this.allowHalfOpen = options?.allowHalfOpen ?? true;
 
+    this.openedIsResolved = false;
+    this.opened = new Promise((resolve, reject) => {
+      this.openedResolve = (info): void => {
+        this.openedIsResolved = true;
+        resolve(info);
+      };
+      this.openedReject = (...args): void => {
+        this.openedIsResolved = true;
+        // eslint-disable-next-line prefer-promise-reject-errors -- ESLint gets this wrong as we are completely forwarding the arguments to reject.
+        reject(...args);
+      };
+    });
+
     this.closed = new Promise((resolve, reject) => {
       this.closedResolve = (...args): void => {
         resolve(...args);
@@ -64,6 +81,22 @@ export class Socket {
       this.socket = new tls.TLSSocket(addressOrSocket);
     }
 
+    if (this.socket instanceof tls.TLSSocket) {
+      this.socket.on('secureConnect', () => {
+        this.openedResolve({
+          remoteAddress: this.socket.remoteAddress,
+          localAddress: this.socket.localAddress,
+        });
+      });
+    } else {
+      this.socket.on('connect', () => {
+        this.openedResolve({
+          remoteAddress: this.socket.remoteAddress,
+          localAddress: this.socket.localAddress,
+        });
+      });
+    }
+
     this.socket.on('close', (hadError) => {
       if (!hadError) {
         this.closedResolve();
@@ -71,11 +104,13 @@ export class Socket {
     });
 
     this.socket.on('error', (err) => {
-      if (err instanceof Error) {
-        this.closedReject(new SocketError(err.message));
-      } else {
-        this.closedReject(new SocketError(err as string));
+      const socketError = new SocketError(
+        err instanceof Error ? err.message : (err as string),
+      );
+      if (!this.openedIsResolved) {
+        this.openedReject(socketError);
       }
+      this.closedReject(socketError);
     });
 
     // types are wrong. fixed based on docs https://nodejs.org/dist/latest/docs/api/stream.html#streamduplextowebstreamduplex
diff --git a/src/types.ts b/src/types.ts
index 42a091b..038fa2e 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -21,3 +21,8 @@ export interface SocketAddress {
   /** The port number to connect to. Example: `5432`. */
   port: number;
 }
+
+export interface SocketInfo {
+  remoteAddress?: string;
+  localAddress?: string;
+}
diff --git a/test/index.test.ts b/test/index.test.ts
index f29ed5c..db071ff 100644
--- a/test/index.test.ts
+++ b/test/index.test.ts
@@ -11,7 +11,7 @@ import {
 void tap.test(
   'Socket connected to tcp server with secureTransport: off',
   async (t) => {
-    t.plan(6);
+    t.plan(7);
     let connectCount = 0;
     const message = 'abcde\r\n';
 
@@ -38,6 +38,8 @@ void tap.test(
       'should pipe message',
     );
 
+    await t.resolveMatch(socket.opened, { localAddress: '::1' });
+
     const close = socket.close();
     t.equal(
       socket.closed,
@@ -137,13 +139,15 @@ for (const data of [
 }
 
 void tap.test('SocketError is thrown on connect failure', async (t) => {
-  t.plan(1);
+  t.plan(2);
 
+  const expectedError = new SocketError('connect ECONNREFUSED 127.0.0.1:1234');
   try {
     const socket = connect('127.0.0.1:1234');
+    socket.opened.catch((err) => t.same(err, expectedError));
     await socket.closed;
   } catch (err) {
-    t.same(err, new SocketError('connect ECONNREFUSED 127.0.0.1:1234'));
+    t.same(err, expectedError);
   } finally {
     t.end();
   }