Add custom socket option - fixes porsager#284

porsager · porsager · commit 5413f0c8bce6 · 2022-04-02T22:25:04.000+02:00
diff --git a/README.md b/README.md
@@ -622,6 +622,7 @@ const sql = postgres('postgres://username:password@host:port/database', {
   onnotice             : fn,            // Defaults to console.log
   onparameter          : fn,            // (key, value) when server param change
   debug                : fn,            // Is called with (connection, query, params, types)
+  socket               : fn,            // fn returning custom socket to use
   transform            : {
     column             : fn,            // Transforms incoming column names
     value              : fn,            // Transforms incoming row values
@@ -768,6 +769,31 @@ const [custom] = sql`
 
 ```
 
+### Custom socket
+
+Easily do in-process ssh tunneling to your database by providing a custom socket for Postgres.js to use. The function (optionally async) must return a socket-like duplex stream.
+
+Here's a sample using [ssh2](https://github.com/mscdex/ssh2)
+
+```js
+import ssh2 from 'ssh2'
+
+const sql = postgres({
+  ...options,
+  socket: ({ hostname, port }) => new Promise((resolve, reject) => {
+    const ssh = new ssh2.Client()
+    ssh
+    .on('error', reject)
+    .on('ready', () => 
+      ssh.forwardOut('127.0.0.1', 12345, hostname, port, 
+        (err, socket) => err ? reject(err) : resolve(socket)
+      )
+    )
+    .connect(sshOptions)
+  })
+})
+```
+
 ## Teardown / Cleanup
 
 To ensure proper teardown and cleanup on server restarts use `await sql.end()` before `process.exit()`.
diff --git a/src/connection.js b/src/connection.js
@@ -74,7 +74,7 @@ function Connection(options, queues = {}, { onopen = noop, onend = noop, onclose
       , lifeTimer = timer(end, options.max_lifetime)
       , connectTimer = timer(connectTimedOut, options.connect_timeout)
 
-  let socket = createSocket()
+  let socket = null
     , result = new Result()
     , incoming = Buffer.alloc(0)
     , needsTypes = options.fetch_types
@@ -122,15 +122,27 @@ function Connection(options, queues = {}, { onopen = noop, onend = noop, onclose
 
   return connection
 
-  function createSocket() {
-    const x = net.Socket()
+  async function createSocket() {
+    let x
+    try {
+      x = options.socket
+        ? (await Promise.resolve(options.socket(options)))
+        : net.Socket()
+    } catch (e) {
+      error(e)
+      return
+    }
     x.on('error', error)
     x.on('close', closed)
     x.on('drain', drain)
     return x
   }
 
-  function cancel({ pid, secret }, resolve, reject) {
+  async function cancel({ pid, secret }, resolve, reject) {
+    socket || (socket = await createSocket())
+    if (!socket)
+      return
+
     socket.removeAllListeners()
     socket = net.Socket()
     socket.on('connect', () => socket.write(b().i32(16).i32(80877102).i32(pid).i32(secret).end(16)))
@@ -324,10 +336,19 @@ function Connection(options, queues = {}, { onopen = noop, onend = noop, onclose
     }
   }
 
-  function connect() {
+  async function connect() {
     terminated = false
     backendParameters = {}
+    socket || (socket = await createSocket())
+
+    if (!socket)
+      return
+
     connectTimer.start()
+
+    if (options.socket)
+      return ssl ? secure() : connected()
+
     socket.on('connect', ssl ? secure : connected)
 
     if (options.path)
@@ -349,7 +370,7 @@ function Connection(options, queues = {}, { onopen = noop, onend = noop, onclose
       statementCount = 1
       lifeTimer.start()
       socket.on('data', data)
-      socket.setKeepAlive(true, 1000 * keep_alive)
+      socket.setKeepAlive && socket.setKeepAlive(true, 1000 * keep_alive)
       const s = StartupMessage()
       write(s)
     } catch (err) {
@@ -397,13 +418,15 @@ function Connection(options, queues = {}, { onopen = noop, onend = noop, onclose
       error(Errors.connection('CONNECTION_DESTROYED', options))
 
     clearImmediate(nextWriteTimer)
-    socket.removeListener('data', data)
-    socket.removeListener('connect', connected)
-    socket.readyState !== 'closed' && socket.end(b().X().end())
+    if (socket) {
+      socket.removeListener('data', data)
+      socket.removeListener('connect', connected)
+      socket.readyState !== 'closed' && socket.end(b().X().end())
+    }
     ended && (ended(), ending = ended = null)
   }
 
-  function closed(hadError) {
+  async function closed(hadError) {
     incoming = Buffer.alloc(0)
     remaining = 0
     incomings = null
@@ -416,7 +439,7 @@ function Connection(options, queues = {}, { onopen = noop, onend = noop, onclose
 
     if (socket.encrypted) {
       socket.removeAllListeners()
-      socket = createSocket()
+      socket = null
     }
 
     if (initial)
diff --git a/src/index.js b/src/index.js
@@ -397,6 +397,7 @@ function parseOptions(a, b) {
     connection      : Object.assign({ application_name: 'postgres.js' }, o.connection),
     target_session_attrs: tsa(o, url, env),
     debug           : o.debug,
+    socket          : o.socket,
     fetch_types     : 'fetch_types' in o ? o.fetch_types : true,
     parameters      : {},
     shared          : { retries: 0, typeArrayMap: {} },
diff --git a/tests/index.js b/tests/index.js
@@ -1954,3 +1954,24 @@ t('Ensure reconnect after max_lifetime with transactions', { timeout: 5000 }, as
 
   return [true, true]
 })
+
+t('Custom socket works', {}, async() => {
+  let result
+  const sql = postgres({
+    socket: () => new Promise((resolve, reject) => {
+      const socket = net.Socket()
+      socket.connect(5432)
+      socket.once('data', x => result = x[0])
+      socket.on('error', reject)
+      socket.on('connect', () => resolve(socket))
+    }),
+    idle_timeout
+  })
+
+  await sql`select 1`
+
+  return [
+    result,
+    82
+  ]
+})
