rustlavel_db/postgres/
connection.rs

1//! A single PostgreSQL connection.
2
3use super::auth::{self, Scram};
4use super::protocol::{
5    self, Authentication, Backend, Buffer, Field, ServerError, TransactionStatus,
6};
7use super::types;
8use crate::config::DatabaseConfig;
9use crate::random;
10use crate::row::{Columns, Row};
11use crate::value::Value;
12use crate::driver::{BoxFuture, Driver, DriverConnection, QueryResult};
13use rustlavel_core::events::Event;
14use rustlavel_core::{Error, Result};
15use std::sync::Arc;
16use std::sync::atomic::{AtomicBool, Ordering};
17use std::time::Instant;
18use tokio::net::TcpStream;
19
20/// Whether query parameters are included in the `db.query` event.
21///
22/// Bindings are what make a slow-query log useful, but they are also where a
23/// password or a token ends up when one is being written. On by default because
24/// the instrumentation bus only has subscribers in development; turned off for
25/// production by the application at boot.
26static LOG_BINDINGS: AtomicBool = AtomicBool::new(true);
27
28pub fn set_log_bindings(enabled: bool) {
29    LOG_BINDINGS.store(enabled, Ordering::Relaxed);
30}
31
32pub fn log_bindings() -> bool {
33    LOG_BINDINGS.load(Ordering::Relaxed)
34}
35
36pub struct Connection {
37    stream: crate::tls::DbStream,
38    /// Bytes read from the socket but not yet consumed as a message.
39    buffer: Vec<u8>,
40    config: DatabaseConfig,
41    process_id: i32,
42    secret: i32,
43    status: TransactionStatus,
44    /// Set when the connection is known to be unusable, so the pool discards it.
45    broken: bool,
46}
47
48impl Connection {
49    /// Open a connection and complete the startup handshake.
50    pub async fn connect(config: &DatabaseConfig) -> Result<Connection> {
51        let address = format!("{}:{}", config.host, config.port);
52
53        let stream = tokio::time::timeout(config.connect_timeout, TcpStream::connect(&address))
54            .await
55            .map_err(|_| {
56                Error::msg(format!(
57                    "timed out connecting to {address}. Is PostgreSQL running and reachable?"
58                ))
59            })?
60            .map_err(|e| {
61                Error::msg(format!(
62                    "cannot connect to {}: {e}",
63                    config.redacted_url()
64                ))
65            })?;
66
67        let _ = stream.set_nodelay(true);
68
69        let mut connection = Connection {
70            stream: crate::tls::DbStream::Plain(stream),
71            buffer: Vec::with_capacity(8 * 1024),
72            config: config.clone(),
73            process_id: 0,
74            secret: 0,
75            status: TransactionStatus::Idle,
76            broken: false,
77        };
78
79        // Before the startup packet, because the startup packet carries the
80        // user name and the database, and after it the password follows. All
81        // of that has to be inside the tunnel, not in front of it.
82        connection.negotiate_tls().await?;
83        connection.startup().await?;
84        Ok(connection)
85    }
86
87    pub fn is_broken(&self) -> bool {
88        self.broken
89    }
90
91    /// True while a transaction is open, so the pool never hands back a
92    /// connection that would leak one.
93    pub fn in_transaction(&self) -> bool {
94        self.status != TransactionStatus::Idle
95    }
96
97    /// Whether this connection is encrypted.
98    pub fn is_encrypted(&self) -> bool {
99        self.stream.is_encrypted()
100    }
101
102    /// Ask PostgreSQL to encrypt the connection, per the SSLRequest exchange in
103    /// the protocol's message formats.
104    ///
105    /// It is eight bytes — a length and a magic number where a normal packet
106    /// would carry its version — and the reply is a single byte outside the
107    /// usual message framing: `S` for yes, `N` for no. That is the whole
108    /// negotiation, and it happens before anything identifying has been sent.
109    async fn negotiate_tls(&mut self) -> Result<()> {
110        let mode = self.config.tls_mode;
111        if !mode.wants_tls() {
112            return Ok(());
113        }
114
115        let mut request = Vec::with_capacity(8);
116        request.extend_from_slice(&8i32.to_be_bytes());
117        request.extend_from_slice(&protocol::SSL_REQUEST_CODE.to_be_bytes());
118        if let Err(e) = self.stream.write_all(&request).await {
119            self.broken = true;
120            return Err(Error::Io(e));
121        }
122        if let Err(e) = self.stream.flush().await {
123            self.broken = true;
124            return Err(Error::Io(e));
125        }
126
127        let mut answer = [0u8; 1];
128        match self.stream.read(&mut answer).await {
129            Ok(1) => {}
130            Ok(_) => {
131                self.broken = true;
132                return Err(Error::msg(
133                    "the server closed the connection when asked about TLS. A PostgreSQL older                      than 8.0 does not understand SSLRequest; set sslmode=disable if that is                      what this is.",
134                ));
135            }
136            Err(e) => {
137                self.broken = true;
138                return Err(Error::Io(e));
139            }
140        }
141
142        match answer[0] {
143            b'S' => {
144                let plain = self.stream.take_plain()?;
145                let encrypted =
146                    crate::tls::upgrade(plain, &self.config.host, &self.config).await?;
147                self.stream = crate::tls::DbStream::Tls(Box::new(encrypted));
148                Ok(())
149            }
150            // The server is willing to talk, but not privately. Under `prefer`
151            // that is accepted; under anything stronger it is the whole point.
152            b'N' if !mode.demands_tls() => Ok(()),
153            b'N' => {
154                self.broken = true;
155                Err(Error::msg(format!(
156                    "sslmode is `{mode}` but this PostgreSQL refused to encrypt the connection.                      Either the server was built without SSL support or `ssl` is off in                      postgresql.conf. Turn it on, or set sslmode=prefer to accept a connection                      in clear text."
157                )))
158            }
159            b'E' => {
160                self.broken = true;
161                Err(Error::msg(
162                    "the server reported an error in response to SSLRequest. This usually means                      it is not actually PostgreSQL.",
163                ))
164            }
165            other => {
166                self.broken = true;
167                Err(Error::msg(format!(
168                    "the server answered SSLRequest with {other:?}, which is not `S` or `N`.                      Whatever is on this port, it is not speaking the PostgreSQL protocol."
169                )))
170            }
171        }
172    }
173
174    async fn startup(&mut self) -> Result<()> {
175        let mut buffer = Buffer::new();
176        buffer.startup(&[
177            ("user", self.config.user.as_str()),
178            ("database", self.config.database.as_str()),
179            ("application_name", self.config.application_name.as_str()),
180            // Timestamps and dates come back in a predictable shape.
181            ("DateStyle", "ISO, MDY"),
182            ("client_encoding", "UTF8"),
183        ]);
184        self.write(buffer).await?;
185
186        let mut scram: Option<Scram> = None;
187
188        loop {
189            match self.read_message().await? {
190                Backend::Authentication(Authentication::Ok) => continue,
191                Backend::Authentication(Authentication::CleartextPassword) => {
192                    let mut buffer = Buffer::new();
193                    buffer.password(&self.config.password);
194                    self.write(buffer).await?;
195                }
196                Backend::Authentication(Authentication::Md5Password { salt }) => {
197                    let digest =
198                        auth::md5_password(&self.config.user, &self.config.password, &salt);
199                    let mut buffer = Buffer::new();
200                    buffer.password(&digest);
201                    self.write(buffer).await?;
202                }
203                Backend::Authentication(Authentication::Sasl { mechanisms }) => {
204                    if !mechanisms.iter().any(|m| m == Scram::MECHANISM) {
205                        return Err(Error::msg(format!(
206                            "the server offers only {mechanisms:?}; this driver implements {}",
207                            Scram::MECHANISM
208                        )));
209                    }
210                    let exchange = Scram::new(&self.config.password, random::nonce(24));
211                    let mut buffer = Buffer::new();
212                    buffer.sasl_initial(Scram::MECHANISM, &exchange.client_first());
213                    self.write(buffer).await?;
214                    scram = Some(exchange);
215                }
216                Backend::Authentication(Authentication::SaslContinue { data }) => {
217                    let exchange = scram
218                        .as_mut()
219                        .ok_or_else(|| Error::Protocol("SASL continue before SASL start".into()))?;
220                    let response = exchange.client_final(&data)?;
221                    let mut buffer = Buffer::new();
222                    buffer.sasl_response(&response);
223                    self.write(buffer).await?;
224                }
225                Backend::Authentication(Authentication::SaslFinal { data }) => {
226                    scram
227                        .as_ref()
228                        .ok_or_else(|| Error::Protocol("SASL final before SASL start".into()))?
229                        .verify(&data)?;
230                }
231                Backend::Authentication(Authentication::Unsupported(code)) => {
232                    return Err(Error::msg(format!(
233                        "the server requested authentication method {code}, which this driver does not implement"
234                    )));
235                }
236                Backend::BackendKeyData { process_id, secret } => {
237                    self.process_id = process_id;
238                    self.secret = secret;
239                }
240                Backend::ParameterStatus { .. } | Backend::Notice(_) => continue,
241                Backend::ReadyForQuery(status) => {
242                    self.status = status;
243                    return Ok(());
244                }
245                Backend::Error(error) => {
246                    self.broken = true;
247                    return Err(authentication_error(error, &self.config));
248                }
249                other => {
250                    return Err(Error::Protocol(format!(
251                        "unexpected message during startup: {other:?}"
252                    )));
253                }
254            }
255        }
256    }
257
258    /// Run a statement with no parameters through the simple query protocol.
259    ///
260    /// Used for DDL and for statements that must run as one unit, such as
261    /// `begin`/`commit`.
262    pub async fn simple_query(&mut self, sql: &str) -> Result<QueryResult> {
263        let started = Instant::now();
264        let mut buffer = Buffer::new();
265        buffer.query(sql);
266        self.write(buffer).await?;
267
268        let result = self.collect(sql).await;
269        self.record(sql, &[], started, &result);
270        result
271    }
272
273    /// Run a parameterised statement through the extended query protocol.
274    ///
275    /// Parameters never enter the SQL text, so a value cannot change the shape
276    /// of the statement — this is what makes SQL injection structurally
277    /// impossible rather than a matter of remembering to escape.
278    pub async fn query(&mut self, sql: &str, params: &[Value]) -> Result<QueryResult> {
279        if params.is_empty() {
280            // Still uses the extended protocol, so a single statement per call
281            // is enforced either way.
282            return self.extended(sql, params).await;
283        }
284        self.extended(sql, params).await
285    }
286
287    async fn extended(&mut self, sql: &str, params: &[Value]) -> Result<QueryResult> {
288        let started = Instant::now();
289        let encoded: Vec<Option<String>> = params.iter().map(Value::to_sql_text).collect();
290
291        let mut buffer = Buffer::new();
292        buffer.parse("", sql);
293        buffer.bind("", "", &encoded);
294        buffer.describe_portal("");
295        buffer.execute("", 0);
296        buffer.sync();
297        self.write(buffer).await?;
298
299        let result = self.collect(sql).await;
300        self.record(sql, params, started, &result);
301        result
302    }
303
304    /// Read messages until `ReadyForQuery`, gathering rows on the way.
305    ///
306    /// The loop always runs to `ReadyForQuery` even after an error, otherwise
307    /// the next query would read this one's leftovers.
308    async fn collect(&mut self, sql: &str) -> Result<QueryResult> {
309        let mut columns: Columns = Arc::new(Vec::new());
310        let mut fields: Vec<Field> = Vec::new();
311        let mut result = QueryResult::default();
312        let mut failure: Option<ServerError> = None;
313
314        loop {
315            match self.read_message().await? {
316                Backend::RowDescription(described) => {
317                    columns = Arc::new(described.iter().map(|f| f.name.clone()).collect());
318                    fields = described;
319                }
320                Backend::DataRow(raw) => {
321                    let values = raw
322                        .iter()
323                        .enumerate()
324                        .map(|(index, bytes)| {
325                            let oid = fields.get(index).map_or(types::TEXT, |f| f.type_oid);
326                            types::decode(oid, bytes.as_deref())
327                        })
328                        .collect();
329                    result.rows.push(Row::new(Arc::clone(&columns), values));
330                }
331                Backend::CommandComplete(tag) => result.affected = affected_rows(&tag),
332                Backend::Error(error) => failure = Some(error),
333                Backend::ReadyForQuery(status) => {
334                    self.status = status;
335                    break;
336                }
337                Backend::EmptyQueryResponse
338                | Backend::ParseComplete
339                | Backend::BindComplete
340                | Backend::CloseComplete
341                | Backend::NoData
342                | Backend::PortalSuspended
343                | Backend::Notice(_)
344                | Backend::ParameterStatus { .. }
345                | Backend::NotificationResponse { .. }
346                | Backend::BackendKeyData { .. }
347                | Backend::Other(_)
348                | Backend::Authentication(_) => {}
349            }
350        }
351
352        match failure {
353            Some(error) => Err(error.into_error(Some(sql))),
354            None => Ok(result),
355        }
356    }
357
358    /// Publish the query on the event bus for Telescope and slow-query logging.
359    fn record(&self, sql: &str, params: &[Value], started: Instant, result: &Result<QueryResult>) {
360        let elapsed = started.elapsed();
361
362        if rustlavel_core::events::has_subscribers() {
363            let bindings = if log_bindings() {
364                params.iter().map(Value::to_display).collect::<Vec<_>>().join(", ")
365            } else {
366                format!("{} value(s) hidden", params.len())
367            };
368            Event::new("db.query")
369                .with("sql", sql)
370                .with("bindings", bindings)
371                .with("rows", result.as_ref().map(|r| r.rows.len()).unwrap_or(0))
372                .with("ok", result.is_ok())
373                .took(elapsed)
374                .dispatch();
375        }
376
377        rustlavel_core::debug!("db: {sql} ({:.1}ms)", elapsed.as_secs_f64() * 1000.0);
378    }
379
380    async fn write(&mut self, buffer: Buffer) -> Result<()> {
381        let bytes = buffer.into_bytes();
382        if let Err(e) = self.stream.write_all(&bytes).await {
383            self.broken = true;
384            return Err(Error::Io(e));
385        }
386        if let Err(e) = self.stream.flush().await {
387            self.broken = true;
388            return Err(Error::Io(e));
389        }
390        Ok(())
391    }
392
393    /// Read exactly one backend message.
394    async fn read_message(&mut self) -> Result<Backend> {
395        // A message is a type byte plus a length that includes itself.
396        self.fill_to(5).await?;
397        let tag = self.buffer[0];
398        let length = i32::from_be_bytes(self.buffer[1..5].try_into().expect("4 bytes")) as usize;
399
400        if length < 4 {
401            self.broken = true;
402            return Err(Error::Protocol("message length is impossibly small".into()));
403        }
404
405        let total = length + 1;
406        self.fill_to(total).await?;
407        let body = self.buffer[5..total].to_vec();
408        self.buffer.drain(..total);
409
410        Backend::parse(tag, &body)
411    }
412
413    /// Read from the socket until the buffer holds at least `wanted` bytes.
414    async fn fill_to(&mut self, wanted: usize) -> Result<()> {
415        while self.buffer.len() < wanted {
416            let mut chunk = [0u8; 8192];
417            let read = match self.stream.read(&mut chunk).await {
418                Ok(read) => read,
419                Err(e) => {
420                    self.broken = true;
421                    return Err(Error::Io(e));
422                }
423            };
424            if read == 0 {
425                self.broken = true;
426                return Err(Error::Protocol(
427                    "the database closed the connection unexpectedly".into(),
428                ));
429            }
430            self.buffer.extend_from_slice(&chunk[..read]);
431        }
432        Ok(())
433    }
434
435    /// Ask the server to close the session politely.
436    pub async fn close(mut self) {
437        let mut buffer = Buffer::new();
438        buffer.terminate();
439        let _ = self.write(buffer).await;
440        let _ = self.stream.shutdown().await;
441    }
442}
443
444/// Opens PostgreSQL connections.
445///
446/// The reference implementation of [`Driver`]: everything above the driver line
447/// is written once, and this is what that line looks like from below.
448pub struct PostgresDriver {
449    config: DatabaseConfig,
450    dialect: Arc<dyn crate::dialect::Dialect>,
451}
452
453impl PostgresDriver {
454    pub fn new(config: DatabaseConfig) -> Self {
455        PostgresDriver { config, dialect: Arc::new(crate::dialect::Postgres) }
456    }
457
458    pub fn config(&self) -> &DatabaseConfig {
459        &self.config
460    }
461}
462
463impl Driver for PostgresDriver {
464    fn generation(&self) -> u64 {
465        self.config.generation()
466    }
467
468    fn dialect(&self) -> Arc<dyn crate::dialect::Dialect> {
469        Arc::clone(&self.dialect)
470    }
471
472    fn connect(&self) -> BoxFuture<'_, Result<Box<dyn DriverConnection>>> {
473        Box::pin(async move {
474            let connection = Connection::connect(&self.config.resolved()).await?;
475            Ok(Box::new(connection) as Box<dyn DriverConnection>)
476        })
477    }
478
479    fn describe(&self) -> String {
480        self.config.redacted_url()
481    }
482
483    fn max_connections(&self) -> usize {
484        self.config.max_connections
485    }
486}
487
488impl DriverConnection for Connection {
489    fn query<'a>(
490        &'a mut self,
491        sql: &'a str,
492        params: &'a [Value],
493    ) -> BoxFuture<'a, Result<QueryResult>> {
494        Box::pin(Connection::query(self, sql, params))
495    }
496
497    fn simple_query<'a>(&'a mut self, sql: &'a str) -> BoxFuture<'a, Result<QueryResult>> {
498        Box::pin(Connection::simple_query(self, sql))
499    }
500
501    fn is_broken(&self) -> bool {
502        Connection::is_broken(self)
503    }
504
505    fn in_transaction(&self) -> bool {
506        Connection::in_transaction(self)
507    }
508
509    fn close(self: Box<Self>) -> BoxFuture<'static, ()> {
510        Box::pin(async move { Connection::close(*self).await })
511    }
512}
513
514/// `INSERT 0 3`, `UPDATE 2`, `SELECT 5` → the trailing count.
515fn affected_rows(tag: &str) -> u64 {
516    tag.split_whitespace().next_back().and_then(|n| n.parse().ok()).unwrap_or(0)
517}
518
519/// Turn an authentication failure into something the developer can act on.
520fn authentication_error(error: ServerError, config: &DatabaseConfig) -> Error {
521    let base = error.clone().into_error(None);
522
523    let advice = match error.code.as_str() {
524        "28P01" => Some(format!(
525            "The password for `{}` was rejected. Check DATABASE_URL in your .env.",
526            config.user
527        )),
528        "3D000" => Some(format!(
529            "Database `{}` does not exist. Create it, or point DATABASE_URL at an existing one.",
530            config.database
531        )),
532        "28000" => Some(
533            "The server rejected this role or host. Check pg_hba.conf allows this connection."
534                .to_string(),
535        ),
536        _ => None,
537    };
538
539    match advice {
540        Some(advice) => Error::msg(format!("{base}\n  {advice}")),
541        None => base,
542    }
543}
544
545#[cfg(test)]
546mod tests {
547    use super::*;
548
549    #[test]
550    fn reads_the_row_count_from_a_command_tag() {
551        assert_eq!(affected_rows("INSERT 0 3"), 3);
552        assert_eq!(affected_rows("UPDATE 2"), 2);
553        assert_eq!(affected_rows("DELETE 0"), 0);
554        assert_eq!(affected_rows("CREATE TABLE"), 0);
555    }
556
557    #[test]
558    fn bindings_can_be_kept_out_of_the_event_stream() {
559        // Restored immediately: the flag is process-wide.
560        assert!(log_bindings());
561        set_log_bindings(false);
562        assert!(!log_bindings());
563        set_log_bindings(true);
564    }
565
566    #[test]
567    fn a_wrong_password_explains_where_to_look() {
568        let error = ServerError {
569            code: "28P01".into(),
570            message: "password authentication failed".into(),
571            ..ServerError::default()
572        };
573        let config = DatabaseConfig { user: "ada".into(), ..DatabaseConfig::default() };
574
575        let rendered = authentication_error(error, &config).to_string();
576        assert!(rendered.contains("DATABASE_URL"));
577        assert!(rendered.contains("`ada`"));
578    }
579
580    #[tokio::test]
581    async fn connecting_to_a_closed_port_names_the_server() {
582        let config = DatabaseConfig { port: 1, ..DatabaseConfig::default() };
583        let error = match Connection::connect(&config).await {
584            Err(error) => error.to_string(),
585            Ok(_) => panic!("nothing should be listening on port 1"),
586        };
587
588        assert!(error.contains("127.0.0.1:1"));
589        // The password must never appear, even in a connection error.
590        assert!(!error.contains("***@") || !error.contains("hunter"));
591    }
592}
rustlavel_db/postgres/connection.rs

rustlavel_db/postgres/
connection.rs
