Integrate new protocol crate to connection

rust-toolchain

@@ -1 +1 @@
-nightly-2019-06-21
+nightly-2019-06-06

sqlx-postgres-protocol/benches/encode.rs

@@ -2,7 +2,7 @@
 extern crate criterion;
 
 use criterion::Criterion;
-use sqlx_postgres_protocol::{Encode, PasswordMessage, StartupMessage, Response, Severity};
+use sqlx_postgres_protocol::{Encode, PasswordMessage, Response, Severity, StartupMessage};
 
 fn criterion_benchmark(c: &mut Criterion) {
     c.bench_function("encode Response(Builder)", |b| {

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(())
+    }
+}

sqlx-postgres/Cargo.toml

@@ -8,6 +8,7 @@ edition = "2018"
 
 [dependencies]
 sqlx-core = { path = "../sqlx-core" }
+sqlx-postgres-protocol = { path = "../sqlx-postgres-protocol" }
 runtime = "=0.3.0-alpha.4"
 futures-preview = "=0.3.0-alpha.16"
 failure = "0.1"

sqlx-postgres/src/connection/establish.rs

@@ -1,11 +1,7 @@
 use super::Connection;
-use crate::protocol::{
-    client::{PasswordMessage, StartupMessage},
-    server::Message as ServerMessage,
-};
 use futures::StreamExt;
 use sqlx_core::ConnectOptions;
-use md5::{Digest, Md5};
+use sqlx_postgres_protocol::{Authentication, Message, PasswordMessage, StartupMessage};
 use std::io;
 
 pub async fn establish<'a, 'b: 'a>(
@@ -14,61 +10,63 @@ pub async fn establish<'a, 'b: 'a>(
 ) -> io::Result<()> {
     // See this doc for more runtime parameters
     // https://www.postgresql.org/docs/12/runtime-config-client.html
-    let params = [
-        ("user", options.user),
-        ("database", options.database),
-        // TODO: Expose this property perhaps?
-        (
-            "application_name",
-            Some(concat!(env!("CARGO_PKG_NAME"), "/v", env!("CARGO_PKG_VERSION"))),
-        ),
+    let mut message = StartupMessage::builder();
+
+    if let Some(user) = options.user {
+        // FIXME: User is technically required. We should default this like psql does.
+        message = message.param("user", user);
+    }
+
+    if let Some(database) = options.database {
+        message = message.param("database", database);
+    }
+
+    let message = message
         // Sets the display format for date and time values,
         // as well as the rules for interpreting ambiguous date input values.
-        ("DateStyle", Some("ISO, MDY")),
+        .param("DateStyle", "ISO, MDY")
         // Sets the display format for interval values.
-        ("IntervalStyle", Some("iso_8601")),
+        .param("IntervalStyle", "iso_8601")
         // Sets the time zone for displaying and interpreting time stamps.
-        ("TimeZone", Some("UTC")),
+        .param("TimeZone", "UTC")
         // Adjust postgres to return percise values for floats
         // NOTE: This is default in postgres 12+
-        ("extra_float_digits", Some("3")),
+        .param("extra_float_digits", "3")
         // Sets the client-side encoding (character set).
-        ("client_encoding", Some("UTF-8")),
-    ];
+        .param("client_encoding", "UTF-8")
+        .build();
 
-    conn.send(StartupMessage { params: &params }).await?;
+    conn.send(message).await?;
 
     // FIXME: This feels like it could be reduced (see other connection flows)
     while let Some(message) = conn.incoming.next().await {
         match message {
-            ServerMessage::AuthenticationOk => {
+            Message::Authentication(Authentication::Ok) => {
                 // Do nothing; server is just telling us that
                 // there is no password needed
             }
 
-            ServerMessage::AuthenticationClearTextPassword => {
-                conn.send(PasswordMessage { password: options.password.unwrap_or_default() })
-                    .await?;
+            Message::Authentication(Authentication::CleartextPassword) => {
+                // FIXME: Should error early (before send) if the user did not supply a password
+                conn.send(PasswordMessage::cleartext(options.password.unwrap_or_default())).await?;
             }
 
-            ServerMessage::AuthenticationMd5Password(body) => {
-                // Hash password|username
-                // FIXME: ConnectOptions should prepare a default user
-                let pass_user =
-                    md5(options.password.unwrap_or_default(), options.user.unwrap_or_default());
-
-                let with_salt = md5(pass_user, body.salt);
-                let password = format!("md5{}", with_salt);
-
-                conn.send(PasswordMessage { password: &password }).await?;
+            Message::Authentication(Authentication::Md5Password { salt }) => {
+                // FIXME: Should error early (before send) if the user did not supply a password
+                conn.send(PasswordMessage::md5(
+                    options.password.unwrap_or_default(),
+                    options.user.unwrap_or_default(),
+                    &salt,
+                ))
+                .await?;
             }
 
-            ServerMessage::BackendKeyData(body) => {
-                conn.process_id = body.process_id;
-                conn.secret_key = body.secret_key;
+            Message::BackendKeyData(body) => {
+                conn.process_id = body.process_id();
+                conn.secret_key = body.secret_key();
             }
 
-            ServerMessage::ReadyForQuery(_) => {
+            Message::ReadyForQuery(_) => {
                 break;
             }
 
@@ -80,8 +78,3 @@ pub async fn establish<'a, 'b: 'a>(
 
     Ok(())
 }
-
-#[inline]
-fn md5(a: impl AsRef<[u8]>, b: impl AsRef<[u8]>) -> String {
-    hex::encode(Md5::new().chain(a).chain(b).result())
-}

sqlx-postgres/src/connection/mod.rs

@@ -1,35 +1,33 @@
-use crate::protocol::{
-    client::{Serialize, Terminate},
-    server::Message as ServerMessage,
-};
 use bytes::BytesMut;
 use futures::{
     channel::mpsc,
     io::{AsyncRead, AsyncReadExt, AsyncWriteExt, ReadHalf, WriteHalf},
-    SinkExt, StreamExt,
+    SinkExt,
 };
-use sqlx_core::ConnectOptions;
 use runtime::{net::TcpStream, task::JoinHandle};
+use sqlx_core::ConnectOptions;
+use sqlx_postgres_protocol::{Encode, Message, Terminate};
 use std::io;
 
 mod establish;
-mod query;
+// mod query;
 
 pub struct Connection {
     writer: WriteHalf<TcpStream>,
-    incoming: mpsc::UnboundedReceiver<ServerMessage>,
+    incoming: mpsc::UnboundedReceiver<Message>,
 
     // Buffer used when serializing outgoing messages
+    // FIXME: Use BytesMut
     wbuf: Vec<u8>,
 
     // Handle to coroutine reading messages from the stream
     receiver: JoinHandle<io::Result<()>>,
 
     // Process ID of the Backend
-    process_id: i32,
+    process_id: u32,
 
     // Backend-unique key to use to send a cancel query message to the server
-    secret_key: i32,
+    secret_key: u32,
 }
 
 impl Connection {
@@ -43,8 +41,8 @@ impl Connection {
             writer,
             receiver,
             incoming: rx,
-            process_id: -1,
-            secret_key: -1,
+            process_id: 0,
+            secret_key: 0,
         };
 
         establish::establish(&mut conn, options).await?;
@@ -52,9 +50,9 @@ impl Connection {
         Ok(conn)
     }
 
-    pub async fn execute<'a, 'b: 'a>(&'a mut self, query: &'b str) -> io::Result<()> {
-        query::query(self, query).await
-    }
+    // pub async fn execute<'a, 'b: 'a>(&'a mut self, query: &'b str) -> io::Result<()> {
+    //     query::query(self, query).await
+    // }
 
     pub async fn close(mut self) -> io::Result<()> {
         self.send(Terminate).await?;
@@ -64,14 +62,16 @@ impl Connection {
         Ok(())
     }
 
-    // Send client-serializable message to the server
-    async fn send<S>(&mut self, message: S) -> io::Result<()>
+    // Send client message to the server
+    async fn send<T>(&mut self, message: T) -> io::Result<()>
     where
-        S: Serialize,
+        T: Encode,
     {
         self.wbuf.clear();
 
-        message.serialize(&mut self.wbuf);
+        message.encode(&mut self.wbuf)?;
+
+        log::trace!("sending: {:?}", bytes::Bytes::from(self.wbuf.clone()));
 
         self.writer.write_all(&self.wbuf).await?;
         self.writer.flush().await?;
@@ -82,7 +82,7 @@ impl Connection {
 
 async fn receiver(
     mut reader: ReadHalf<TcpStream>,
-    mut sender: mpsc::UnboundedSender<ServerMessage>,
+    mut sender: mpsc::UnboundedSender<Message>,
 ) -> io::Result<()> {
     let mut rbuf = BytesMut::with_capacity(0);
     let mut len = 0;
@@ -107,6 +107,7 @@ async fn receiver(
         }
 
         // TODO: Need a select! on a channel that I can trigger to cancel this
+        log::trace!("waiting to read ...");
         let cnt = reader.read(&mut rbuf[len..]).await?;
 
         if cnt > 0 {
@@ -117,28 +118,29 @@ async fn receiver(
         }
 
         while len > 0 {
+            log::trace!("{} in rbuf", len);
+            log::trace!("rbuf: {:?}", rbuf);
+
             let size = rbuf.len();
-            let message = ServerMessage::deserialize(&mut rbuf)?;
+            let message = Message::decode(&mut rbuf)?;
             len -= size - rbuf.len();
 
-            // TODO: Some messages should be kept behind here
             match message {
-                Some(ServerMessage::ParameterStatus(body)) => {
-                    log::debug!("parameter {} = {}", body.name()?, body.value()?);
+                Some(Message::ParameterStatus(body)) => {
+                    log::debug!("parameter: {} = {}", body.name(), body.value());
                 }
 
-                Some(ServerMessage::NoticeResponse(body)) => {
-                    log::warn!("notice: {:?}", body);
+                Some(Message::Response(body)) => {
+                    log::warn!("response: {:?}", body);
                 }
 
                 Some(message) => {
-                    // TODO: Handle this error?
                     sender.send(message).await.unwrap();
                 }
 
                 None => {
                     // Did not receive enough bytes to
-                    // deserialize a complete message
+                    // decode a complete message
                     break;
                 }
             }

sqlx-postgres/src/lib.rs

@@ -1,5 +1,5 @@
 #![feature(non_exhaustive, async_await)]
 #![allow(clippy::needless_lifetimes)]
 
-// mod connection;
-// pub use connection::Connection;
+mod connection;
+pub use connection::Connection;

src/main.rs

@@ -1,17 +1,15 @@
 #![feature(async_await)]
 
-//use sqlx::{pg::Connection, ConnectOptions};
+use sqlx::{pg::Connection, ConnectOptions};
 
 #[runtime::main]
 async fn main() -> Result<(), failure::Error> {
     env_logger::try_init()?;
 
-    //    let mut conn =
-    //        Connection::establish(ConnectOptions::new().user("postgres").password("password")).await?;
-    //
-    //    conn.execute("CREATE EXTENSION IF NOT EXISTS \"uuid-ossp\";").await?;
-    //
-    //    conn.close().await?;
+    let conn =
+        Connection::establish(ConnectOptions::new().user("postgres").password("password")).await?;
+
+    conn.close().await?;
 
     Ok(())
 }