Postgres wire protocol deep dive

Building a Postgres proxy means speaking the wire protocol fluently. Here's what we learned implementing it in Rust.

Why a proxy?

We considered other approaches:

• Postgres extension. Maximum performance, but harder to deploy and version. Users would need to install our extension on their database.
• Application-level interception. Modify the ORM or driver. But then you're tied to specific languages and frameworks.
• Proxy. Sits between client and server, speaks native Postgres. Works with any language, any ORM, any tool that can connect to Postgres.

The proxy approach won because it's transparent. Point your connection string at Astapa instead of your database, and everything just works.

Protocol basics

The Postgres wire protocol is a binary, message-based protocol. Every message has:

1. A single-byte type tag (except startup messages)
2. A 4-byte length (including itself)
3. Payload

[type: 1 byte][length: 4 bytes][payload: length-4 bytes]

There are two directions:

• Frontend messages — client to server (Query, Parse, Bind, Execute, etc.)
• Backend messages — server to client (RowDescription, DataRow, ReadyForQuery, etc.)

Startup sequence

Connection startup is the trickiest part. Here's the flow:

Client                     Proxy                      Postgres
   |                         |                           |
   |-- SSLRequest ---------->|                           |
   |<--------- 'N' (no SSL) -|                           |
   |                         |                           |
   |-- StartupMessage ------>|                           |
   |   (user, database,      |-- StartupMessage -------->|
   |    astapa_session)     |   (user, database)        |
   |                         |                           |
   |                         |<-- AuthenticationOk ------|
   |<-- AuthenticationOk ----|                           |
   |                         |                           |
   |                         |<-- ParameterStatus[] -----|
   |<-- ParameterStatus[] ---|                           |
   |                         |                           |
   |                         |<-- ReadyForQuery ---------|
   |<-- ReadyForQuery -------|                           |

Key points:

• We reject SSL (for now) and let clients fall back to plaintext
• We strip astapa_session from parameters before forwarding to Postgres
• We relay all authentication messages unchanged

Simple vs Extended query protocol

Postgres has two query modes:

Simple query protocol — Send a Query message with raw SQL, get results back.

Frontend: Query("SELECT * FROM users")
Backend: RowDescription, DataRow[], CommandComplete, ReadyForQuery

Extended query protocol — Separate parse, bind, and execute steps. Used by prepared statements.

Frontend: Parse(name="s1", query="SELECT * FROM users WHERE id = $1")
Backend: ParseComplete
Frontend: Bind(portal="", statement="s1", params=[42])
Backend: BindComplete
Frontend: Execute(portal="", max_rows=0)
Backend: DataRow[], CommandComplete
Frontend: Sync
Backend: ReadyForQuery

We handle both. The rewriting happens at Parse time for extended protocol, at Query time for simple protocol.

Message parsing in Rust

Here's how we parse frontend messages:

pub fn parse_frontend_message(buf: &mut BytesMut) -> Result<Option<FrontendMessage>> {
    if buf.len() < 5 {
        return Ok(None); // Need more data
    }

    let tag = buf[0];
    let len = i32::from_be_bytes([buf[1], buf[2], buf[3], buf[4]]) as usize;

    if buf.len() < 1 + len {
        return Ok(None); // Need more data
    }

    let payload = buf.split_to(1 + len);

    match tag {
        b'Q' => parse_query(&payload[5..]),
        b'P' => parse_parse(&payload[5..]),
        b'B' => parse_bind(&payload[5..]),
        b'E' => parse_execute(&payload[5..]),
        b'S' => Ok(FrontendMessage::Sync),
        b'X' => Ok(FrontendMessage::Terminate),
        _ => Ok(FrontendMessage::Other { tag, payload }),
    }
}

Startup messages are special — they don't have a type tag. We detect them by checking for the SSL request magic number (80877103) or protocol version (196608 = 3.0).

Query rewriting

When we intercept a Query or Parse message, we:

1. Extract the SQL string
2. Parse it with sqlparser-rs
3. Walk the AST, rewriting table references
4. Serialize back to SQL
5. Forward the modified query

pub fn rewrite(&self, sql: &str) -> Result<(String, QueryType)> {
    let dialect = PostgreSqlDialect {};
    let ast = Parser::parse_sql(&dialect, sql)?;

    let mut rewritten = ast.clone();
    let query_type = self.analyze_and_rewrite(&mut rewritten)?;

    let new_sql = rewritten
        .iter()
        .map(|s| s.to_string())
        .collect::<Vec<_>>()
        .join("; ");

    Ok((new_sql, query_type))
}

Edge cases we hit

JSON operators. Postgres uses -> and ->> for JSON access. These look like operators to the parser, not special syntax.

Array syntax. ANY(ARRAY[1,2,3]) and array_col[1] require careful handling.

Dollar quoting. Function bodies use $$..$$ strings. The parser handles this, but it's easy to mess up escaping.

Identifier quoting. We always quote identifiers in output to handle reserved words and special characters: "session_abc123"."_view_users".

Performance: buffered I/O

The naive approach — read one message, process it, write it — is slow. We use buffered I/O:

let mut client_buffer = BytesMut::with_capacity(16384);
let mut backend_buffer = BytesMut::with_capacity(16384);

loop {
    tokio::select! {
        result = client.read_buf(&mut client_buffer) => {
            // Process all complete messages in buffer
            while let Some(msg) = parse_frontend_message(&mut client_buffer)? {
                // Rewrite and forward
            }
        }
        result = backend.read_buf(&mut backend_buffer) => {
            // Relay all complete messages to client
            while let Some(msg) = parse_backend_message(&mut backend_buffer)? {
                client.write_all(&encode(msg)).await?;
            }
        }
    }
}

The 16KB buffers mean we can batch multiple small messages into single syscalls.

Connection pooling compatibility

PgBouncer and other poolers multiplex connections. A single backend connection serves multiple clients. This breaks session-based isolation unless you handle it.

Our solution: the SET astapa_session command.

SET astapa_session TO 'abc-123';  -- Switch to session
-- ... queries run in that session ...
RESET astapa_session;              -- Back to passthrough

When we see this command, we swap the rewriter context. The client can switch sessions without reconnecting.

What we'd do differently

1. Start with extended protocol only. Simple query protocol is... simple, but it's a separate code path to maintain.

2. Use tokio-postgres for parsing. We wrote our own parser initially. tokio-postgres's protocol handling is battle-tested.

3. Add SSL earlier. We currently reject SSL requests. Adding TLS termination is on the roadmap.

Resources

If you're building your own Postgres proxy:

• Postgres Protocol Documentation — The official spec
• pgwire — Rust library for Postgres wire protocol
• sqlparser-rs — SQL parser in Rust

Want to see our implementation? It's on GitHub.