pgsync_postgres/replication/
conn.rs

1//! Minimal PostgreSQL wire-protocol client for the replication connection.
2//!
3//! The published `tokio-postgres` cannot open a `replication=database` session, so this
4//! module implements just enough of the protocol: startup, authentication
5//! (trust / cleartext / MD5 / SCRAM-SHA-256), simple queries, and the CopyBoth stream used
6//! by `START_REPLICATION`. Message encoding/decoding is delegated to `postgres-protocol`.
7
8use crate::errors::{from_io, kind_for_sqlstate};
9use crate::tls;
10use crate::url::{ConnectionInfo, TlsPolicy};
11use bytes::{Buf, Bytes, BytesMut};
12use fallible_iterator::FallibleIterator;
13use pgsync_core::config::SafetyConfig;
14use pgsync_core::{Error, ErrorKind, Result};
15use postgres_protocol::authentication::sasl::{ChannelBinding, ScramSha256};
16use postgres_protocol::message::backend::Message;
17use postgres_protocol::message::frontend;
18use std::pin::Pin;
19use std::time::Duration;
20use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
21use tokio::net::TcpStream;
22use tracing::debug;
23
24type Io = Pin<Box<dyn AsyncReadWrite + Send>>;
25
26pub trait AsyncReadWrite: AsyncRead + AsyncWrite {}
27impl<T: AsyncRead + AsyncWrite> AsyncReadWrite for T {}
28
29const COPY_BOTH_RESPONSE_TAG: u8 = b'W';
30
31/// A decoded simple-query result set.
32#[derive(Debug, Default, Clone)]
33pub struct SimpleResult {
34    pub columns: Vec<String>,
35    pub rows: Vec<Vec<Option<String>>>,
36}
37
38impl SimpleResult {
39    pub fn get(&self, row: usize, col: &str) -> Option<&str> {
40        let idx = self.columns.iter().position(|c| c == col)?;
41        self.rows.get(row)?.get(idx)?.as_deref()
42    }
43}
44
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46enum Mode {
47    Idle,
48    CopyBoth,
49    Closed,
50}
51
52pub struct ReplicationConnection {
53    io: Io,
54    read_buf: BytesMut,
55    write_buf: BytesMut,
56    mode: Mode,
57    label: String,
58    pub server_version: Option<String>,
59    /// Largest backend message accepted (`resources.max_message_bytes`). Enforced before
60    /// the message is buffered, so memory never silently exceeds the budget.
61    max_message_bytes: usize,
62}
63
64enum Backend {
65    CopyBothResponse,
66    Msg(Message),
67}
68
69impl ReplicationConnection {
70    /// Open a logical-replication connection to `info.dbname`.
71    pub async fn connect(
72        info: &ConnectionInfo,
73        safety: &SafetyConfig,
74        timeout: Duration,
75    ) -> Result<Self> {
76        let label = info.redacted();
77        let fut = Self::connect_inner(info, safety, &label);
78        tokio::time::timeout(timeout, fut).await.map_err(|_| {
79            Error::transient(format!(
80                "{label}: replication connection timed out after {timeout:?}"
81            ))
82        })?
83    }
84
85    async fn connect_inner(
86        info: &ConnectionInfo,
87        safety: &SafetyConfig,
88        label: &str,
89    ) -> Result<Self> {
90        let addr = format!("{}:{}", info.host, info.port);
91        let mut tcp = TcpStream::connect(&addr)
92            .await
93            .map_err(|e| from_io(&format!("{label}: connect"), e))?;
94        tcp.set_nodelay(true).ok();
95
96        let io: Io = if info.tls.uses_tls() {
97            let mut buf = BytesMut::new();
98            frontend::ssl_request(&mut buf);
99            tcp.write_all(&buf).await.map_err(|e| from_io(label, e))?;
100            let mut answer = [0u8; 1];
101            tcp.read_exact(&mut answer)
102                .await
103                .map_err(|e| from_io(label, e))?;
104            match (answer[0], info.tls) {
105                (b'S', _) => {
106                    let config = tls::client_config(info.tls, safety, info.tls_ca_file.as_deref())?;
107                    let connector = tokio_rustls::TlsConnector::from(config);
108                    let name = tls::server_name(&info.host)?;
109                    let stream = connector.connect(name, tcp).await.map_err(|e| {
110                        Error::config(format!(
111                            "{label}: TLS handshake failed: {}",
112                            crate::errors::error_chain(&e)
113                        ))
114                    })?;
115                    Box::pin(stream)
116                }
117                (b'N', TlsPolicy::Required { .. }) => {
118                    return Err(Error::config(format!(
119                        "{label}: server does not support TLS but TLS is required"
120                    )));
121                }
122                (b'N', _) => {
123                    debug!(
124                        target = label,
125                        "server declined TLS; continuing in plain text (loopback)"
126                    );
127                    Box::pin(tcp)
128                }
129                (other, _) => {
130                    return Err(Error::internal(format!(
131                        "{label}: unexpected SSL negotiation byte {other:#x}"
132                    )));
133                }
134            }
135        } else {
136            Box::pin(tcp)
137        };
138
139        let mut conn = ReplicationConnection {
140            io,
141            read_buf: BytesMut::with_capacity(64 * 1024),
142            write_buf: BytesMut::with_capacity(8 * 1024),
143            mode: Mode::Idle,
144            label: label.to_owned(),
145            server_version: None,
146            max_message_bytes: usize::MAX,
147        };
148        conn.startup(info).await?;
149        Ok(conn)
150    }
151
152    async fn startup(&mut self, info: &ConnectionInfo) -> Result<()> {
153        let params = [
154            ("user", info.user.as_str()),
155            ("database", info.dbname.as_str()),
156            ("replication", "database"),
157            ("application_name", "pgsync"),
158            ("client_encoding", "UTF8"),
159        ];
160        frontend::startup_message(params, &mut self.write_buf)
161            .map_err(|e| from_io(&self.label, e))?;
162        self.flush().await?;
163
164        let password = info.password().map(|p| p.to_vec());
165        loop {
166            match self.read_backend().await? {
167                Backend::Msg(Message::AuthenticationOk) => break,
168                Backend::Msg(Message::AuthenticationCleartextPassword) => {
169                    let pw = password.as_deref().ok_or_else(Self::need_password)?;
170                    frontend::password_message(pw, &mut self.write_buf)
171                        .map_err(|e| from_io(&self.label, e))?;
172                    self.flush().await?;
173                }
174                Backend::Msg(Message::AuthenticationMd5Password(body)) => {
175                    let pw = password.as_deref().ok_or_else(Self::need_password)?;
176                    let hashed = postgres_protocol::authentication::md5_hash(
177                        info.user.as_bytes(),
178                        pw,
179                        body.salt(),
180                    );
181                    frontend::password_message(hashed.as_bytes(), &mut self.write_buf)
182                        .map_err(|e| from_io(&self.label, e))?;
183                    self.flush().await?;
184                }
185                Backend::Msg(Message::AuthenticationSasl(body)) => {
186                    let pw = password.as_deref().ok_or_else(Self::need_password)?;
187                    let mut mechanisms = body.mechanisms();
188                    let mut has_scram = false;
189                    while let Some(m) = mechanisms.next().map_err(|e| from_io(&self.label, e))? {
190                        if m == "SCRAM-SHA-256" {
191                            has_scram = true;
192                        }
193                    }
194                    if !has_scram {
195                        return Err(Error::config(format!(
196                            "{}: server offers no supported SASL mechanism",
197                            self.label
198                        )));
199                    }
200                    // Channel binding is declared unsupported ("n,,"); the server accepts this
201                    // unless it is configured to require it.
202                    let mut scram = ScramSha256::new(pw, ChannelBinding::unsupported());
203                    frontend::sasl_initial_response(
204                        "SCRAM-SHA-256",
205                        scram.message(),
206                        &mut self.write_buf,
207                    )
208                    .map_err(|e| from_io(&self.label, e))?;
209                    self.flush().await?;
210                    match self.read_backend().await? {
211                        Backend::Msg(Message::AuthenticationSaslContinue(body)) => {
212                            scram.update(body.data()).map_err(|e| {
213                                Error::with_source(
214                                    ErrorKind::Config,
215                                    format!("{}: SCRAM failed", self.label),
216                                    e,
217                                )
218                            })?;
219                        }
220                        other => return Err(self.unexpected(other, "AuthenticationSASLContinue")),
221                    }
222                    frontend::sasl_response(scram.message(), &mut self.write_buf)
223                        .map_err(|e| from_io(&self.label, e))?;
224                    self.flush().await?;
225                    match self.read_backend().await? {
226                        Backend::Msg(Message::AuthenticationSaslFinal(body)) => {
227                            scram.finish(body.data()).map_err(|e| {
228                                Error::with_source(
229                                    ErrorKind::Config,
230                                    format!("{}: SCRAM verification failed", self.label),
231                                    e,
232                                )
233                            })?;
234                        }
235                        other => return Err(self.unexpected(other, "AuthenticationSASLFinal")),
236                    }
237                }
238                Backend::Msg(Message::ErrorResponse(body)) => {
239                    return Err(self.db_error(body.fields()));
240                }
241                other => return Err(self.unexpected(other, "authentication")),
242            }
243        }
244
245        loop {
246            match self.read_backend().await? {
247                Backend::Msg(Message::ParameterStatus(body)) => {
248                    if body.name().map_err(|e| from_io(&self.label, e))? == "server_version" {
249                        self.server_version = Some(
250                            body.value()
251                                .map_err(|e| from_io(&self.label, e))?
252                                .to_owned(),
253                        );
254                    }
255                }
256                Backend::Msg(Message::BackendKeyData(_)) => {}
257                Backend::Msg(Message::NoticeResponse(_)) => {}
258                Backend::Msg(Message::ReadyForQuery(_)) => return Ok(()),
259                Backend::Msg(Message::ErrorResponse(body)) => {
260                    return Err(self.db_error(body.fields()));
261                }
262                other => return Err(self.unexpected(other, "ReadyForQuery")),
263            }
264        }
265    }
266
267    /// Refuse backend messages larger than `bytes` (see `resources.max_message_bytes`).
268    pub fn set_max_message_bytes(&mut self, bytes: usize) {
269        self.max_message_bytes = bytes.max(1024);
270    }
271
272    fn need_password() -> Error {
273        Error::config("server requested a password but the connection URL has none")
274    }
275
276    /// Run a simple-query command (replication commands or SQL) and collect the result.
277    pub async fn simple_query(&mut self, sql: &str) -> Result<SimpleResult> {
278        self.ensure_idle()?;
279        frontend::query(sql, &mut self.write_buf).map_err(|e| from_io(&self.label, e))?;
280        self.flush().await?;
281        let mut result = SimpleResult::default();
282        let mut error: Option<Error> = None;
283        loop {
284            match self.read_backend().await? {
285                Backend::Msg(Message::RowDescription(body)) => {
286                    let mut fields = body.fields();
287                    result.columns.clear();
288                    while let Some(f) = fields.next().map_err(|e| from_io(&self.label, e))? {
289                        result.columns.push(f.name().to_owned());
290                    }
291                }
292                Backend::Msg(Message::DataRow(body)) => {
293                    let mut row = Vec::with_capacity(result.columns.len());
294                    let buf = body.buffer();
295                    let mut ranges = body.ranges();
296                    while let Some(r) = ranges.next().map_err(|e| from_io(&self.label, e))? {
297                        row.push(r.map(|r| String::from_utf8_lossy(&buf[r]).into_owned()));
298                    }
299                    result.rows.push(row);
300                }
301                Backend::Msg(Message::CommandComplete(_))
302                | Backend::Msg(Message::EmptyQueryResponse) => {}
303                Backend::Msg(Message::NoticeResponse(body)) => {
304                    self.log_notice(body.fields());
305                }
306                Backend::Msg(Message::ErrorResponse(body)) => {
307                    error = Some(self.db_error(body.fields()));
308                }
309                Backend::Msg(Message::ReadyForQuery(_)) => break,
310                Backend::CopyBothResponse => {
311                    return Err(Error::internal(format!(
312                        "{}: unexpected CopyBothResponse to {sql:?}",
313                        self.label
314                    )));
315                }
316                other => return Err(self.unexpected(other, "simple query response")),
317            }
318        }
319        match error {
320            Some(e) => Err(e),
321            None => Ok(result),
322        }
323    }
324
325    /// Issue `START_REPLICATION ...` and switch to CopyBoth mode.
326    pub async fn start_replication(&mut self, command: &str) -> Result<()> {
327        self.ensure_idle()?;
328        frontend::query(command, &mut self.write_buf).map_err(|e| from_io(&self.label, e))?;
329        self.flush().await?;
330        loop {
331            match self.read_backend().await? {
332                Backend::CopyBothResponse => {
333                    self.mode = Mode::CopyBoth;
334                    return Ok(());
335                }
336                Backend::Msg(Message::NoticeResponse(body)) => self.log_notice(body.fields()),
337                Backend::Msg(Message::ErrorResponse(body)) => {
338                    let err = self.db_error(body.fields());
339                    // Consume the trailing ReadyForQuery so the connection stays usable.
340                    loop {
341                        if let Backend::Msg(Message::ReadyForQuery(_)) = self.read_backend().await?
342                        {
343                            break;
344                        }
345                    }
346                    return Err(err);
347                }
348                other => return Err(self.unexpected(other, "CopyBothResponse")),
349            }
350        }
351    }
352
353    /// Next CopyData payload from the server. Cancel-safe: a partially received frame
354    /// stays buffered. Returns an error when the stream ends.
355    pub async fn next_copy_data(&mut self) -> Result<Bytes> {
356        if self.mode != Mode::CopyBoth {
357            return Err(Error::internal(format!(
358                "{}: not in CopyBoth mode",
359                self.label
360            )));
361        }
362        loop {
363            match self.read_backend().await? {
364                Backend::Msg(Message::CopyData(body)) => {
365                    return Ok(body.into_bytes());
366                }
367                // A shutting-down logical walsender ends the stream with CommandComplete
368                // (no CopyDone); a physical one sends CopyDone first. Both mean "reconnect".
369                Backend::Msg(Message::CopyDone)
370                | Backend::Msg(Message::CommandComplete(_))
371                | Backend::Msg(Message::ReadyForQuery(_)) => {
372                    self.mode = Mode::Closed;
373                    return Err(Error::transient(format!(
374                        "{}: server ended the replication stream (shutdown or timeline change)",
375                        self.label
376                    )));
377                }
378                Backend::Msg(Message::NoticeResponse(body)) => self.log_notice(body.fields()),
379                Backend::Msg(Message::ErrorResponse(body)) => {
380                    self.mode = Mode::Closed;
381                    return Err(self.db_error(body.fields()));
382                }
383                other => return Err(self.unexpected(other, "CopyData")),
384            }
385        }
386    }
387
388    /// Send a CopyData frame (standby status update) to the server.
389    pub async fn send_copy_data(&mut self, payload: &[u8]) -> Result<()> {
390        if self.mode != Mode::CopyBoth {
391            return Err(Error::internal(format!(
392                "{}: not in CopyBoth mode",
393                self.label
394            )));
395        }
396        let msg = frontend::CopyData::new(payload).map_err(|e| from_io(&self.label, e))?;
397        msg.write(&mut self.write_buf);
398        self.flush().await
399    }
400
401    fn ensure_idle(&self) -> Result<()> {
402        match self.mode {
403            Mode::Idle => Ok(()),
404            Mode::CopyBoth => Err(Error::internal(format!(
405                "{}: connection is streaming",
406                self.label
407            ))),
408            Mode::Closed => Err(Error::transient(format!(
409                "{}: connection closed",
410                self.label
411            ))),
412        }
413    }
414
415    /// Cancel-safe: bytes not yet written stay in `write_buf` and are sent by the next call.
416    async fn flush(&mut self) -> Result<()> {
417        while !self.write_buf.is_empty() {
418            let n = self
419                .io
420                .write(&self.write_buf)
421                .await
422                .map_err(|e| from_io(&self.label, e))?;
423            if n == 0 {
424                return Err(Error::transient(format!(
425                    "{}: write returned 0 bytes",
426                    self.label
427                )));
428            }
429            self.write_buf.advance(n);
430        }
431        self.io.flush().await.map_err(|e| from_io(&self.label, e))
432    }
433
434    async fn read_backend(&mut self) -> Result<Backend> {
435        loop {
436            if self.read_buf.len() >= 5 {
437                let tag = self.read_buf[0];
438                let len = u32::from_be_bytes([
439                    self.read_buf[1],
440                    self.read_buf[2],
441                    self.read_buf[3],
442                    self.read_buf[4],
443                ]) as usize;
444                if len < 4 {
445                    return Err(Error::internal(format!(
446                        "{}: invalid message length {len}",
447                        self.label
448                    )));
449                }
450                if len + 1 > self.max_message_bytes {
451                    self.mode = Mode::Closed;
452                    return Err(Error::resource(format!(
453                        "{}: server sent a {} byte message, above resources.max_message_bytes ({}); \
454                         nothing was acknowledged. Raise the limit (and memory_budget_mb if needed) \
455                         and `run` to apply the transaction",
456                        self.label,
457                        len + 1,
458                        self.max_message_bytes
459                    )));
460                }
461                if self.read_buf.len() > len {
462                    if tag == COPY_BOTH_RESPONSE_TAG {
463                        self.read_buf.advance(len + 1);
464                        return Ok(Backend::CopyBothResponse);
465                    }
466                    let msg = Message::parse(&mut self.read_buf).map_err(|e| {
467                        Error::with_source(
468                            ErrorKind::Internal,
469                            format!("{}: protocol error", self.label),
470                            e,
471                        )
472                    })?;
473                    if let Some(m) = msg {
474                        return Ok(Backend::Msg(m));
475                    }
476                    return Err(Error::internal(format!(
477                        "{}: parser consumed nothing",
478                        self.label
479                    )));
480                }
481                self.read_buf.reserve(len + 1 - self.read_buf.len());
482            } else {
483                self.read_buf.reserve(64 * 1024);
484            }
485            let n = self
486                .io
487                .read_buf(&mut self.read_buf)
488                .await
489                .map_err(|e| from_io(&self.label, e))?;
490            if n == 0 {
491                self.mode = Mode::Closed;
492                return Err(Error::transient(format!(
493                    "{}: connection closed by server",
494                    self.label
495                )));
496            }
497        }
498    }
499
500    fn db_error(&self, fields: postgres_protocol::message::backend::ErrorFields<'_>) -> Error {
501        let (code, message, detail) = collect_error_fields(fields);
502        let kind = stream_error_kind(&code, &message);
503        let mut msg = format!("{}: {message} (SQLSTATE {code})", self.label);
504        if let Some(d) = detail {
505            msg.push_str(&format!(": {d}"));
506        }
507        Error::new(kind, msg)
508    }
509
510    fn log_notice(&self, fields: postgres_protocol::message::backend::ErrorFields<'_>) {
511        let (code, message, _) = collect_error_fields(fields);
512        debug!(target = %self.label, code, "server notice: {message}");
513    }
514
515    fn unexpected(&self, got: Backend, expected: &str) -> Error {
516        let what = match got {
517            Backend::CopyBothResponse => "CopyBothResponse".to_owned(),
518            Backend::Msg(m) => message_name(&m).to_owned(),
519        };
520        Error::internal(format!("{}: expected {expected}, got {what}", self.label))
521    }
522
523    pub async fn close(mut self) {
524        let mut buf = BytesMut::new();
525        frontend::terminate(&mut buf);
526        let _ = self.io.write_all(&buf).await;
527        let _ = self.io.shutdown().await;
528    }
529}
530
531/// Errors on the replication connection that mean the slot can no longer deliver every
532/// transaction: the slot was invalidated (`max_slot_wal_keep_size`) or its WAL is gone.
533/// Both need a resynchronisation, not a retry.
534pub fn stream_error_kind(code: &str, message: &str) -> ErrorKind {
535    match code {
536        // object_not_in_prerequisite_state: "can no longer get changes from replication
537        // slot ... This slot has been invalidated because ..."
538        "55000" if message.contains("replication slot") || message.contains("invalidated") => {
539            ErrorKind::ResyncRequired
540        }
541        // undefined_file: "requested WAL segment ... has already been removed"
542        "58P01" => ErrorKind::ResyncRequired,
543        _ => kind_for_sqlstate(code),
544    }
545}
546
547fn collect_error_fields(
548    mut fields: postgres_protocol::message::backend::ErrorFields<'_>,
549) -> (String, String, Option<String>) {
550    let mut code = String::from("XX000");
551    let mut message = String::from("unknown error");
552    let mut detail = None;
553    while let Ok(Some(f)) = fields.next() {
554        match f.type_() {
555            b'C' => code = String::from_utf8_lossy(f.value_bytes()).into_owned(),
556            b'M' => message = String::from_utf8_lossy(f.value_bytes()).into_owned(),
557            b'D' => detail = Some(String::from_utf8_lossy(f.value_bytes()).into_owned()),
558            _ => {}
559        }
560    }
561    (code, message, detail)
562}
563
564fn message_name(m: &Message) -> &'static str {
565    match m {
566        Message::AuthenticationOk => "AuthenticationOk",
567        Message::AuthenticationCleartextPassword => "AuthenticationCleartextPassword",
568        Message::AuthenticationMd5Password(_) => "AuthenticationMD5Password",
569        Message::AuthenticationSasl(_) => "AuthenticationSASL",
570        Message::AuthenticationSaslContinue(_) => "AuthenticationSASLContinue",
571        Message::AuthenticationSaslFinal(_) => "AuthenticationSASLFinal",
572        Message::BackendKeyData(_) => "BackendKeyData",
573        Message::CommandComplete(_) => "CommandComplete",
574        Message::CopyData(_) => "CopyData",
575        Message::CopyDone => "CopyDone",
576        Message::CopyInResponse(_) => "CopyInResponse",
577        Message::CopyOutResponse(_) => "CopyOutResponse",
578        Message::DataRow(_) => "DataRow",
579        Message::EmptyQueryResponse => "EmptyQueryResponse",
580        Message::ErrorResponse(_) => "ErrorResponse",
581        Message::NoticeResponse(_) => "NoticeResponse",
582        Message::ParameterStatus(_) => "ParameterStatus",
583        Message::ReadyForQuery(_) => "ReadyForQuery",
584        Message::RowDescription(_) => "RowDescription",
585        _ => "other",
586    }
587}
588
589#[cfg(test)]
590mod tests {
591    use super::*;
592
593    #[test]
594    fn invalidated_slot_and_removed_wal_need_resync() {
595        assert_eq!(
596            stream_error_kind(
597                "55000",
598                "can no longer get changes from replication slot \"s\""
599            ),
600            ErrorKind::ResyncRequired
601        );
602        assert_eq!(
603            stream_error_kind(
604                "58P01",
605                "requested WAL segment 000000010000000000000003 has already been removed"
606            ),
607            ErrorKind::ResyncRequired
608        );
609        assert_eq!(
610            stream_error_kind("55006", "replication slot \"s\" is active for PID 1"),
611            ErrorKind::Transient
612        );
613        assert_eq!(
614            stream_error_kind("55000", "something else"),
615            ErrorKind::Internal
616        );
617        assert_eq!(
618            stream_error_kind("57P01", "terminating connection"),
619            ErrorKind::Transient
620        );
621    }
622}
pgsync_postgres/replication/conn.rs

pgsync_postgres/replication/
conn.rs
