fix clientIp message interception

Author: jgoux

apps/browser-proxy/src/create-message.ts

@@ -1,22 +1,48 @@
-export function createParameterStatusMessage(name: string, value: string): ArrayBuffer {
+export function createStartupMessage(
+  user: string,
+  database: string,
+  additionalParams: Record<string, string> = {}
+): ArrayBuffer {
   const encoder = new TextEncoder()
-  const nameBuffer = encoder.encode(name + '\0')
-  const valueBuffer = encoder.encode(value + '\0')
 
-  const messageLength = 4 + nameBuffer.length + valueBuffer.length
-  const message = new ArrayBuffer(1 + messageLength)
+  // Protocol version number (3.0)
+  const protocolVersion = 196608
+
+  // Combine required and additional parameters
+  const params = {
+    user,
+    database,
+    ...additionalParams,
+  }
+
+  // Calculate total message length
+  let messageLength = 4 // Protocol version
+  for (const [key, value] of Object.entries(params)) {
+    messageLength += key.length + 1 + value.length + 1
+  }
+  messageLength += 1 // Null terminator
+
+  const message = new ArrayBuffer(4 + messageLength)
   const view = new DataView(message)
   const uint8Array = new Uint8Array(message)
 
   let offset = 0
-  view.setUint8(offset++, 'S'.charCodeAt(0)) // Message type
-  view.setUint32(offset, messageLength, false) // Message length (big-endian)
+  view.setInt32(offset, messageLength + 4, false) // Total message length (including itself)
+  offset += 4
+  view.setInt32(offset, protocolVersion, false) // Protocol version number
   offset += 4
 
-  uint8Array.set(nameBuffer, offset)
-  offset += nameBuffer.length
+  // Write key-value pairs
+  for (const [key, value] of Object.entries(params)) {
+    uint8Array.set(encoder.encode(key), offset)
+    offset += key.length
+    uint8Array.set([0], offset++) // Null terminator for key
+    uint8Array.set(encoder.encode(value), offset)
+    offset += value.length
+    uint8Array.set([0], offset++) // Null terminator for value
+  }
 
-  uint8Array.set(valueBuffer, offset)
+  uint8Array.set([0], offset) // Final null terminator
 
   return message
 }

apps/browser-proxy/src/index.ts

@@ -6,7 +6,7 @@ import makeDebug from 'debug'
 import * as tls from 'node:tls'
 import { extractDatabaseId, isValidServername } from './servername.ts'
 import { getTls } from './tls.ts'
-import { createParameterStatusMessage } from './create-message.ts'
+import { createStartupMessage } from './create-message.ts'
 import { extractIP } from './extract-ip.ts'
 
 const debug = makeDebug('browser-proxy')
@@ -147,10 +147,9 @@ tcpServer.on('connection', (socket) => {
         return
       }
 
-      const clientIpMessage = createParameterStatusMessage(
-        'client_ip',
-        extractIP(socket.remoteAddress!)
-      )
+      const clientIpMessage = createStartupMessage('postgres', 'postgres', {
+        client_ip: extractIP(connection.socket.remoteAddress!),
+      })
       websocket.send(clientIpMessage)
     },
     onMessage(message, state) {
@@ -182,7 +181,7 @@ tcpServer.on('connection', (socket) => {
     if (databaseId) {
       tcpConnections.delete(databaseId)
       const websocket = websocketConnections.get(databaseId)
-      websocket?.send(createParameterStatusMessage('client_ip', ''))
+      websocket?.send(createStartupMessage('postgres', 'postgres', { client_ip: '' }))
     }
   })
 })

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

@@ -17,7 +17,7 @@ import {
 } from 'react'
 import { DbManager } from '~/lib/db'
 import { useAsyncMemo } from '~/lib/hooks'
-import { parseParameterStatus } from '~/lib/pg-wire-util'
+import { isStartupMessage, parseStartupMessage } from '~/lib/pg-wire-util'
 import { createClient } from '~/utils/supabase/client'
 
 export type AppProps = PropsWithChildren
