working q

src/messages.rs

@@ -1,59 +1,11 @@
 
 use tokio::net::TcpStream;
-use tokio::io::{AsyncReadExt, AsyncWriteExt};
-use bytes::{Buf, BufMut, BytesMut};
+use tokio::net::tcp::OwnedReadHalf;
+use tokio::io::{AsyncWriteExt, BufReader, AsyncReadExt};
+use bytes::{BufMut, BytesMut};
 
 use crate::errors::Error;
 
-/// Handle the startup phase for the client.
-/// This one is special because Startup and SSLRequest
-/// packages don't start with a u8 letter code.
-pub async fn handle_client_startup(stream: &mut TcpStream) -> Result<(), Error> {
-    loop {
-        // Could be StartupMessage or SSLRequest
-        // which makes this variable length.
-        let len = match stream.read_i32().await {
-            Ok(len) => len,
-            Err(_) => return Err(Error::ClientBadStartup),
-        };
-
-        // Read whatever is left.
-        let mut startup = vec![0u8; len as usize - 4];
-
-        match stream.read_exact(&mut startup).await {
-            Ok(_) => (),
-            Err(_) => return Err(Error::ClientBadStartup),
-        };
-
-        let mut bytes = BytesMut::from(&startup[..]);
-        let code = bytes.get_i32();
-
-        match code {
-            // Client wants SSL. We don't support it at the moment.
-            80877103 => {
-                let mut no = BytesMut::with_capacity(1);
-                no.put_u8(b'N');
-
-                write_all(stream, no).await?;
-            },
-
-            // Regular startup message.
-            196608 => {
-                // TODO: perform actual auth.
-                // TODO: record startup parameters client sends over.
-                auth_ok(stream).await?;
-                ready_for_query(stream).await?;
-                return Ok(());
-            },
-
-            _ => {
-                return Err(Error::ProtocolSyncError);
-            }
-        };
-    }
-}
-
-
 pub async fn auth_ok(stream: &mut TcpStream) -> Result<(), Error> {
     let mut auth_ok = BytesMut::with_capacity(9);
 
@@ -79,4 +31,32 @@ pub async fn write_all(stream: &mut TcpStream, buf: BytesMut) -> Result<(), Erro
         Ok(_) => Ok(()),
         Err(_) => return Err(Error::SocketError),
     }
+}
+
+/// Read a complete message from the socket.
+pub async fn read_message(stream: &mut BufReader<OwnedReadHalf>) -> Result<BytesMut, Error> {
+    let code = match stream.read_u8().await {
+        Ok(code) => code,
+        Err(_) => return Err(Error::SocketError),
+    };
+
+    let len = match stream.read_i32().await {
+        Ok(len) => len,
+        Err(_) => return Err(Error::SocketError),
+    };
+
+    let mut buf = vec![0u8; len as usize - 4];
+
+    match stream.read_exact(&mut buf).await {
+        Ok(_) => (),
+        Err(_) => return Err(Error::SocketError),
+    };
+
+    let mut bytes = BytesMut::with_capacity(len as usize + 1);
+
+    bytes.put_u8(code);
+    bytes.put_i32(len);
+    bytes.put_slice(&buf);
+
+    Ok(bytes)
 }