Recipe: The PostgreSQL Wire Protocol
A Zig program can use the PostgreSQL wire protocol with no driver. After the startup message, each message has three parts: a one-byte type, a big-endian Int32 length that counts itself, and the payload.
The problem
A Postgres driver is mostly a codec for one TCP protocol. When you see that protocol without a driver, you see the real cost of a query, and the cause of a hang.
This recipe implements a part of both ends, sufficient to complete a startup handshake and run one query. Thus the exact bytes are on the page.
The other database recipes run a mock backend in the same process, and this recipe does too. Thus CI runs the full exchange and compares the output, without a postgres that runs. The framing is the real protocol. Only the authentication and the query engine are stubs.
The plan
- All of the integers are big-endian. Write one
putIntand onegetInt, and use them everywhere. - Send a
StartupMessage. It has no type byte, only a length and a list of parameters that ends in a null. - Read the reply stream of the backend until
ReadyForQuery(Z). Decode each message by its type byte. - Send a
Query(Q). Then readRowDescription(T),DataRow(D),CommandComplete(C) and the nextReadyForQuery.
const std = @import("std");
// PostgreSQL frames all integers big-endian (network order).
fn putInt(comptime T: type, w: *std.Io.Writer, v: T) !void {
var b: [@divExact(@typeInfo(T).int.bits, 8)]u8 = undefined;
std.mem.writeInt(T, &b, v, .big);
try w.writeAll(&b);
}
fn getInt(comptime T: type, r: *std.Io.Reader) !T {
var b: [@divExact(@typeInfo(T).int.bits, 8)]u8 = undefined;
try r.readSliceAll(&b);
return std.mem.readInt(T, &b, .big);
}
// After startup, each message starts with a type byte. Then an Int32 length
// follows. The length counts itself but not the type byte. Then the payload
// follows: that length minus four bytes.
fn sendMsg(w: *std.Io.Writer, tag: u8, payload: []const u8) !void {
try w.writeAll(&.{tag});
try putInt(i32, w, @intCast(payload.len + 4));
try w.writeAll(payload);
}
const Msg = struct { tag: u8, payload: []const u8 };
fn readMsg(r: *std.Io.Reader, buf: []u8) !Msg {
var tb: [1]u8 = undefined;
try r.readSliceAll(&tb);
const len = try getInt(i32, r);
const payload = buf[0..@intCast(len - 4)];
try r.readSliceAll(payload);
return .{ .tag = tb[0], .payload = payload };
}
// A small reader for a message payload.
const Cursor = struct {
b: []const u8,
i: usize = 0,
fn int(self: *Cursor, comptime T: type) T {
const n = @divExact(@typeInfo(T).int.bits, 8);
const v = std.mem.readInt(T, self.b[self.i..][0..n], .big);
self.i += n;
return v;
}
fn cstr(self: *Cursor) []const u8 {
const start = self.i;
while (self.b[self.i] != 0) self.i += 1;
defer self.i += 1; // move past the terminator
return self.b[start..self.i];
}
};
// --- the mock backend ---
fn serve(server: *std.Io.net.Server, io: std.Io) void {
serveImpl(server, io) catch {};
}
fn sendParam(w: *std.Io.Writer, name: []const u8, val: []const u8) !void {
var buf: [128]u8 = undefined;
var fb = std.Io.Writer.fixed(&buf);
try fb.writeAll(name);
try fb.writeAll("\x00");
try fb.writeAll(val);
try fb.writeAll("\x00");
try sendMsg(w, 'S', fb.buffered());
}
fn serveImpl(server: *std.Io.net.Server, io: std.Io) !void {
var conn = try server.accept(io);
defer conn.close(io);
var rbuf: [1024]u8 = undefined;
var wbuf: [1024]u8 = undefined;
var reader = conn.reader(io, &rbuf);
var writer = conn.writer(io, &wbuf);
const r = &reader.interface;
const w = &writer.interface;
// StartupMessage has no type byte: an Int32 length, then the body. The demo
// does no authentication. Thus it reads the body and discards it.
const slen = try getInt(i32, r);
try r.discardAll(@intCast(slen - 4));
var p: [128]u8 = undefined;
var fb = std.Io.Writer.fixed(&p);
try putInt(i32, &fb, 0); // AuthenticationOk
try sendMsg(w, 'R', fb.buffered());
try sendParam(w, "server_version", "16.0");
try sendParam(w, "client_encoding", "UTF8");
fb = std.Io.Writer.fixed(&p);
try putInt(i32, &fb, 4242); // BackendKeyData: pid
try putInt(i32, &fb, 99887766); // secret
try sendMsg(w, 'K', fb.buffered());
try sendMsg(w, 'Z', "I"); // ReadyForQuery: idle
try w.flush();
// One simple Query. The reply has one text column and one row.
var mbuf: [1024]u8 = undefined;
const q = try readMsg(r, &mbuf);
if (q.tag == 'Q') {
fb = std.Io.Writer.fixed(&p);
try putInt(i16, &fb, 1); // RowDescription: one field
try fb.writeAll("greeting\x00");
try putInt(i32, &fb, 0); // table oid
try putInt(i16, &fb, 0); // column number
try putInt(i32, &fb, 25); // type oid: text
try putInt(i16, &fb, -1); // type is variable length
try putInt(i32, &fb, -1); // no type modifier
try putInt(i16, &fb, 0); // text format
try sendMsg(w, 'T', fb.buffered());
fb = std.Io.Writer.fixed(&p);
try putInt(i16, &fb, 1); // DataRow: one column
try putInt(i32, &fb, 3); // its length
try fb.writeAll("zig");
try sendMsg(w, 'D', fb.buffered());
try sendMsg(w, 'C', "SELECT 1\x00"); // CommandComplete
try sendMsg(w, 'Z', "I");
try w.flush();
}
_ = readMsg(r, &mbuf) catch {}; // Terminate, or the client closes the connection
}
// --- the client ---
fn sendStartup(w: *std.Io.Writer, user: []const u8) !void {
var buf: [256]u8 = undefined;
var fb = std.Io.Writer.fixed(&buf);
try putInt(i32, &fb, 196608); // protocol 3.0 == 3 << 16
try fb.writeAll("user\x00");
try fb.writeAll(user);
try fb.writeAll("\x00");
try fb.writeAll("\x00"); // end of parameters
const body = fb.buffered();
try putInt(i32, w, @intCast(4 + body.len)); // length counts itself
try w.writeAll(body);
try w.flush();
}
fn statusName(b: u8) []const u8 {
return switch (b) {
'I' => "idle",
'T' => "in transaction",
'E' => "failed transaction",
else => "?",
};
}
pub fn main(init: std.process.Init) !void {
const io = init.io;
var buf: [256]u8 = undefined;
var file_writer = std.Io.File.stdout().writerStreaming(io, &buf);
const out = &file_writer.interface;
const any_port = try std.Io.net.IpAddress.parse("127.0.0.1", 0);
var server = try any_port.listen(io, .{});
defer server.deinit(io);
const thread = try std.Thread.spawn(.{}, serve, .{ &server, io });
var stream = try server.socket.address.connect(io, .{ .mode = .stream });
defer stream.close(io);
var rbuf: [1024]u8 = undefined;
var wbuf: [1024]u8 = undefined;
var reader = stream.reader(io, &rbuf);
var writer = stream.writer(io, &wbuf);
const r = &reader.interface;
const w = &writer.interface;
var mbuf: [1024]u8 = undefined;
try out.writeAll("> startup (user=zig)\n");
try sendStartup(w, "zig");
while (true) {
const m = try readMsg(r, &mbuf);
var c = Cursor{ .b = m.payload };
switch (m.tag) {
'R' => if (c.int(i32) == 0) try out.writeAll("auth: ok\n"),
'S' => try out.print("param: {s}={s}\n", .{ c.cstr(), c.cstr() }),
'K' => try out.print("backend: pid={d}\n", .{c.int(i32)}),
'Z' => {
try out.print("ready: {s}\n", .{statusName(m.payload[0])});
break;
},
else => {},
}
}
try out.writeAll("> query: SELECT 'zig' AS greeting\n");
try sendMsg(w, 'Q', "SELECT 'zig' AS greeting\x00");
try w.flush();
while (true) {
const m = try readMsg(r, &mbuf);
var c = Cursor{ .b = m.payload };
switch (m.tag) {
'T' => {
const n = c.int(i16);
try out.print("fields: {d} ({s})\n", .{ n, c.cstr() });
},
'D' => {
_ = c.int(i16); // column count
const len = c.int(i32);
try out.print("row: {s}\n", .{c.b[c.i..][0..@intCast(len)]});
},
'C' => try out.print("done: {s}\n", .{c.cstr()}),
'Z' => {
try out.print("ready: {s}\n", .{statusName(m.payload[0])});
break;
},
else => {},
}
}
try sendMsg(w, 'X', ""); // Terminate
try w.flush();
thread.join();
try out.flush();
}Run it yourself:zig run postgres-wire.zig
Two message shapes
After the first message, each Postgres message has the same frame:
- a one-byte type
- an
Int32length that counts itself but not the type byte length - 4bytes of payload
sendMsg and readMsg encode exactly that frame. Thus the other code works only with payloads.
The startup message is the one exception. It is older than the type-byte convention. Thus it is only a length and a body. Because of this difference, sendStartup and readMsg are separate.
The handshake is a stream, not a call
When you send the startup message, you do not get one reply. The backend replies with a sequence of messages:
AuthenticationOk.- One
ParameterStatusfor each server setting (server_version,client_encoding, and usually about twelve more). BackendKeyData, with the process id and the secret that a client uses to cancel queries.ReadyForQuery.
The client loops until it sees Z. That message tells the client that the connection is ready to use. A real driver caches the parameters and the cancel key here.
Read a result
A query replies with the same stream shape:
RowDescriptiongives the names of the columns and their types.DataRowcontains one row as column values, each with its length before it. A length of -1 means SQLNULL.CommandCompletereports the tag.ReadyForQuerycloses the batch.
The Cursor helper moves through a payload with an index. It reads integers and null-terminated strings. These messages need only that parse.
Flush, or the program hangs
Each send here has a flush after it, because the first version of this recipe hung without the flush. The socket writer has a buffer. Thus a message that you “sent” stays in your buffer until you flush.
If both sides then block on a read, neither side can continue. This is a common deadlock in wire protocols. It is easy to cause this deadlock when a flush does not immediately follow a send.
Variations
- Real authentication: the backend can reply to the startup with
AuthenticationCleartextPasswordor with an MD5 or SCRAM challenge, in place ofAuthenticationOk. The client replies with aPasswordMessage, and the loop continues. - Prepared statements: the extended-query protocol replaces
QwithParse,BindandExecute. With this protocol, parameters stay out of the SQL text. - A local database: to keep data without a server or a protocol, SQLite from Zig links the C library directly.