@@ -131,19 +131,20 @@ export default function AppProvider({ children }: AppProps) {
       ws.onopen = () => {
         setLiveSharedDatabaseId(databaseId)
       }
+
       ws.onmessage = async (event) => {
         const message = new Uint8Array(await event.data)
 
-        const messageType = String.fromCharCode(message[0])
-        if (messageType === 'S') {
-          const { name, value } = parseParameterStatus(message)
-          if (name === 'client_ip') {
-            setConnectedClientIp(value === '' ? null : value)
-            return
+        if (isStartupMessage(message)) {
+          const parameters = parseStartupMessage(message)
+          if ('client_ip' in parameters) {
+            setConnectedClientIp(parameters.client_ip === '' ? null : parameters.client_ip)
           }
+          return
         }
 
         const response = await db.execProtocolRaw(message)
+
         ws.send(response)
       }
       ws.onclose = (event) => {

apps/postgres-new/electric-sql-pglite-0.2.6.tgz

@@ -1,21 +1,80 @@
-export function parseParameterStatus(data: Uint8Array): { name: string; value: string } {
+export function isStartupMessage(message: Uint8Array): boolean {
+  if (message.length < 8) {
+    return false // Message is too short to be a valid startup message
+  }
+
+  const view = new DataView(message.buffer, message.byteOffset, message.byteLength)
+
+  // Get the message length (first 4 bytes)
+  const messageLength = view.getInt32(0, false) // big-endian
+
+  // Check if the message length matches the actual buffer length
+  if (messageLength !== message.length) {
+    return false
+  }
+
+  // Get the protocol version (next 4 bytes)
+  const protocolVersion = view.getInt32(4, false) // big-endian
+
+  // Check if the protocol version is valid (3.0)
+  // Protocol version 3.0 is represented as 196608 (0x00030000)
+  if (protocolVersion !== 196608) {
+    return false
+  }
+
+  // Additional check: ensure the message ends with a null terminator
+  if (message[message.length - 1] !== 0) {
+    return false
+  }
+
+  return true
+}
+
+export function parseStartupMessage(message: Uint8Array): {
+  user: string
+  database: string
+  [key: string]: string
+} {
+  const view = new DataView(message.buffer, message.byteOffset, message.byteLength)
   const decoder = new TextDecoder()
+  const params: {
+    user: string
+    database: string
+    [key: string]: string
+  } = {}
+
+  // Skip the message length (4 bytes) and protocol version (4 bytes)
+  let offset = 8
+
+  while (offset < message.length) {
+    let keyStart = offset
+    let keyEnd = keyStart
+    while (message[keyEnd] !== 0 && keyEnd < message.length) {
+      keyEnd++
+    }
+    if (keyEnd === message.length) break // End of message
+
+    const key = decoder.decode(message.subarray(keyStart, keyEnd))
+    offset = keyEnd + 1 // Skip null terminator
 
-  // Skip message type (1 byte) and length (4 bytes)
-  let offset = 5
+    let valueStart = offset
+    let valueEnd = valueStart
+    while (message[valueEnd] !== 0 && valueEnd < message.length) {
+      valueEnd++
+    }
+    if (valueEnd === message.length) break // End of message
 
-  // Find the null terminator for the name
-  let nameEnd = offset
-  while (data[nameEnd] !== 0) nameEnd++
+    const value = decoder.decode(message.subarray(valueStart, valueEnd))
+    offset = valueEnd + 1 // Skip null terminator
 
-  const name = decoder.decode(data.subarray(offset, nameEnd))
-  offset = nameEnd + 1
+    params[key] = value
 
-  // Find the null terminator for the value
-  let valueEnd = offset
-  while (data[valueEnd] !== 0) valueEnd++
+    if (message[offset] === 0) break // Final null terminator
+  }
 
-  const value = decoder.decode(data.subarray(offset, valueEnd))
+  if (!params.user || !params.database) {
+    throw new Error('user or database not found in startup message')
+  }
 
-  return { name, value }
+  return params
 }

apps/postgres-new/lib/pg-wire-util.ts

@@ -12,7 +12,7 @@
   "dependencies": {
     "@ai-sdk/openai": "^0.0.21",
     "@dagrejs/dagre": "^1.1.2",
-    "@electric-sql/pglite": "0.2.6",
+    "@electric-sql/pglite": "file:./electric-sql-pglite-0.2.6.tgz",
     "@gregnr/postgres-meta": "^0.82.0-dev.2",
     "@monaco-editor/react": "^4.6.0",
     "@radix-ui/react-accordion": "^1.2.0",

apps/postgres-new/package.json

@@ -3,9 +3,7 @@
   "scripts": {
     "dev": "npm run dev --workspace postgres-new"
   },
-  "workspaces": [
-    "apps/*"
-  ],
+  "workspaces": ["apps/*"],
   "devDependencies": {
     "supabase": "^1.191.3"
   }