orez 0.0.38 → 0.0.40

package/dist/pg-proxy.js

@@ -1,8 +1,9 @@
 /**
  * tcp proxy that makes pglite speak postgresql wire protocol.
  *
- * uses pg-gateway to handle protocol lifecycle for regular connections,
- * and directly handles the raw socket for replication connections.
+ * handles the postgresql wire protocol directly using raw tcp sockets,
+ * avoiding pg-gateway's Duplex.toWeb() which deadlocks under concurrent
+ * connections with large responses.
  *
  * regular connections: forwarded to pglite via execProtocolRaw()
  * replication connections: intercepted, replication protocol faked
@@ -12,7 +13,6 @@
  * query interleaving that causes CVR concurrent modification errors.
  */
 import { createServer } from 'node:net';
-import { fromNodeSocket } from 'pg-gateway/node';
 import { log } from './log.js';
 import { Mutex } from './mutex.js';
 import { handleReplicationQuery, handleStartReplication } from './replication/handler.js';
@@ -55,6 +55,7 @@ const QUERY_REWRITES = [
 // parameter status messages sent during connection handshake
 // pg_restore and other tools read these to determine server capabilities
 const SERVER_PARAMS = [
+    ['server_version', '16.4'],
     ['server_encoding', 'UTF8'],
     ['client_encoding', 'UTF8'],
     ['DateStyle', 'ISO, MDY'],
@@ -63,7 +64,10 @@ const SERVER_PARAMS = [
     ['TimeZone', 'UTC'],
     ['IntervalStyle', 'postgres'],
 ];
-// build a ParameterStatus wire protocol message (type 'S', 0x53)
+// queries to intercept and return no-op success (synthetic SET response)
+// pglite rejects SET TRANSACTION if any query (e.g. SET search_path) ran first
+const NOOP_QUERY_PATTERNS = [/^\s*SET\s+TRANSACTION\b/i, /^\s*SET\s+SESSION\b/i];
+// ── wire protocol helpers ──
 function buildParameterStatus(name, value) {
     const encoder = new TextEncoder();
     const nameBytes = encoder.encode(name);
@@ -81,12 +85,58 @@ function buildParameterStatus(name, value) {
     buf[pos] = 0;
     return buf;
 }
-// queries to intercept and return no-op success (synthetic SET response)
-// pglite rejects SET TRANSACTION if any query (e.g. SET search_path) ran first
-const NOOP_QUERY_PATTERNS = [/^\s*SET\s+TRANSACTION\b/i, /^\s*SET\s+SESSION\b/i];
-/**
- * extract query text from a Parse message (0x50).
- */
+function buildAuthOk() {
+    const buf = new Uint8Array(9);
+    buf[0] = 0x52; // 'R' AuthenticationOk
+    new DataView(buf.buffer).setInt32(1, 8);
+    new DataView(buf.buffer).setInt32(5, 0); // auth ok
+    return buf;
+}
+function buildAuthCleartextPassword() {
+    const buf = new Uint8Array(9);
+    buf[0] = 0x52; // 'R'
+    new DataView(buf.buffer).setInt32(1, 8);
+    new DataView(buf.buffer).setInt32(5, 3); // cleartext password
+    return buf;
+}
+function buildBackendKeyData() {
+    const buf = new Uint8Array(13);
+    buf[0] = 0x4b; // 'K'
+    new DataView(buf.buffer).setInt32(1, 12);
+    new DataView(buf.buffer).setInt32(5, process.pid);
+    new DataView(buf.buffer).setInt32(9, 0);
+    return buf;
+}
+function buildReadyForQuery(status = 0x49) {
+    const buf = new Uint8Array(6);
+    buf[0] = 0x5a; // 'Z'
+    new DataView(buf.buffer).setInt32(1, 5);
+    buf[5] = status; // 'I' = idle
+    return buf;
+}
+function buildErrorResponse(message) {
+    const encoder = new TextEncoder();
+    const msgBytes = encoder.encode(message);
+    // S(ERROR) + C(code) + M(message) + terminator
+    const sField = new Uint8Array([0x53, ...encoder.encode('ERROR'), 0]);
+    const cField = new Uint8Array([0x43, ...encoder.encode('08006'), 0]);
+    const mField = new Uint8Array([0x4d, ...msgBytes, 0]);
+    const terminator = new Uint8Array([0]);
+    const bodyLen = 4 + sField.length + cField.length + mField.length + terminator.length;
+    const buf = new Uint8Array(1 + bodyLen);
+    buf[0] = 0x45; // 'E'
+    new DataView(buf.buffer).setInt32(1, bodyLen);
+    let pos = 5;
+    buf.set(sField, pos);
+    pos += sField.length;
+    buf.set(cField, pos);
+    pos += cField.length;
+    buf.set(mField, pos);
+    pos += mField.length;
+    buf.set(terminator, pos);
+    return buf;
+}
+// ── query helpers ──
 function extractParseQuery(data) {
     if (data[0] !== 0x50)
         return null;
@@ -99,9 +149,6 @@ function extractParseQuery(data) {
         offset++;
     return new TextDecoder().decode(data.subarray(queryStart, offset));
 }
-/**
- * rebuild a Parse message with a modified query string.
- */
 function rebuildParseMessage(data, newQuery) {
     let offset = 5;
     while (offset < data.length && data[offset] !== 0)
@@ -129,9 +176,6 @@ function rebuildParseMessage(data, newQuery) {
     result.set(suffix, pos);
     return result;
 }
-/**
- * rebuild a Simple Query message with a modified query string.
- */
 function rebuildSimpleQuery(newQuery) {
     const encoder = new TextEncoder();
     const queryBytes = encoder.encode(newQuery + '\0');
@@ -141,9 +185,6 @@ function rebuildSimpleQuery(newQuery) {
     buf.set(queryBytes, 5);
     return buf;
 }
-/**
- * intercept and rewrite query messages to make pglite look like real postgres.
- */
 function interceptQuery(data) {
     const msgType = data[0];
     if (msgType === 0x51) {
@@ -183,9 +224,6 @@ function interceptQuery(data) {
     }
     return data;
 }
-/**
- * check if a query should be intercepted as a no-op.
- */
 function isNoopQuery(data) {
     let query = null;
     if (data[0] === 0x51) {
@@ -200,9 +238,6 @@ function isNoopQuery(data) {
         return false;
     return NOOP_QUERY_PATTERNS.some((p) => p.test(query));
 }
-/**
- * build a synthetic "SET" command complete response.
- */
 function buildSetCompleteResponse() {
     const encoder = new TextEncoder();
     const tag = encoder.encode('SET\0');
@@ -219,18 +254,12 @@ function buildSetCompleteResponse() {
     result.set(rfq, cc.length);
     return result;
 }
-/**
- * build a synthetic ParseComplete response for extended protocol no-ops.
- */
 function buildParseCompleteResponse() {
     const pc = new Uint8Array(5);
     pc[0] = 0x31; // ParseComplete
     new DataView(pc.buffer).setInt32(1, 4);
     return pc;
 }
-/**
- * strip ReadyForQuery messages from a response buffer.
- */
 function stripReadyForQuery(data) {
     if (data.length === 0)
         return data;
@@ -260,6 +289,257 @@ function stripReadyForQuery(data) {
     }
     return result;
 }
+// ── socket write with backpressure ──
+function socketWrite(socket, data) {
+    if (data.length === 0 || socket.destroyed)
+        return Promise.resolve();
+    return new Promise((resolve, reject) => {
+        const ok = socket.write(data, (err) => (err ? reject(err) : resolve()));
+        // if buffer is full, the callback still fires when flushed
+        if (!ok)
+            void 0;
+    });
+}
+// ── startup handshake ──
+// parse startup message from raw bytes.
+// handles SSLRequest (8 bytes, code 80877103) and StartupMessage.
+function parseStartupMessage(buf) {
+    const dv = new DataView(buf.buffer, buf.byteOffset, buf.byteLength);
+    const len = dv.getInt32(0);
+    const code = dv.getInt32(4);
+    // SSL request: length=8, code=80877103
+    if (len === 8 && code === 80877103) {
+        return { isSSL: true, params: {} };
+    }
+    // startup message: length, protocol(196608=3.0), then key=value pairs
+    const params = {};
+    let offset = 8;
+    while (offset < len) {
+        const keyStart = offset;
+        while (offset < buf.length && buf[offset] !== 0)
+            offset++;
+        const key = buf.subarray(keyStart, offset).toString();
+        offset++;
+        if (!key)
+            break; // double-null = end of params
+        const valStart = offset;
+        while (offset < buf.length && buf[offset] !== 0)
+            offset++;
+        params[key] = buf.subarray(valStart, offset).toString();
+        offset++;
+    }
+    return { isSSL: false, params };
+}
+// read exactly `n` bytes from socket
+function readBytes(socket, n) {
+    return new Promise((resolve, reject) => {
+        let collected = Buffer.alloc(0);
+        const onData = (chunk) => {
+            collected = Buffer.concat([collected, chunk]);
+            if (collected.length >= n) {
+                socket.removeListener('data', onData);
+                socket.removeListener('error', onError);
+                socket.removeListener('close', onClose);
+                socket.pause();
+                resolve(collected);
+            }
+        };
+        const onError = (err) => {
+            socket.removeListener('data', onData);
+            socket.removeListener('close', onClose);
+            reject(err);
+        };
+        const onClose = () => {
+            socket.removeListener('data', onData);
+            socket.removeListener('error', onError);
+            reject(new Error('socket closed'));
+        };
+        socket.on('data', onData);
+        socket.on('error', onError);
+        socket.on('close', onClose);
+        socket.resume();
+    });
+}
+// perform the startup handshake (SSL negotiation, auth, parameter status)
+async function performHandshake(socket, config) {
+    // read initial message length (first 4 bytes)
+    let buf = await readBytes(socket, 8);
+    // check for SSL request
+    const startup = parseStartupMessage(buf);
+    if (startup.isSSL) {
+        // reject SSL, client will reconnect without it
+        socket.write(Buffer.from('N'));
+        buf = await readBytes(socket, 8);
+    }
+    // now we have startup message header - read the rest if needed
+    const dv = new DataView(buf.buffer, buf.byteOffset, buf.byteLength);
+    const msgLen = dv.getInt32(0);
+    if (buf.length < msgLen) {
+        const rest = await readBytes(socket, msgLen - buf.length);
+        buf = Buffer.concat([buf, rest]);
+    }
+    const { params } = parseStartupMessage(buf);
+    // request cleartext password
+    socket.write(buildAuthCleartextPassword());
+    // read password message: type(1) + len(4) + password + null
+    const pwBuf = await readBytes(socket, 5);
+    const pwDv = new DataView(pwBuf.buffer, pwBuf.byteOffset, pwBuf.byteLength);
+    const pwLen = pwDv.getInt32(1);
+    let fullPwBuf = pwBuf;
+    if (fullPwBuf.length < 1 + pwLen) {
+        const rest = await readBytes(socket, 1 + pwLen - fullPwBuf.length);
+        fullPwBuf = Buffer.concat([fullPwBuf, rest]);
+    }
+    const password = fullPwBuf.subarray(5, 1 + pwLen - 1).toString();
+    // validate credentials
+    if (params.user !== config.pgUser || password !== config.pgPassword) {
+        socket.write(buildErrorResponse('authentication failed'));
+        socket.write(buildReadyForQuery());
+        socket.destroy();
+        throw new Error('auth failed');
+    }
+    // auth ok
+    socket.write(buildAuthOk());
+    // send parameter status messages
+    for (const [name, value] of SERVER_PARAMS) {
+        socket.write(buildParameterStatus(name, value));
+    }
+    // backend key data
+    socket.write(buildBackendKeyData());
+    // ready for query
+    socket.write(buildReadyForQuery());
+    return { params };
+}
+// ── message loop ──
+// process messages from a connected, authenticated client.
+// uses callback-based 'data' events instead of async iterators
+// for reliable behavior across runtimes (node.js, bun).
+function messageLoop(socket, db, mutex, isReplicationConnection, replicationDb, replicationMutex) {
+    return new Promise((resolve, reject) => {
+        let buffer = Buffer.alloc(0);
+        let processing = false;
+        async function processBuffer() {
+            if (processing)
+                return;
+            processing = true;
+            socket.pause();
+            try {
+                while (buffer.length >= 5) {
+                    const msgType = buffer[0];
+                    const dv = new DataView(buffer.buffer, buffer.byteOffset, buffer.byteLength);
+                    const msgLen = dv.getInt32(1);
+                    const totalLen = 1 + msgLen;
+                    if (buffer.length < totalLen)
+                        break; // need more data
+                    // copy message out before modifying buffer
+                    const message = new Uint8Array(buffer.buffer.slice(buffer.byteOffset, buffer.byteOffset + totalLen));
+                    buffer = buffer.subarray(totalLen);
+                    // handle Terminate message
+                    if (msgType === 0x58) {
+                        resolve();
+                        return;
+                    }
+                    // handle replication connections
+                    if (isReplicationConnection) {
+                        await handleReplicationMsg(message, socket, replicationDb, replicationMutex);
+                        continue;
+                    }
+                    // handle regular messages
+                    await handleRegularMessage(message, socket, db, mutex);
+                }
+            }
+            catch (err) {
+                reject(err);
+                return;
+            }
+            processing = false;
+            socket.resume();
+        }
+        socket.on('data', (chunk) => {
+            buffer = buffer.length > 0 ? Buffer.concat([buffer, chunk]) : chunk;
+            processBuffer();
+        });
+        socket.on('end', () => resolve());
+        socket.on('error', (err) => reject(err));
+        socket.on('close', () => resolve());
+        socket.resume();
+    });
+}
+async function handleRegularMessage(data, socket, db, mutex) {
+    // check for no-op queries
+    if (isNoopQuery(data)) {
+        if (data[0] === 0x51) {
+            await socketWrite(socket, buildSetCompleteResponse());
+            return;
+        }
+        else if (data[0] === 0x50) {
+            await socketWrite(socket, buildParseCompleteResponse());
+            return;
+        }
+    }
+    // intercept and rewrite queries
+    data = interceptQuery(data);
+    // serialize pglite access
+    await mutex.acquire();
+    let result;
+    try {
+        result = await db.execProtocolRaw(data, { throwOnError: false });
+    }
+    catch (err) {
+        mutex.release();
+        throw err;
+    }
+    // strip ReadyForQuery from non-Sync/non-SimpleQuery responses
+    if (data[0] !== 0x53 && data[0] !== 0x51) {
+        result = stripReadyForQuery(result);
+    }
+    mutex.release();
+    // write response directly to socket
+    await socketWrite(socket, result);
+}
+async function handleReplicationMsg(data, socket, db, mutex) {
+    if (data[0] !== 0x51)
+        return;
+    const view = new DataView(data.buffer, data.byteOffset, data.byteLength);
+    const len = view.getInt32(1);
+    const query = new TextDecoder().decode(data.subarray(5, 1 + len - 1)).replace(/\0$/, '');
+    const upper = query.trim().toUpperCase();
+    log.debug.proxy(`repl query: ${query.slice(0, 200)}`);
+    if (upper.startsWith('START_REPLICATION')) {
+        const writer = {
+            write(chunk) {
+                if (!socket.destroyed) {
+                    socket.write(chunk);
+                }
+            },
+        };
+        // drain incoming standby status updates
+        socket.on('data', (_chunk) => { });
+        socket.on('close', () => socket.destroy());
+        // this runs indefinitely until the socket closes
+        await handleStartReplication(query, writer, db, mutex).catch((err) => {
+            log.debug.proxy(`replication stream ended: ${err}`);
+        });
+        return;
+    }
+    // handle replication queries + fallthrough to pglite
+    await mutex.acquire();
+    try {
+        const response = await handleReplicationQuery(query, db);
+        if (response) {
+            await socketWrite(socket, response);
+            return;
+        }
+        // apply query rewrites before forwarding
+        data = interceptQuery(data);
+        const result = await db.execProtocolRaw(data, { throwOnError: false });
+        await socketWrite(socket, result);
+    }
+    finally {
+        mutex.release();
+    }
+}
+// ── main entry point ──
 export async function startPgProxy(dbInput, config) {
     // normalize input: single PGlite instance = use it for all databases (backwards compat for tests)
     const instances = 'postgres' in dbInput
@@ -271,7 +551,6 @@ export async function startPgProxy(dbInput, config) {
         cvr: new Mutex(),
         cdb: new Mutex(),
     };
-    // helper to get instance + mutex for a database name
     function getDbContext(dbName) {
         if (dbName === 'zero_cvr')
             return { db: instances.cvr, mutex: mutexes.cvr };
@@ -280,105 +559,39 @@ export async function startPgProxy(dbInput, config) {
         return { db: instances.postgres, mutex: mutexes.postgres };
     }
     const server = createServer(async (socket) => {
-        // prevent idle timeouts from killing connections
         socket.setKeepAlive(true, 30000);
         socket.setTimeout(0);
+        socket.setNoDelay(true);
         let dbName = 'postgres';
         let isReplicationConnection = false;
-        // clean up pglite transaction state when a client disconnects
-        socket.on('close', async () => {
-            const { db, mutex } = getDbContext(dbName);
-            await mutex.acquire();
-            try {
-                await db.exec('ROLLBACK');
-            }
-            catch {
-                // no transaction to rollback
-            }
-            finally {
-                mutex.release();
-            }
-        });
         try {
-            const connection = await fromNodeSocket(socket, {
-                serverVersion: '16.4',
-                auth: {
-                    method: 'password',
-                    getClearTextPassword() {
-                        return config.pgPassword;
-                    },
-                    validateCredentials(credentials) {
-                        return (credentials.password === credentials.clearTextPassword &&
-                            credentials.username === config.pgUser);
-                    },
-                },
-                // send ParameterStatus messages that standard postgres tools expect
-                // pg-gateway sends server_version via the serverVersion option above,
-                // but tools like pg_restore also need encoding, datestyle, etc.
-                onAuthenticated() {
-                    for (const [name, value] of SERVER_PARAMS) {
-                        socket.write(buildParameterStatus(name, value));
-                    }
-                },
-                async onStartup(state) {
-                    const params = state.clientParams;
-                    if (params?.replication === 'database') {
-                        isReplicationConnection = true;
-                    }
-                    dbName = params?.database || 'postgres';
-                    log.debug.proxy(`connection: db=${dbName} user=${params?.user} replication=${params?.replication || 'none'}`);
-                    const { db } = getDbContext(dbName);
-                    await db.waitReady;
-                },
-                async onMessage(data, state) {
-                    if (!state.isAuthenticated)
-                        return;
-                    // handle replication connections (always go to postgres instance)
-                    if (isReplicationConnection) {
-                        if (data[0] === 0x51) {
-                            const view = new DataView(data.buffer, data.byteOffset, data.byteLength);
-                            const len = view.getInt32(1);
-                            const query = new TextDecoder()
-                                .decode(data.subarray(5, 1 + len - 1))
-                                .replace(/\0$/, '');
-                            log.debug.proxy(`repl query: ${query.slice(0, 200)}`);
-                        }
-                        return handleReplicationMessage(data, socket, instances.postgres, mutexes.postgres, connection);
-                    }
-                    // check for no-op queries
-                    if (isNoopQuery(data)) {
-                        if (data[0] === 0x51) {
-                            return buildSetCompleteResponse();
-                        }
-                        else if (data[0] === 0x50) {
-                            return buildParseCompleteResponse();
-                        }
-                    }
-                    // intercept and rewrite queries
-                    data = interceptQuery(data);
-                    // message-level locking on the connection's pglite instance
-                    const { db, mutex } = getDbContext(dbName);
-                    await mutex.acquire();
-                    let result;
-                    try {
-                        result = await db.execProtocolRaw(data, {
-                            throwOnError: false,
-                        });
-                    }
-                    catch (err) {
-                        mutex.release();
-                        throw err;
-                    }
-                    // strip ReadyForQuery from non-Sync/non-SimpleQuery responses
-                    if (data[0] !== 0x53 && data[0] !== 0x51) {
-                        result = stripReadyForQuery(result);
-                    }
-                    mutex.release();
-                    return result;
-                },
+            // perform startup handshake
+            const { params } = await performHandshake(socket, config);
+            dbName = params.database || 'postgres';
+            isReplicationConnection = params.replication === 'database';
+            log.debug.proxy(`connection: db=${dbName} user=${params.user} replication=${params.replication || 'none'}`);
+            const { db } = getDbContext(dbName);
+            await db.waitReady;
+            // clean up pglite transaction state when client disconnects
+            socket.on('close', async () => {
+                const { db: closeDb, mutex: closeMutex } = getDbContext(dbName);
+                await closeMutex.acquire();
+                try {
+                    await closeDb.exec('ROLLBACK');
+                }
+                catch {
+                    // no transaction to rollback
+                }
+                finally {
+                    closeMutex.release();
+                }
             });
+            // enter message processing loop
+            const { db: msgDb, mutex: msgMutex } = getDbContext(dbName);
+            await messageLoop(socket, msgDb, msgMutex, isReplicationConnection, instances.postgres, mutexes.postgres);
         }
         catch (err) {
+            // connection error during handshake or message loop
             if (!socket.destroyed) {
                 socket.destroy();
             }
@@ -392,48 +605,4 @@ export async function startPgProxy(dbInput, config) {
         server.on('error', reject);
     });
 }
-async function handleReplicationMessage(data, socket, db, mutex, connection) {
-    if (data[0] !== 0x51)
-        return undefined;
-    const view = new DataView(data.buffer, data.byteOffset, data.byteLength);
-    const len = view.getInt32(1);
-    const query = new TextDecoder().decode(data.subarray(5, 1 + len - 1)).replace(/\0$/, '');
-    const upper = query.trim().toUpperCase();
-    // check if this is a START_REPLICATION command
-    if (upper.startsWith('START_REPLICATION')) {
-        await connection.detach();
-        const writer = {
-            write(chunk) {
-                if (!socket.destroyed) {
-                    socket.write(chunk);
-                }
-            },
-        };
-        // drain incoming standby status updates
-        socket.on('data', (_chunk) => { });
-        socket.on('close', () => {
-            socket.destroy();
-        });
-        handleStartReplication(query, writer, db, mutex).catch((err) => {
-            log.debug.proxy(`replication stream ended: ${err}`);
-        });
-        return undefined;
-    }
-    // handle replication queries + fallthrough to pglite, all under mutex
-    await mutex.acquire();
-    try {
-        const response = await handleReplicationQuery(query, db);
-        if (response)
-            return response;
-        // apply query rewrites before forwarding
-        data = interceptQuery(data);
-        // fall through to pglite for unrecognized queries
-        return await db.execProtocolRaw(data, {
-            throwOnError: false,
-        });
-    }
-    finally {
-        mutex.release();
-    }
-}
 //# sourceMappingURL=pg-proxy.js.map