Integrate new protocol crate to connection

sqlx-postgres-protocol/src/authentication.rs

@@ -1,4 +1,6 @@
+use crate::Decode;
 use bytes::Bytes;
+use std::io;
 
 #[derive(Debug)]
 pub enum Authentication {
@@ -36,3 +38,13 @@ pub enum Authentication {
     /// SASL authentication has completed.
     SaslFinal { data: Bytes },
 }
+
+impl Decode for Authentication {
+    fn decode(src: Bytes) -> io::Result<Self> {
+        Ok(match src[0] {
+            0 => Authentication::Ok,
+
+            token => unimplemented!("decode not implemented for token: {}", token),
+        })
+    }
+}

sqlx-postgres-protocol/src/backend_key_data.rs

@@ -1,10 +1,32 @@
-use bytes::Bytes;
+use crate::Decode;
+use bytes::{Buf, Bytes};
+use std::io::{self, Cursor};
 
 #[derive(Debug)]
 pub struct BackendKeyData {
     /// The process ID of this backend.
-    pub process_id: u32,
+    process_id: u32,
 
     /// The secret key of this backend.
-    pub secret_key: u32,
+    secret_key: u32,
+}
+
+impl BackendKeyData {
+    pub fn process_id(&self) -> u32 {
+        self.process_id
+    }
+
+    pub fn secret_key(&self) -> u32 {
+        self.secret_key
+    }
+}
+
+impl Decode for BackendKeyData {
+    fn decode(src: Bytes) -> io::Result<Self> {
+        let mut reader = Cursor::new(src);
+        let process_id = reader.get_u32_be();
+        let secret_key = reader.get_u32_be();
+
+        Ok(Self { process_id, secret_key })
+    }
 }

sqlx-postgres-protocol/src/decode.rs

@@ -16,3 +16,11 @@ pub(crate) fn get_str(src: &[u8]) -> io::Result<&str> {
 
     Ok(s)
 }
+
+#[inline]
+pub(crate) fn get_str_bytes_unchecked(src: &Bytes) -> Bytes {
+    let end = memchr(b'\0', &src).unwrap();
+    let buf = src.slice_to(end);
+
+    buf
+}

sqlx-postgres-protocol/src/encode.rs

@@ -2,7 +2,9 @@ use std::io;
 
 pub trait Encode {
     // TODO: Remove
-    fn size_hint(&self) -> usize { 0 }
+    fn size_hint(&self) -> usize {
+        0
+    }
 
     // FIXME: Use BytesMut and not Vec<u8> (also remove the error type here)
     fn encode(&self, buf: &mut Vec<u8>) -> io::Result<()>;

sqlx-postgres-protocol/src/lib.rs

@@ -5,17 +5,23 @@ mod backend_key_data;
 mod decode;
 mod encode;
 mod message;
+mod parameter_status;
 mod password_message;
 mod ready_for_query;
 mod response;
 mod startup_message;
+mod terminate;
 
 pub use self::{
+    authentication::Authentication,
+    backend_key_data::BackendKeyData,
     decode::Decode,
     encode::Encode,
     message::Message,
+    parameter_status::ParameterStatus,
     password_message::PasswordMessage,
     ready_for_query::{ReadyForQuery, TransactionStatus},
     response::{Response, ResponseBuilder, Severity},
     startup_message::StartupMessage,
+    terminate::Terminate,
 };

sqlx-postgres-protocol/src/message.rs

@@ -1,49 +1,53 @@
-use crate::{Decode, Encode, ReadyForQuery, Response};
-use byteorder::{BigEndian, ReadBytesExt};
-use bytes::Bytes;
-use std::io::{self, Cursor};
+use crate::{Authentication, BackendKeyData, Decode, ParameterStatus, ReadyForQuery, Response};
+use byteorder::{BigEndian, ByteOrder};
+use bytes::BytesMut;
+use std::io;
 
 #[derive(Debug)]
 pub enum Message {
+    Authentication(Authentication),
+    ParameterStatus(ParameterStatus),
+    BackendKeyData(BackendKeyData),
     ReadyForQuery(ReadyForQuery),
     Response(Response),
 }
 
-impl Encode for Message {
-    fn size_hint(&self) -> usize {
-        match self {
-            Message::ReadyForQuery(body) => body.size_hint(),
-            Message::Response(body) => body.size_hint(),
-        }
-    }
-
-    fn encode(&self, buf: &mut Vec<u8>) -> io::Result<()> {
-        match self {
-            Message::ReadyForQuery(body) => body.encode(buf),
-            Message::Response(body) => body.encode(buf),
-        }
-    }
-}
-
-impl Decode for Message {
-    fn decode(src: Bytes) -> io::Result<Self>
+impl Message {
+    // FIXME: `Message::decode` shares the name of the remaining message type `::decode` despite being very
+    //        different
+    pub fn decode(src: &mut BytesMut) -> io::Result<Option<Self>>
     where
         Self: Sized,
     {
-        let mut buf = Cursor::new(&src);
+        if src.len() < 5 {
+            // No message is less than 5 bytes
+            return Ok(None);
+        }
 
-        let token = buf.read_u8()?;
-        let len = buf.read_u32::<BigEndian>()? as usize;
-        let pos = buf.position() as usize;
+        let token = src[0];
+        if token == 0 {
+            // FIXME: Handle end-of-stream
+            return Err(io::ErrorKind::InvalidData)?;
+        }
 
-        // `len` includes the size of the length u32
-        let src = src.slice(pos, pos + len - 4);
+        // FIXME: What happens if len(u32) < len(usize) ?
+        let len = BigEndian::read_u32(&src[1..5]) as usize;
 
-        Ok(match token {
+        if src.len() < len {
+            // We don't have enough in the stream yet
+            return Ok(None);
+        }
+
+        let src = src.split_to(len + 1).freeze().slice_from(5);
+
+        Ok(Some(match token {
             b'N' | b'E' => Message::Response(Response::decode(src)?),
+            b'S' => Message::ParameterStatus(ParameterStatus::decode(src)?),
             b'Z' => Message::ReadyForQuery(ReadyForQuery::decode(src)?),
+            b'R' => Message::Authentication(Authentication::decode(src)?),
+            b'K' => Message::BackendKeyData(BackendKeyData::decode(src)?),
 
             _ => unimplemented!("decode not implemented for token: {}", token as char),
-        })
+        }))
     }
 }

sqlx-postgres-protocol/src/parameter_status.rs

@@ -0,0 +1,31 @@
+use crate::{decode::get_str_bytes_unchecked, Decode};
+use bytes::Bytes;
+use std::{io, str};
+
+// FIXME: Use &str functions for a custom Debug
+#[derive(Debug)]
+pub struct ParameterStatus {
+    name: Bytes,
+    value: Bytes,
+}
+
+impl ParameterStatus {
+    #[inline]
+    pub fn name(&self) -> &str {
+        unsafe { str::from_utf8_unchecked(&self.name) }
+    }
+
+    #[inline]
+    pub fn value(&self) -> &str {
+        unsafe { str::from_utf8_unchecked(&self.value) }
+    }
+}
+
+impl Decode for ParameterStatus {
+    fn decode(src: Bytes) -> io::Result<Self> {
+        let name = get_str_bytes_unchecked(&src);
+        let value = get_str_bytes_unchecked(&src.slice_from(name.len() + 1));
+
+        Ok(Self { name, value })
+    }
+}

sqlx-postgres-protocol/src/response.rs

@@ -2,8 +2,7 @@ use crate::{decode::get_str, Decode, Encode};
 use byteorder::{BigEndian, WriteBytesExt};
 use bytes::Bytes;
 use std::{
-    fmt,
-    io::{self, Write},
+    fmt, io,
     ops::Range,
     pin::Pin,
     ptr::NonNull,

sqlx-postgres-protocol/src/startup_message.rs

@@ -1,4 +1,4 @@
-use crate::{Decode, Encode};
+use crate::Encode;
 use bytes::{BufMut, Bytes, BytesMut};
 use std::io;
 
@@ -25,7 +25,7 @@ impl StartupMessage {
 
 impl Encode for StartupMessage {
     fn encode(&self, buf: &mut Vec<u8>) -> io::Result<()> {
-        let len = self.params.len() + 9;
+        let len = self.params.len() + 8;
         buf.reserve(len);
         buf.put_u32_be(len as u32);
         buf.put_u16_be(self.version.0);
@@ -36,12 +36,9 @@ impl Encode for StartupMessage {
     }
 }
 
+// TODO: Impl Iterator to iter over params
 pub struct StartupMessageParams<'a>(&'a [u8]);
 
-// impl Iterator for StartupMessageParams {
-
-// }
-
 pub struct StartupMessageBuilder {
     // (major, minor)
     version: (u16, u16),
@@ -89,7 +86,7 @@ mod tests {
     use crate::Encode;
     use std::io;
 
-    const STARTUP_MESSAGE: &[u8] = b"\0\0\0*\0\x03\0\0user\0postgres\0database\0postgres\0\0";
+    const STARTUP_MESSAGE: &[u8] = b"\0\0\0)\0\x03\0\0user\0postgres\0database\0postgres\0\0";
 
     #[test]
     fn it_encodes_startup_message() -> io::Result<()> {

sqlx-postgres-protocol/src/terminate.rs

@@ -0,0 +1,15 @@
+use crate::Encode;
+use bytes::BufMut;
+use std::io;
+
+pub struct Terminate;
+
+impl Encode for Terminate {
+    fn encode(&self, buf: &mut Vec<u8>) -> io::Result<()> {
+        buf.reserve(5);
+        buf.put_u8(b'X');
+        buf.put_u32_be(4);
+
+        Ok(())
+    }
+}