working q

src/client.rs

@@ -2,7 +2,8 @@
 /// We are pretending to be the backend.
 
 use tokio::net::TcpStream;
-use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
+use tokio::io::{AsyncReadExt, BufReader};
 
 use bytes::{BytesMut, Buf, BufMut};
 
@@ -10,7 +11,8 @@ use crate::errors::Error;
 use crate::messages::*;
 
 pub struct Client {
-    stream: TcpStream,
+    read: BufReader<OwnedReadHalf>,
+    write: OwnedWriteHalf,
 }
 
 impl Client {
@@ -50,8 +52,11 @@ impl Client {
                     auth_ok(&mut stream).await?;
                     ready_for_query(&mut stream).await?;
 
+                    let (read, write) = stream.into_split();
+
                     return Ok(Client {
-                        stream: stream,
+                        read: BufReader::new(read),
+                        write: write,
                     });
                 },
 
@@ -61,4 +66,23 @@ impl Client {
             };
         }
     }
+
+    pub async fn handle(&mut self) -> Result<(), Error> {
+        loop {
+            let mut message = read_message(&mut self.read).await?;
+            let original = message.clone(); // To be forwarded to the server
+            let code = message.get_u8() as char;
+            let len = message.get_i32() as usize;
+
+            match code {
+                'Q' => {
+                    println!(">>> Query: {:?}", message);
+                },
+
+                _ => {
+                    println!(">>> Unexpected code: {}", code);
+                },
+            }
+        }
+    }
 }

src/main.rs

@@ -20,8 +20,8 @@ async fn main() {
     };
 
     loop {
-        let (mut socket, addr) = match listener.accept().await {
+        let (socket, addr) = match listener.accept().await {
-            Ok((mut socket, addr)) => (socket, addr),
+            Ok((socket, addr)) => (socket, addr),
             Err(err) => {
                 println!("> Listener: {:?}", err);
                 continue;
@@ -33,8 +33,9 @@ async fn main() {
             println!(">> Client {:?} connected", addr);
 
             match client::Client::startup(socket).await {
-                Ok(client) => {
+                Ok(mut client) => {
                     println!(">> Client {:?} connected successfully!", addr);
+                    client.handle().await;
                 },
 
                 Err(err) => {

src/messages.rs

@@ -1,59 +1,11 @@
 
 use tokio::net::TcpStream;
-use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::net::tcp::OwnedReadHalf;
-use bytes::{Buf, BufMut, BytesMut};
+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)
 }