md5 auth for postgres_backend.rs

proxy/src/main.rs

@@ -0,0 +1,47 @@
+///
+/// Postgres protocol proxy/router.
+///
+/// This service listens psql port and can check auth via external service
+/// (control plane API in our case) and can create new databases and accounts
+/// in somewhat transparent manner (again via communication with control plane API).
+///
+use std::{
+    net::{SocketAddr, TcpListener},
+    thread,
+};
+
+mod cplane_api;
+mod proxy;
+
+pub struct ProxyConf {
+    /// main entrypoint for users to connect to
+    pub proxy_address: SocketAddr,
+
+    /// http management endpoint. Upon user account creation control plane
+    /// will notify us here, so that we can 'unfreeze' user session.
+    pub mgmt_address: SocketAddr,
+
+    /// control plane address where we check auth and create clusters.
+    pub cplane_address: SocketAddr,
+}
+
+fn main() -> anyhow::Result<()> {
+    let conf = ProxyConf {
+        proxy_address: "0.0.0.0:4000".parse()?,
+        mgmt_address: "0.0.0.0:8080".parse()?,
+        cplane_address: "127.0.0.1:3000".parse()?,
+    };
+    let conf: &'static ProxyConf = Box::leak(Box::new(conf));
+
+    // Check that we can bind to address before further initialization
+    println!("Starting proxy on {}", conf.proxy_address);
+    let pageserver_listener = TcpListener::bind(conf.proxy_address)?;
+
+    // Spawn a thread to listen for connections. It will spawn further threads
+    // for each connection.
+    let proxy_listener_thread = thread::Builder::new()
+        .name("Proxy thread".into())
+        .spawn(move || proxy::thread_main(&conf, pageserver_listener))?;
+
+    proxy_listener_thread.join().unwrap()
+}

proxy/src/proxy.rs

@@ -0,0 +1,98 @@
+use crate::ProxyConf;
+use anyhow::bail;
+use bytes::Bytes;
+use std::{
+    net::{TcpListener, TcpStream},
+    thread,
+};
+use zenith_utils::postgres_backend::PostgresBackend;
+use zenith_utils::{
+    postgres_backend,
+    pq_proto::{BeMessage, HELLO_WORLD_ROW, SINGLE_COL_ROWDESC},
+};
+
+///
+/// Main proxy listener loop.
+///
+/// Listens for connections, and launches a new handler thread for each.
+///
+pub fn thread_main(conf: &'static ProxyConf, listener: TcpListener) -> anyhow::Result<()> {
+    loop {
+        let (socket, peer_addr) = listener.accept()?;
+        println!("accepted connection from {}", peer_addr);
+        socket.set_nodelay(true).unwrap();
+
+        thread::spawn(move || {
+            if let Err(err) = proxy_conn_main(conf, socket) {
+                println!("error: {}", err);
+            }
+        });
+    }
+}
+
+pub fn proxy_conn_main(conf: &'static ProxyConf, socket: TcpStream) -> anyhow::Result<()> {
+    let mut conn_handler = ProxyHandler { conf };
+    let mut pgbackend = PostgresBackend::new(socket, postgres_backend::AuthType::MD5)?;
+    pgbackend.run(&mut conn_handler)
+}
+
+struct ProxyHandler {
+    conf: &'static ProxyConf,
+}
+
+// impl ProxyHandler {
+// }
+
+impl postgres_backend::Handler for ProxyHandler {
+    fn process_query(
+        &mut self,
+        pgb: &mut PostgresBackend,
+        query_string: Bytes,
+    ) -> anyhow::Result<()> {
+        println!("Got query: {:?}", query_string);
+        pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
+            .write_message_noflush(&HELLO_WORLD_ROW)?
+            .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
+        pgb.flush()?;
+        Ok(())
+    }
+
+    fn startup(
+        &mut self,
+        _pgb: &mut PostgresBackend,
+        sm: &zenith_utils::pq_proto::FeStartupMessage,
+    ) -> anyhow::Result<()> {
+        println!("Got startup: {:?}", sm);
+        Ok(())
+    }
+
+    fn check_auth_md5(
+        &mut self,
+        pgb: &mut PostgresBackend,
+        md5_response: &[u8],
+    ) -> anyhow::Result<()> {
+        let user = "stask";
+        let pass = "mypassword";
+        let stored_hash = format!(
+            "{:x}",
+            md5::compute([pass.as_bytes(), user.as_bytes()].concat())
+        );
+        let salted_stored_hash = format!(
+            "md5{:x}",
+            md5::compute([stored_hash.as_bytes(), &pgb.md5_salt].concat())
+        );
+
+        let received_hash = std::str::from_utf8(&md5_response)?;
+
+        println!(
+            "check_auth_md5: {:?} vs {}, salt {:?}",
+            received_hash, salted_stored_hash, &pgb.md5_salt
+        );
+
+        if received_hash == salted_stored_hash {
+            Ok(())
+        } else {
+            bail!("Auth failed")
+        }
+    }
+}