commit c0b747ba3454de1990549d09103424f1e85bf0e7 Author: Lev Kokotov Date: Thu Feb 3 13:35:40 2022 -0800 working startup diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ea8c4bf --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +/target diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000..4cc882b --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,272 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 3 + +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + +[[package]] +name = "bytes" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4872d67bab6358e59559027aa3b9157c53d9358c51423c17554809a8858e0f8" + +[[package]] +name = "cfg-if" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" + +[[package]] +name = "hermit-abi" +version = "0.1.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62b467343b94ba476dcb2500d242dadbb39557df889310ac77c5d99100aaac33" +dependencies = [ + "libc", +] + +[[package]] +name = "instant" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a5bbe824c507c5da5956355e86a746d82e0e1464f65d862cc5e71da70e94b2c" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "libc" +version = "0.2.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e74d72e0f9b65b5b4ca49a346af3976df0f9c61d550727f349ecd559f251a26c" + +[[package]] +name = "lock_api" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88943dd7ef4a2e5a4bfa2753aaab3013e34ce2533d1996fb18ef591e315e2b3b" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51b9bbe6c47d51fc3e1a9b945965946b4c44142ab8792c50835a980d362c2710" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "memchr" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "308cc39be01b73d0d18f82a0e7b2a3df85245f84af96fdddc5d202d27e47b86a" + +[[package]] +name = "mio" +version = "0.7.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8067b404fe97c70829f082dec8bcf4f71225d7eaea1d8645349cb76fa06205cc" +dependencies = [ + "libc", + "log", + "miow", + "ntapi", + "winapi", +] + +[[package]] +name = "miow" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9f1c5b025cda876f66ef43a113f91ebc9f4ccef34843000e0adf6ebbab84e21" +dependencies = [ + "winapi", +] + +[[package]] +name = "ntapi" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f6bb902e437b6d86e03cce10a7e2af662292c5dfef23b65899ea3ac9354ad44" +dependencies = [ + "winapi", +] + +[[package]] +name = "num_cpus" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19e64526ebdee182341572e50e9ad03965aa510cd94427a4549448f285e957a1" +dependencies = [ + "hermit-abi", + "libc", +] + +[[package]] +name = "once_cell" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da32515d9f6e6e489d7bc9d84c71b060db7247dc035bbe44eac88cf87486d8d5" + +[[package]] +name = "parking_lot" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d17b78036a60663b797adeaee46f5c9dfebb86948d1255007a1d6be0271ff99" +dependencies = [ + "instant", + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d76e8e1493bcac0d2766c42737f34458f1c8c50c0d23bcb24ea953affb273216" +dependencies = [ + "cfg-if", + "instant", + "libc", + "redox_syscall", + "smallvec", + "winapi", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e280fbe77cc62c91527259e9442153f4688736748d24660126286329742b4c6c" + +[[package]] +name = "proc-macro2" +version = "1.0.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7342d5883fbccae1cc37a2353b09c87c9b0f3afd73f5fb9bba687a1f733b029" +dependencies = [ + "unicode-xid", +] + +[[package]] +name = "quote" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "864d3e96a899863136fc6e99f3d7cae289dafe43bf2c5ac19b70df7210c0a145" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "rabbit" +version = "0.1.0" +dependencies = [ + "bytes", + "tokio", +] + +[[package]] +name = "redox_syscall" +version = "0.2.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8383f39639269cde97d255a32bdb68c047337295414940c68bdd30c2e13203ff" +dependencies = [ + "bitflags", +] + +[[package]] +name = "scopeguard" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd" + +[[package]] +name = "signal-hook-registry" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e51e73328dc4ac0c7ccbda3a494dfa03df1de2f46018127f60c693f2648455b0" +dependencies = [ + "libc", +] + +[[package]] +name = "smallvec" +version = "1.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2dd574626839106c320a323308629dcb1acfc96e32a8cba364ddc61ac23ee83" + +[[package]] +name = "syn" +version = "1.0.86" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a65b3f4ffa0092e9887669db0eae07941f023991ab58ea44da8fe8e2d511c6b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-xid", +] + +[[package]] +name = "tokio" +version = "1.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c27a64b625de6d309e8c57716ba93021dccf1b3b5c97edd6d3dd2d2135afc0a" +dependencies = [ + "bytes", + "libc", + "memchr", + "mio", + "num_cpus", + "once_cell", + "parking_lot", + "pin-project-lite", + "signal-hook-registry", + "tokio-macros", + "winapi", +] + +[[package]] +name = "tokio-macros" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b557f72f448c511a979e2564e55d74e6c4432fc96ff4f6241bc6bded342643b7" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "unicode-xid" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ccb82d61f80a663efe1f787a51b16b5a51e3314d6ac365b08639f52387b33f3" + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..d3772c2 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "rabbit" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +tokio = { version = "1", features = ["full"] } +bytes = "1" \ No newline at end of file diff --git a/src/client.rs b/src/client.rs new file mode 100644 index 0000000..0142dd3 --- /dev/null +++ b/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 { + 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); + } + }; + } + } +} \ No newline at end of file diff --git a/src/errors.rs b/src/errors.rs new file mode 100644 index 0000000..4e3e026 --- /dev/null +++ b/src/errors.rs @@ -0,0 +1,7 @@ +#[derive(Debug, PartialEq)] +pub enum Error { + SocketError, + ClientDisconneted, + ClientBadStartup, + ProtocolSyncError, +} \ No newline at end of file diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..cd6e69b --- /dev/null +++ b/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); + } + }; + }); + } +} diff --git a/src/messages.rs b/src/messages.rs new file mode 100644 index 0000000..c43da68 --- /dev/null +++ b/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), + } +} \ No newline at end of file