working startup

src/client.rs

@@ -0,0 +1,64 @@
+/// PostgreSQL client (frontend).
+/// We are pretending to be the backend.
+
+use tokio::net::TcpStream;
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
+use bytes::{BytesMut, Buf, BufMut};
+
+use crate::errors::Error;
+use crate::messages::*;
+
+pub struct Client {
+    stream: TcpStream,
+}
+
+impl Client {
+    pub async fn startup(mut stream: TcpStream) -> Result<Client, 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(&mut stream, no).await?;
+                },
+
+                // Regular startup message.
+                196608 => {
+                    // TODO: perform actual auth.
+                    // TODO: record startup parameters client sends over.
+                    auth_ok(&mut stream).await?;
+                    ready_for_query(&mut stream).await?;
+
+                    return Ok(Client {
+                        stream: stream,
+                    });
+                },
+
+                _ => {
+                    return Err(Error::ProtocolSyncError);
+                }
+            };
+        }
+    }
+}

src/errors.rs

@@ -0,0 +1,7 @@
+#[derive(Debug, PartialEq)]
+pub enum Error {
+    SocketError,
+    ClientDisconneted,
+    ClientBadStartup,
+    ProtocolSyncError,
+}

src/main.rs

@@ -0,0 +1,46 @@
+extern crate bytes;
+extern crate tokio;
+
+use tokio::net::TcpListener;
+
+mod errors;
+mod messages;
+mod client;
+
+#[tokio::main]
+async fn main() {
+    println!("> Welcome to PgRabbit");
+
+    let listener = match TcpListener::bind("0.0.0.0:5433").await {
+        Ok(sock) => sock,
+        Err(err) => {
+            println!("> Error: {:?}", err);
+            return;
+        }
+    };
+
+    loop {
+        let (mut socket, addr) = match listener.accept().await {
+            Ok((mut socket, addr)) => (socket, addr),
+            Err(err) => {
+                println!("> Listener: {:?}", err);
+                continue;
+            }
+        };
+
+        // Client goes to another thread, bye.
+        tokio::task::spawn(async move {
+            println!(">> Client {:?} connected", addr);
+
+            match client::Client::startup(socket).await {
+                Ok(client) => {
+                    println!(">> Client {:?} connected successfully!", addr);
+                },
+
+                Err(err) => {
+                    println!(">> Error: {:?}", err);
+                }
+            };
+        });
+    }
+}

src/messages.rs

@@ -0,0 +1,82 @@
+
+use tokio::net::TcpStream;
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use bytes::{Buf, 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);
+
+    auth_ok.put_u8(b'R');
+    auth_ok.put_i32(8);
+    auth_ok.put_i32(0);
+
+    Ok(write_all(stream, auth_ok).await?)
+}
+
+pub async fn ready_for_query(stream: &mut TcpStream) -> Result<(), Error> {
+    let mut bytes = BytesMut::with_capacity(5);
+
+    bytes.put_u8(b'Z');
+    bytes.put_i32(5);
+    bytes.put_u8(b'I'); // Idle
+
+    Ok(write_all(stream, bytes).await?)
+}
+
+pub async fn write_all(stream: &mut TcpStream, buf: BytesMut) -> Result<(), Error> {
+    match stream.write_all(&buf).await {
+        Ok(_) => Ok(()),
+        Err(_) => return Err(Error::SocketError),
+    }
+}