decoding logic for replication protocol

Co-authored-by: Petros Angelatos <[email protected]>

Author: jeff-davis

postgres-protocol/src/message/backend.rs

@@ -11,6 +11,7 @@ use std::str;
 
 use crate::Oid;
 
+// top-level message tags
 pub const PARSE_COMPLETE_TAG: u8 = b'1';
 pub const BIND_COMPLETE_TAG: u8 = b'2';
 pub const CLOSE_COMPLETE_TAG: u8 = b'3';
@@ -22,6 +23,7 @@ pub const DATA_ROW_TAG: u8 = b'D';
 pub const ERROR_RESPONSE_TAG: u8 = b'E';
 pub const COPY_IN_RESPONSE_TAG: u8 = b'G';
 pub const COPY_OUT_RESPONSE_TAG: u8 = b'H';
+pub const COPY_BOTH_RESPONSE_TAG: u8 = b'W';
 pub const EMPTY_QUERY_RESPONSE_TAG: u8 = b'I';
 pub const BACKEND_KEY_DATA_TAG: u8 = b'K';
 pub const NO_DATA_TAG: u8 = b'n';
@@ -33,6 +35,10 @@ pub const PARAMETER_DESCRIPTION_TAG: u8 = b't';
 pub const ROW_DESCRIPTION_TAG: u8 = b'T';
 pub const READY_FOR_QUERY_TAG: u8 = b'Z';
 
+// replication message tags
+pub const XLOG_DATA_TAG: u8 = b'w';
+pub const PRIMARY_KEEPALIVE_TAG: u8 = b'k';
+
 #[derive(Debug, Copy, Clone)]
 pub struct Header {
     tag: u8,
@@ -93,6 +99,7 @@ pub enum Message {
     CopyDone,
     CopyInResponse(CopyInResponseBody),
     CopyOutResponse(CopyOutResponseBody),
+    CopyBothResponse(CopyBothResponseBody),
     DataRow(DataRowBody),
     EmptyQueryResponse,
     ErrorResponse(ErrorResponseBody),
@@ -190,6 +197,16 @@ impl Message {
                     storage,
                 })
             }
+            COPY_BOTH_RESPONSE_TAG => {
+                let format = buf.read_u8()?;
+                let len = buf.read_u16::<BigEndian>()?;
+                let storage = buf.read_all();
+                Message::CopyBothResponse(CopyBothResponseBody {
+                    format,
+                    len,
+                    storage,
+                })
+            }
             EMPTY_QUERY_RESPONSE_TAG => Message::EmptyQueryResponse,
             BACKEND_KEY_DATA_TAG => {
                 let process_id = buf.read_i32::<BigEndian>()?;
@@ -278,6 +295,59 @@ impl Message {
     }
 }
 
+/// An enum representing Postgres backend replication messages.
+#[non_exhaustive]
+#[derive(Debug)]
+pub enum ReplicationMessage<D> {
+    XLogData(XLogDataBody<D>),
+    PrimaryKeepAlive(PrimaryKeepAliveBody),
+}
+
+impl ReplicationMessage<Bytes> {
+    #[inline]
+    pub fn parse(buf: &Bytes) -> io::Result<Self> {
+        let mut buf = Buffer {
+            bytes: buf.clone(),
+            idx: 0,
+        };
+
+        let tag = buf.read_u8()?;
+
+        let replication_message = match tag {
+            XLOG_DATA_TAG => {
+                let wal_start = buf.read_u64::<BigEndian>()?;
+                let wal_end = buf.read_u64::<BigEndian>()?;
+                let timestamp = buf.read_i64::<BigEndian>()?;
+                let data = buf.read_all();
+                ReplicationMessage::XLogData(XLogDataBody {
+                    wal_start,
+                    wal_end,
+                    timestamp,
+                    data,
+                })
+            }
+            PRIMARY_KEEPALIVE_TAG => {
+                let wal_end = buf.read_u64::<BigEndian>()?;
+                let timestamp = buf.read_i64::<BigEndian>()?;
+                let reply = buf.read_u8()?;
+                ReplicationMessage::PrimaryKeepAlive(PrimaryKeepAliveBody {
+                    wal_end,
+                    timestamp,
+                    reply,
+                })
+            }
+            tag => {
+                return Err(io::Error::new(
+                    io::ErrorKind::InvalidInput,
+                    format!("unknown replication message tag `{}`", tag),
+                ));
+            }
+        };
+
+        Ok(replication_message)
+    }
+}
+
 struct Buffer {
     bytes: Bytes,
     idx: usize,
@@ -524,6 +594,27 @@ impl CopyOutResponseBody {
     }
 }
 
+pub struct CopyBothResponseBody {
+    storage: Bytes,
+    len: u16,
+    format: u8,
+}
+
+impl CopyBothResponseBody {
+    #[inline]
+    pub fn format(&self) -> u8 {
+        self.format
+    }
+
+    #[inline]
+    pub fn column_formats(&self) -> ColumnFormats<'_> {
+        ColumnFormats {
+            remaining: self.len,
+            buf: &self.storage,
+        }
+    }
+}
+
 pub struct DataRowBody {
     storage: Bytes,
     len: u16,
@@ -776,6 +867,78 @@ impl RowDescriptionBody {
     }
 }
 
+#[derive(Debug)]
+pub struct XLogDataBody<D> {
+    wal_start: u64,
+    wal_end: u64,
+    timestamp: i64,
+    data: D,
+}
+
+impl<D> XLogDataBody<D> {
+    #[inline]
+    pub fn wal_start(&self) -> u64 {
+        self.wal_start
+    }
+
+    #[inline]
+    pub fn wal_end(&self) -> u64 {
+        self.wal_end
+    }
+
+    #[inline]
+    pub fn timestamp(&self) -> i64 {
+        self.timestamp
+    }
+
+    #[inline]
+    pub fn data(&self) -> &D {
+        &self.data
+    }
+
+    #[inline]
+    pub fn into_data(self) -> D {
+        self.data
+    }
+
+    pub fn map_data<F, D2, E>(self, f: F) -> Result<XLogDataBody<D2>, E>
+    where
+        F: Fn(D) -> Result<D2, E>,
+    {
+        let data = f(self.data)?;
+        Ok(XLogDataBody {
+            wal_start: self.wal_start,
+            wal_end: self.wal_end,
+            timestamp: self.timestamp,
+            data,
+        })
+    }
+}
+
+#[derive(Debug)]
+pub struct PrimaryKeepAliveBody {
+    wal_end: u64,
+    timestamp: i64,
+    reply: u8,
+}
+
+impl PrimaryKeepAliveBody {
+    #[inline]
+    pub fn wal_end(&self) -> u64 {
+        self.wal_end
+    }
+
+    #[inline]
+    pub fn timestamp(&self) -> i64 {
+        self.timestamp
+    }
+
+    #[inline]
+    pub fn reply(&self) -> u8 {
+        self.reply
+    }
+}
+
 pub struct Fields<'a> {
     buf: &'a [u8],
     remaining: u16,