chore: port browser-proxy to pg-gateway Web API

apps/browser-proxy/package.json

@@ -9,8 +9,10 @@
   "dependencies": {
     "@aws-sdk/client-s3": "^3.645.0",
     "debug": "^4.3.7",
+    "expiry-map": "^2.0.0",
     "findhit-proxywrap": "^0.3.13",
-    "pg-gateway": "0.3.0-alpha.7",
+    "p-memoize": "^7.1.1",
+    "pg-gateway": "^0.3.0-alpha.7",
     "ws": "^8.18.0"
   },
   "devDependencies": {

apps/browser-proxy/src/index.ts

@@ -1,43 +1,34 @@
 import * as nodeNet from 'node:net'
 import * as https from 'node:https'
-import { PostgresConnection } from 'pg-gateway'
+import { BackendError, PostgresConnection } from 'pg-gateway'
+import { fromNodeSocket } from 'pg-gateway/node'
 import { WebSocketServer, type WebSocket } from 'ws'
 import makeDebug from 'debug'
-import * as tls from 'node:tls'
 import { extractDatabaseId, isValidServername } from './servername.ts'
-import { getTls } from './tls.ts'
+import { getTls, setSecureContext } from './tls.ts'
 import { createStartupMessage } from './create-message.ts'
 import { extractIP } from './extract-ip.ts'
 
 const debug = makeDebug('browser-proxy')
 
-const tcpConnections = new Map<string, nodeNet.Socket>()
+const tcpConnections = new Map<string, PostgresConnection>()
 const websocketConnections = new Map<string, WebSocket>()
 
-let tlsOptions = await getTls()
-
-// refresh the TLS certificate every week
-setInterval(
-  async () => {
-    tlsOptions = await getTls()
-    httpsServer.setSecureContext(tlsOptions)
-  },
-  1000 * 60 * 60 * 24 * 7
-)
-
 const httpsServer = https.createServer({
-  ...tlsOptions,
   SNICallback: (servername, callback) => {
     debug('SNICallback', servername)
     if (isValidServername(servername)) {
       debug('SNICallback', 'valid')
-      callback(null, tls.createSecureContext(tlsOptions))
+      callback(null)
     } else {
       debug('SNICallback', 'invalid')
       callback(new Error('invalid SNI'))
     }
   },
 })
+await setSecureContext(httpsServer)
+// reset the secure context every week to pick up any new TLS certificates
+setInterval(() => setSecureContext(httpsServer), 1000 * 60 * 60 * 24 * 7)
 
 const websocketServer = new WebSocketServer({
   server: httpsServer,
@@ -70,8 +61,8 @@ websocketServer.on('connection', (socket, request) => {
 
   socket.on('message', (data: Buffer) => {
     debug('websocket message', data.toString('hex'))
-    const tcpSocket = tcpConnections.get(databaseId)
-    tcpSocket?.write(data)
+    const tcpConnection = tcpConnections.get(databaseId)
+    tcpConnection?.streamWriter?.write(data)
   })
 
   socket.on('close', () => {
@@ -86,50 +77,41 @@ const net = (
 
 const tcpServer = net.createServer()
 
-tcpServer.on('connection', (socket) => {
+tcpServer.on('connection', async (socket) => {
   let databaseId: string | undefined
 
-  const connection = new PostgresConnection(socket, {
-    tls: tlsOptions,
+  const connection = await fromNodeSocket(socket, {
+    tls: getTls,
     onTlsUpgrade(state) {
-      if (!state.tlsInfo?.sniServerName || !isValidServername(state.tlsInfo.sniServerName)) {
-        // connection.detach()
-        connection.sendError({
+      if (!state.tlsInfo?.serverName || !isValidServername(state.tlsInfo.serverName)) {
+        throw BackendError.create({
           code: '08006',
           message: 'invalid SNI',
           severity: 'FATAL',
         })
-        connection.end()
-        return
       }
 
-      const _databaseId = extractDatabaseId(state.tlsInfo.sniServerName!)
+      const _databaseId = extractDatabaseId(state.tlsInfo.serverName!)
 
       if (!websocketConnections.has(_databaseId!)) {
-        // connection.detach()
-        connection.sendError({
+        throw BackendError.create({
           code: 'XX000',
           message: 'the browser is not sharing the database',
           severity: 'FATAL',
         })
-        connection.end()
-        return
       }
 
       if (tcpConnections.has(_databaseId)) {
-        // connection.detach()
-        connection.sendError({
+        throw BackendError.create({
           code: '53300',
           message: 'sorry, too many clients already',
           severity: 'FATAL',
         })
-        connection.end()
-        return
       }
 
       // only set the databaseId after we've verified the connection
       databaseId = _databaseId
-      tcpConnections.set(databaseId!, connection.socket)
+      tcpConnections.set(databaseId!, connection)
     },
     serverVersion() {
       return '16.3'
@@ -138,13 +120,11 @@ tcpServer.on('connection', (socket) => {
       const websocket = websocketConnections.get(databaseId!)
 
       if (!websocket) {
-        connection.sendError({
+        throw BackendError.create({
           code: 'XX000',
           message: 'the browser is not sharing the database',
           severity: 'FATAL',
         })
-        connection.end()
-        return
       }
 
       const clientIpMessage = createStartupMessage('postgres', 'postgres', {
@@ -160,13 +140,11 @@ tcpServer.on('connection', (socket) => {
       const websocket = websocketConnections.get(databaseId!)
 
       if (!websocket) {
-        connection.sendError({
+        throw BackendError.create({
           code: 'XX000',
           message: 'the browser is not sharing the database',
           severity: 'FATAL',
         })
-        connection.end()
-        return
       }
 
       debug('tcp message', { message })

apps/browser-proxy/src/tls.ts

@@ -1,9 +1,12 @@
 import { Buffer } from 'node:buffer'
 import { GetObjectCommand, S3Client } from '@aws-sdk/client-s3'
+import pMemoize from 'p-memoize'
+import ExpiryMap from 'expiry-map'
+import type { Server } from 'node:https'
 
 const s3Client = new S3Client({ forcePathStyle: true })
 
-export async function getTls() {
+async function _getTls() {
   const cert = await s3Client
     .send(
       new GetObjectCommand({
@@ -31,3 +34,12 @@ export async function getTls() {
     key: Buffer.from(key),
   }
 }
+
+// cache the TLS certificate for 1 week
+const cache = new ExpiryMap(1000 * 60 * 60 * 24 * 7)
+export const getTls = pMemoize(_getTls, { cache })
+
+export async function setSecureContext(httpsServer: Server) {
+  const tlsOptions = await getTls()
+  httpsServer.setSecureContext(tlsOptions)
+}

apps/postgres-new/components/app-provider.tsx

@@ -138,7 +138,17 @@ export default function AppProvider({ children }: AppProps) {
         if (isStartupMessage(message)) {
           const parameters = parseStartupMessage(message)
           if ('client_ip' in parameters) {
-            setConnectedClientIp(parameters.client_ip === '' ? null : parameters.client_ip)
+            // client disconnected
+            if (parameters.client_ip === '') {
+              setConnectedClientIp(null)
+              // we ensure we're not in a transaction block first
+              await db.sql`rollback;`.catch()
+              // we clean the session state, see: https://www.pgbouncer.org/faq.html#how-to-use-prepared-statements-with-session-pooling
+              // we do this to avoid having old prepared statements in the session
+              await db.sql`discard all;`
+            } else {
+              setConnectedClientIp(parameters.client_ip)
+            }
           }
           return
         }

apps/postgres-new/components/live-share-icon.tsx

@@ -11,8 +11,8 @@ export const LiveShareIcon = (props: { size?: number; className?: string }) => (
     className={cx('lucide lucide-live-share', props.className)}
   >
     <path
-      fill-rule="evenodd"
-      clip-rule="evenodd"
+      fillRule="evenodd"
+      clipRule="evenodd"
       d="M13.735 1.694L15.178 1l8.029 6.328v1.388l-8.029 6.072-1.443-.694v-2.776h-.59c-4.06-.02-6.71.104-10.61 5.163l-1.534-.493a8.23 8.23 0 0 1 .271-2.255 11.026 11.026 0 0 1 3.92-6.793 11.339 11.339 0 0 1 7.502-2.547h1.04v-2.7zm1.804 7.917v2.776l5.676-4.281-5.648-4.545v2.664h-2.86A9.299 9.299 0 0 0 5.77 8.848a10.444 10.444 0 0 0-2.401 4.122c3.351-3.213 6.19-3.359 9.798-3.359h2.373zm-7.647 5.896a4.31 4.31 0 1 1 4.788 7.166 4.31 4.31 0 0 1-4.788-7.166zm.955 5.728a2.588 2.588 0 1 0 2.878-4.302 2.588 2.588 0 0 0-2.878 4.302z"
     />
   </svg>