ironwork_rt/sql/postgres/
wire.rs

1//! PostgreSQL's frontend/backend protocol, version 3: startup and authentication, the simple query
2//! cycle, and the extended one (Parse, Describe, Bind, Execute, Sync) with text values.
3
4use super::scram::{Scram, nonce};
5use std::io::{BufReader, Read, Write};
6use std::net::TcpStream;
7use std::path::{Path, PathBuf};
8
9/// A duplex byte stream: a socket, or TLS over one.
10pub trait Stream: Read + Write + Send {}
11
12impl<T: Read + Write + Send> Stream for T {}
13
14/// Wraps a connected socket in TLS, verifying the server's certificate chain against `roots` (a PEM
15/// file) or built-in roots, and its name as `host`. ironwork's own build has no implementation, so
16/// that it keeps no dependencies; the build in `tls/` supplies one.
17pub trait Tls: Send + Sync {
18    fn wrap(&self, socket: TcpStream, host: &str, roots: Option<&Path>) -> std::io::Result<Box<dyn Stream>>;
19}
20
21/// The URL's `sslmode`: the two of libpq's modes that ironwork offers.
22#[derive(Clone, Copy, Debug, PartialEq, Eq)]
23pub enum SslMode {
24    Disable,
25    /// TLS, with the server's certificate chain and name verified.
26    VerifyFull,
27}
28
29/// Where and as whom to connect.
30#[derive(Clone, Debug, PartialEq, Eq)]
31pub struct Target {
32    pub host: String,
33    pub port: u16,
34    pub user: String,
35    pub password: Option<String>,
36    pub database: String,
37    /// None where the URL does not say.
38    pub ssl: Option<SslMode>,
39    pub root_cert: Option<PathBuf>,
40}
41
42impl Target {
43    /// `postgres://user[:password]@host[:port]/database[?option=value&...]`, with the options
44    /// `host` (a directory, for a Unix socket), `sslmode` and `sslrootcert`. The password may come
45    /// from PGPASSWORD instead.
46    pub fn parse(url: &str) -> Result<Self, String> {
47        let rest = url.strip_prefix("postgres://").or_else(|| url.strip_prefix("postgresql://")).ok_or("a database URL starts postgres://")?;
48        let (rest, query) = rest.split_once('?').unwrap_or((rest, ""));
49        let (authority, database) = rest.split_once('/').unwrap_or((rest, ""));
50        let (userinfo, hostport) = authority.rsplit_once('@').map_or((None, authority), |(u, h)| (Some(u), h));
51        let (user, password) = match userinfo.map(|u| u.split_once(':').map_or((u, None), |(u, p)| (u, Some(p)))) {
52            Some((u, p)) => (Some(decode(u)?), p.map(decode).transpose()?),
53            None => (None, None),
54        };
55        let (mut host, port) = match hostport.rsplit_once(':') {
56            Some((h, p)) => (h.to_owned(), p.parse().map_err(|_| format!("{p} is not a port"))?),
57            None => (hostport.to_owned(), 5432),
58        };
59        let (mut ssl, mut root_cert) = (None, None);
60        for pair in query.split('&').filter(|p| !p.is_empty()) {
61            match pair.split_once('=') {
62                Some(("host", h)) => host = decode(h)?,
63                Some(("sslmode", "disable")) => ssl = Some(SslMode::Disable),
64                Some(("sslmode", "verify-full")) => ssl = Some(SslMode::VerifyFull),
65                Some(("sslmode", other)) => {
66                    return Err(format!("sslmode={other} is not offered: disable, or verify-full, which checks the server's certificate and name"));
67                }
68                Some(("sslrootcert", path)) => root_cert = Some(PathBuf::from(decode(path)?)),
69                _ => return Err(format!("{pair} is not a URL option ironwork reads")),
70            }
71        }
72        let user = match user {
73            Some(u) => u,
74            None => std::env::var("USER").map_err(|_| "the URL names no user, and USER is not set")?,
75        };
76        Ok(Self {
77            host: if host.is_empty() { "localhost".into() } else { decode(&host)? },
78            port,
79            database: if database.is_empty() { user.clone() } else { decode(database)? },
80            password: password.or_else(|| std::env::var("PGPASSWORD").ok()),
81            user,
82            ssl,
83            root_cert,
84        })
85    }
86}
87
88fn decode(text: &str) -> Result<String, String> {
89    let bytes = text.as_bytes();
90    let (mut out, mut i) = (Vec::new(), 0);
91    while i < bytes.len() {
92        if bytes[i] == b'%' {
93            let hex = text.get(i + 1..i + 3).and_then(|h| u8::from_str_radix(h, 16).ok()).ok_or("the URL has a bad % escape")?;
94            out.push(hex);
95            i += 3;
96        } else {
97            out.push(bytes[i]);
98            i += 1;
99        }
100    }
101    String::from_utf8(out).map_err(|_| "the URL is not UTF-8 once decoded".to_owned())
102}
103
104/// Why a request failed: the server refused it, naming an SQLSTATE, or the connection broke.
105#[derive(Debug, PartialEq, Eq)]
106pub enum Failure {
107    Refused { state: String, message: String },
108    Broken(String),
109}
110
111impl From<std::io::Error> for Failure {
112    fn from(e: std::io::Error) -> Self {
113        Failure::Broken(format!("the connection to PostgreSQL failed: {e}"))
114    }
115}
116
117/// What Execute returned: each row's columns as text, and the command tag.
118#[derive(Debug, Default)]
119pub struct Executed {
120    pub rows: Vec<Vec<Option<String>>>,
121    pub tag: String,
122}
123
124/// A prepared statement's parameter and column types.
125#[derive(Clone, Debug, Default)]
126pub struct Described {
127    pub parameters: Vec<u32>,
128    pub columns: Vec<u32>,
129}
130
131pub struct Connection {
132    stream: BufReader<Box<dyn Stream>>,
133    /// Messages waiting for the next flush.
134    out: Vec<u8>,
135    pub server_version: String,
136    pub encrypted: bool,
137    /// ReadyForQuery's transaction status: `I` idle, `T` in a transaction, `E` in a failed one.
138    pub status: u8,
139}
140
141struct Body<'a> {
142    bytes: &'a [u8],
143    at: usize,
144}
145
146impl<'a> Body<'a> {
147    fn take(&mut self, n: usize) -> Result<&'a [u8], Failure> {
148        let got = self.bytes.get(self.at..self.at + n).ok_or_else(|| Failure::Broken("PostgreSQL sent a short message".into()))?;
149        self.at += n;
150        Ok(got)
151    }
152    fn i16(&mut self) -> Result<i16, Failure> {
153        self.take(2).map(|b| i16::from_be_bytes([b[0], b[1]]))
154    }
155    fn i32(&mut self) -> Result<i32, Failure> {
156        self.take(4).map(|b| i32::from_be_bytes([b[0], b[1], b[2], b[3]]))
157    }
158    fn cstr(&mut self) -> Result<String, Failure> {
159        let end = self.bytes[self.at..].iter().position(|&b| b == 0).ok_or_else(|| Failure::Broken("PostgreSQL sent an unterminated string".into()))?;
160        let s = String::from_utf8_lossy(&self.bytes[self.at..self.at + end]).into_owned();
161        self.at += end + 1;
162        Ok(s)
163    }
164    fn rest(&mut self) -> &'a [u8] {
165        let r = &self.bytes[self.at..];
166        self.at = self.bytes.len();
167        r
168    }
169}
170
171fn cstr(out: &mut Vec<u8>, s: &str) {
172    out.extend_from_slice(s.as_bytes());
173    out.push(0);
174}
175
176/// An ErrorResponse's SQLSTATE and message.
177fn refusal(body: &[u8]) -> Failure {
178    let mut b = Body { bytes: body, at: 0 };
179    let (mut state, mut message) = (String::new(), String::new());
180    while let Ok(field) = b.take(1) {
181        if field[0] == 0 {
182            break;
183        }
184        let Ok(value) = b.cstr() else { break };
185        match field[0] {
186            b'C' => state = value,
187            b'M' => message = value,
188            _ => {}
189        }
190    }
191    Failure::Refused { state, message }
192}
193
194impl Connection {
195    /// Connects, over TLS where `sslmode` asks for it. Without an sslmode, a TCP connection uses TLS
196    /// when this build has it, and a Unix socket never does.
197    pub fn open(target: &Target, tls: Option<&dyn Tls>) -> Result<Self, String> {
198        let at = format!("PostgreSQL at {}:{}", target.host, target.port);
199        let broken = |e: std::io::Error| format!("cannot reach {at}: {e}");
200        let socket = target.host.starts_with('/');
201        let mode = match (target.ssl, tls) {
202            (Some(mode), _) => mode,
203            (None, Some(_)) if !socket => SslMode::VerifyFull,
204            (None, _) => SslMode::Disable,
205        };
206        let stream: Box<dyn Stream> = match (mode, socket, tls) {
207            (SslMode::Disable, true, _) => unix_socket(target).map_err(broken)?,
208            (SslMode::VerifyFull, true, _) => return Err("sslmode=verify-full is for TCP; a Unix socket needs no TLS".into()),
209            (SslMode::VerifyFull, false, None) => {
210                return Err("sslmode=verify-full needs TLS, which this build of ironwork leaves out to keep it free of dependencies; build tls/ for it".into());
211            }
212            (SslMode::Disable, false, _) => Box::new(tcp(target).map_err(broken)?),
213            (SslMode::VerifyFull, false, Some(tls)) => {
214                let mut socket = tcp(target).map_err(broken)?;
215                socket.write_all(&[0, 0, 0, 8, 0x04, 0xD2, 0x16, 0x2F]).map_err(broken)?;
216                let mut answer = [0u8];
217                socket.read_exact(&mut answer).map_err(broken)?;
218                if answer[0] != b'S' {
219                    return Err(format!("{at} does not offer TLS; give sslmode=disable to connect without it"));
220                }
221                tls.wrap(socket, &target.host, target.root_cert.as_deref()).map_err(|e| format!("TLS with {at} failed: {e}"))?
222            }
223        };
224        let mut conn = Self { stream: BufReader::new(stream), out: Vec::new(), server_version: String::new(), encrypted: mode == SslMode::VerifyFull, status: b'I' };
225        conn.start(target).map_err(|f| match f {
226            Failure::Refused { state, message } => format!("PostgreSQL refused the connection ({state}): {message}"),
227            Failure::Broken(m) => m,
228        })?;
229        Ok(conn)
230    }
231
232    fn start(&mut self, target: &Target) -> Result<(), Failure> {
233        let mut body = 196_608i32.to_be_bytes().to_vec();
234        let params = [
235            ("user", target.user.as_str()),
236            ("database", &target.database),
237            ("client_encoding", "UTF8"),
238            ("DateStyle", "ISO"),
239            ("TimeZone", "UTC"),
240            ("extra_float_digits", "3"),
241            ("application_name", "ironwork"),
242        ];
243        for (name, value) in params {
244            cstr(&mut body, name);
245            cstr(&mut body, value);
246        }
247        body.push(0);
248        self.out.extend_from_slice(&(body.len() as i32 + 4).to_be_bytes());
249        self.out.extend_from_slice(&body);
250        self.flush()?;
251        let password = || target.password.clone().ok_or_else(|| Failure::Broken("PostgreSQL asks for a password: give one in the URL or PGPASSWORD".into()));
252        let mut scram = None;
253        loop {
254            let (kind, body) = self.receive()?;
255            let mut b = Body { bytes: &body, at: 0 };
256            match kind {
257                b'R' => match b.i32()? {
258                    0 => {}
259                    3 => {
260                        let mut reply = Vec::new();
261                        cstr(&mut reply, &password()?);
262                        self.send(b'p', &reply)?;
263                        self.flush()?;
264                    }
265                    10 => {
266                        let mechanisms: Vec<String> = std::iter::from_fn(|| b.cstr().ok().filter(|m| !m.is_empty())).collect();
267                        if !mechanisms.iter().any(|m| m == "SCRAM-SHA-256") {
268                            return Err(Failure::Broken(format!("PostgreSQL offers SASL {mechanisms:?}, and ironwork speaks SCRAM-SHA-256")));
269                        }
270                        let mut exchange = Scram::new(&password()?, nonce());
271                        let first = exchange.client_first("");
272                        let mut reply = Vec::new();
273                        cstr(&mut reply, "SCRAM-SHA-256");
274                        reply.extend_from_slice(&(first.len() as i32).to_be_bytes());
275                        reply.extend_from_slice(first.as_bytes());
276                        self.send(b'p', &reply)?;
277                        self.flush()?;
278                        scram = Some(exchange);
279                    }
280                    11 => {
281                        let exchange = scram.as_mut().ok_or_else(|| Failure::Broken("PostgreSQL continued a SASL exchange that had not begun".into()))?;
282                        let reply = exchange.client_final(&String::from_utf8_lossy(b.rest())).map_err(Failure::Broken)?;
283                        self.send(b'p', reply.as_bytes())?;
284                        self.flush()?;
285                    }
286                    12 => {
287                        let exchange = scram.as_ref().ok_or_else(|| Failure::Broken("PostgreSQL ended a SASL exchange that had not begun".into()))?;
288                        exchange.verify(&String::from_utf8_lossy(b.rest())).map_err(Failure::Broken)?;
289                    }
290                    5 => return Err(Failure::Broken("PostgreSQL asks for MD5 authentication; ironwork speaks SCRAM-SHA-256".into())),
291                    other => return Err(Failure::Broken(format!("PostgreSQL asks for authentication method {other}, which ironwork does not speak"))),
292                },
293                b'Z' => {
294                    self.status = body.first().copied().unwrap_or(b'I');
295                    return Ok(());
296                }
297                b'E' => return Err(refusal(&body)),
298                _ => self.note(kind, &body),
299            }
300        }
301    }
302
303    fn send(&mut self, kind: u8, body: &[u8]) -> Result<(), Failure> {
304        self.out.push(kind);
305        self.out.extend_from_slice(&(body.len() as i32 + 4).to_be_bytes());
306        self.out.extend_from_slice(body);
307        Ok(())
308    }
309
310    fn flush(&mut self) -> Result<(), Failure> {
311        let stream = self.stream.get_mut();
312        stream.write_all(&self.out)?;
313        stream.flush()?;
314        self.out.clear();
315        Ok(())
316    }
317
318    fn receive(&mut self) -> Result<(u8, Vec<u8>), Failure> {
319        let mut head = [0u8; 5];
320        self.stream.read_exact(&mut head)?;
321        let len = i32::from_be_bytes([head[1], head[2], head[3], head[4]]);
322        let mut body = vec![0u8; (len.max(4) - 4) as usize];
323        self.stream.read_exact(&mut body)?;
324        Ok((head[0], body))
325    }
326
327    /// Messages that may arrive at any time: ParameterStatus, notices and notifications.
328    fn note(&mut self, kind: u8, body: &[u8]) {
329        if kind == b'S' {
330            let mut b = Body { bytes: body, at: 0 };
331            if let (Ok(name), Ok(value)) = (b.cstr(), b.cstr())
332                && name == "server_version"
333            {
334                self.server_version = value;
335            }
336        }
337    }
338
339    /// Reads to ReadyForQuery, handing each message to `each`, and fails with the first refusal.
340    fn until_ready(&mut self, mut each: impl FnMut(u8, &[u8]) -> Result<(), Failure>) -> Result<(), Failure> {
341        let mut refused = None;
342        loop {
343            let (kind, body) = self.receive()?;
344            match kind {
345                b'Z' => {
346                    self.status = body.first().copied().unwrap_or(b'I');
347                    return refused.map_or(Ok(()), Err);
348                }
349                b'E' => refused = refused.or(Some(refusal(&body))),
350                b'S' | b'N' | b'A' => self.note(kind, &body),
351                _ if refused.is_none() => each(kind, &body)?,
352                _ => {}
353            }
354        }
355    }
356
357    pub fn simple(&mut self, sql: &str) -> Result<(), Failure> {
358        let mut body = Vec::new();
359        cstr(&mut body, sql);
360        self.send(b'Q', &body)?;
361        self.flush()?;
362        self.until_ready(|_, _| Ok(()))
363    }
364
365    pub fn prepare(&mut self, name: &str, sql: &str) -> Result<Described, Failure> {
366        let mut parse = Vec::new();
367        cstr(&mut parse, name);
368        cstr(&mut parse, sql);
369        parse.extend_from_slice(&0i16.to_be_bytes());
370        self.send(b'P', &parse)?;
371        let mut describe = vec![b'S'];
372        cstr(&mut describe, name);
373        self.send(b'D', &describe)?;
374        self.send(b'S', &[])?;
375        self.flush()?;
376        let mut described = Described::default();
377        self.until_ready(|kind, body| {
378            let mut b = Body { bytes: body, at: 0 };
379            match kind {
380                b't' => {
381                    let n = b.i16()?;
382                    described.parameters = (0..n).map(|_| b.i32().map(|o| o as u32)).collect::<Result<_, _>>()?;
383                }
384                b'T' => {
385                    let n = b.i16()?;
386                    for _ in 0..n {
387                        b.cstr()?;
388                        b.take(6)?;
389                        described.columns.push(b.i32()? as u32);
390                        b.take(8)?;
391                    }
392                }
393                _ => {}
394            }
395            Ok(())
396        })?;
397        Ok(described)
398    }
399
400    /// Binds text parameters (None is NULL) to a prepared statement and executes it, returning at
401    /// most `max_rows` rows (0 for all).
402    pub fn execute(&mut self, name: &str, parameters: &[Option<String>], max_rows: i32) -> Result<Executed, Failure> {
403        let mut bind = Vec::new();
404        cstr(&mut bind, "");
405        cstr(&mut bind, name);
406        bind.extend_from_slice(&0i16.to_be_bytes());
407        bind.extend_from_slice(&(parameters.len() as i16).to_be_bytes());
408        for p in parameters {
409            match p {
410                None => bind.extend_from_slice(&(-1i32).to_be_bytes()),
411                Some(text) => {
412                    bind.extend_from_slice(&(text.len() as i32).to_be_bytes());
413                    bind.extend_from_slice(text.as_bytes());
414                }
415            }
416        }
417        bind.extend_from_slice(&0i16.to_be_bytes());
418        self.send(b'B', &bind)?;
419        let mut execute = Vec::new();
420        cstr(&mut execute, "");
421        execute.extend_from_slice(&max_rows.to_be_bytes());
422        self.send(b'E', &execute)?;
423        self.send(b'S', &[])?;
424        self.flush()?;
425        let mut executed = Executed::default();
426        self.until_ready(|kind, body| {
427            let mut b = Body { bytes: body, at: 0 };
428            match kind {
429                b'D' => {
430                    let n = b.i16()?;
431                    let mut row = Vec::new();
432                    for _ in 0..n {
433                        let len = b.i32()?;
434                        row.push(if len < 0 { None } else { Some(String::from_utf8_lossy(b.take(len as usize)?).into_owned()) });
435                    }
436                    executed.rows.push(row);
437                }
438                b'C' => executed.tag = b.cstr()?,
439                _ => {}
440            }
441            Ok(())
442        })?;
443        Ok(executed)
444    }
445}
446
447fn tcp(target: &Target) -> std::io::Result<TcpStream> {
448    let socket = TcpStream::connect((target.host.as_str(), target.port))?;
449    socket.set_nodelay(true)?;
450    Ok(socket)
451}
452
453#[cfg(unix)]
454fn unix_socket(target: &Target) -> std::io::Result<Box<dyn Stream>> {
455    Ok(Box::new(std::os::unix::net::UnixStream::connect(format!("{}/.s.PGSQL.{}", target.host, target.port))?))
456}
457
458#[cfg(not(unix))]
459fn unix_socket(_: &Target) -> std::io::Result<Box<dyn Stream>> {
460    Err(std::io::Error::new(std::io::ErrorKind::Unsupported, "Unix sockets need a Unix system"))
461}
462
463#[cfg(test)]
464mod tests {
465    use super::*;
466
467    #[test]
468    fn urls() {
469        let t = Target::parse("postgres://ironwork:p%40ss@db.example:6543/payroll").unwrap();
470        let expected = Target { host: "db.example".into(), port: 6543, user: "ironwork".into(), password: Some("p@ss".into()), database: "payroll".into(), ssl: None, root_cert: None };
471        assert_eq!(t, expected);
472        let t = Target::parse("postgres://me@db/payroll?sslmode=verify-full&sslrootcert=/etc/ca.pem").unwrap();
473        assert_eq!((t.ssl, t.root_cert.as_deref()), (Some(SslMode::VerifyFull), Some(Path::new("/etc/ca.pem"))));
474        assert!(Target::parse("postgres://me@db/payroll?sslmode=require").unwrap_err().contains("verify-full"));
475        let t = Target::parse("postgres://me@/payroll?host=/var/run/postgresql").unwrap();
476        assert_eq!((t.host.as_str(), t.port, t.database.as_str()), ("/var/run/postgresql", 5432, "payroll"));
477        assert!(Target::parse("mysql://x").is_err());
478        assert!(Target::parse("postgres://me@h:port/db").is_err());
479    }
480}
ironwork_rt/sql/postgres/wire.rs

ironwork_rt/sql/postgres/
wire.rs
