distinguish between new and old users

proxy/src/cplane_api.rs

@@ -0,0 +1,54 @@
+use anyhow::{bail, Result};
+use std::{collections::HashMap, net::SocketAddr};
+
+pub struct CPlaneApi {
+    address: SocketAddr,
+}
+
+// mock cplane api
+impl CPlaneApi {
+    pub fn new(address: &SocketAddr) -> CPlaneApi {
+        CPlaneApi {
+            address: address.clone(),
+        }
+    }
+
+    pub fn check_auth(&self, user: &str, md5_response: &[u8], salt: &[u8; 4]) -> Result<()> {
+        // passwords for both is "mypass"
+        let auth_map: HashMap<_, &str> = vec![
+            ("stas@zenith", "716ee6e1c4a9364d66285452c47402b1"),
+            ("stas2@zenith", "3996f75df64c16a8bfaf01301b61d582"),
+        ]
+        .into_iter()
+        .collect();
+
+        let stored_hash = auth_map
+            .get(&user)
+            .ok_or_else(|| anyhow::Error::msg("user not found"))?;
+        let salted_stored_hash = format!(
+            "md5{:x}",
+            md5::compute([stored_hash.as_bytes(), salt].concat())
+        );
+
+        let received_hash = std::str::from_utf8(&md5_response)?;
+
+        println!(
+            "auth: {} rh={} sh={} ssh={} {:?}",
+            user, received_hash, stored_hash, salted_stored_hash, salt
+        );
+
+        if received_hash == salted_stored_hash {
+            Ok(())
+        } else {
+            bail!("Auth failed")
+        }
+    }
+
+    fn get_database_uri(_user: String, _database: String) -> Option<String> {
+        Some("postgresql://localhost/stas".to_string())
+    }
+
+    fn create_database(_user: String, _database: String) -> Option<String> {
+        Some("postgresql://localhost/stas".to_string())
+    }
+}

proxy/src/proxy.rs

@@ -1,14 +1,14 @@
-use crate::ProxyConf;
-use anyhow::bail;
+use crate::{cplane_api::CPlaneApi, ProxyConf};
+
 use bytes::Bytes;
 use std::{
     net::{TcpListener, TcpStream},
     thread,
 };
-use zenith_utils::postgres_backend::PostgresBackend;
+use zenith_utils::postgres_backend::{AuthType, PostgresBackend};
 use zenith_utils::{
     postgres_backend,
-    pq_proto::{BeMessage, HELLO_WORLD_ROW, SINGLE_COL_ROWDESC},
+    pq_proto::{BeMessage, SINGLE_COL_ROWDESC},
 };
 
 ///
@@ -31,13 +31,24 @@ pub fn thread_main(conf: &'static ProxyConf, listener: TcpListener) -> anyhow::R
 }
 
 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)?;
+    let mut conn_handler = ProxyHandler {
+        conf,
+        existing_user: false,
+        cplane: CPlaneApi::new(&conf.cplane_address),
+        user: "".into(),
+        database: "".into(),
+    };
+    let mut pgbackend = PostgresBackend::new(socket, postgres_backend::AuthType::Trust)?;
     pgbackend.run(&mut conn_handler)
 }
 
 struct ProxyHandler {
     conf: &'static ProxyConf,
+    existing_user: bool,
+    cplane: CPlaneApi,
+
+    user: String,
+    database: String,
 }
 
 // impl ProxyHandler {
@@ -50,19 +61,51 @@ impl postgres_backend::Handler for ProxyHandler {
         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"))?;
+
+        if !self.existing_user {
+            pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
+                .write_message_noflush(&BeMessage::DataRow(&[Some(b"new user scenario")]))?
+                .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
+        } else {
+            pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
+                .write_message_noflush(&BeMessage::DataRow(&[Some(b"existing user")]))?
+                .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
+        }
+
         pgb.flush()?;
         Ok(())
     }
 
     fn startup(
         &mut self,
-        _pgb: &mut PostgresBackend,
+        pgb: &mut PostgresBackend,
         sm: &zenith_utils::pq_proto::FeStartupMessage,
     ) -> anyhow::Result<()> {
         println!("Got startup: {:?}", sm);
+
+        self.user = sm
+            .params
+            .get("user")
+            .ok_or_else(|| anyhow::Error::msg("user is required in startup packet"))?
+            .into();
+        self.database = sm
+            .params
+            .get("database")
+            .ok_or_else(|| anyhow::Error::msg("database is required in startup packet"))?
+            .into();
+
+        // We use '@zenith' in username as an indicator that user already created
+        // this database and not logging in with his system username.
+        //
+        // With that approach we can create new databases on demand with something like
+        // psql -h zenith.tech -U stas@zenith my_new_db (assuming .pgpass is set). That is
+        // especially helpful if one is setting configuration files for some app that requires
+        // database -- he can just fill config and run initial migration without any other actions.
+        if self.user.ends_with("@zenith") {
+            pgb.auth_type = AuthType::MD5;
+            self.existing_user = true;
+        }
+
         Ok(())
     }
 
@@ -71,28 +114,8 @@ impl postgres_backend::Handler for ProxyHandler {
         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")
-        }
+        assert!(self.existing_user);
+        self.cplane
+            .check_auth(self.user.as_str(), md5_response, &pgb.md5_salt)
     }
 }

zenith_utils/src/postgres_backend.rs

@@ -21,6 +21,10 @@ pub trait Handler {
     fn process_query(&mut self, pgb: &mut PostgresBackend, query_string: Bytes) -> Result<()>;
 
     /// Called on startup packet receival, allows to process params.
+    ///
+    /// If Ok(false) is returned postgres_backend will skip auth -- that is needed for new users
+    /// creation is the proxy code. That is quite hacky and ad-hoc solution, may be we could allow
+    /// to override whole init logic in implementations.
     fn startup(&mut self, _pgb: &mut PostgresBackend, _sm: &FeStartupMessage) -> Result<()> {
         Ok(())
     }
@@ -54,8 +58,9 @@ pub struct PostgresBackend {
     buf_out: BytesMut,
 
     state: ProtoState,
+
     pub md5_salt: [u8; 4],
-    auth_type: AuthType,
+    pub auth_type: AuthType,
 }
 
 // In replication.rs a separate thread is reading keepalives from the
@@ -157,14 +162,16 @@ impl PostgresBackend {
                 Some(FeMessage::StartupMessage(m)) => {
                     trace!("got startup message {:?}", m);
 
-                    handler.startup(self, &m)?;
-
                     match m.kind {
                         StartupRequestCode::NegotiateGss | StartupRequestCode::NegotiateSsl => {
                             info!("SSL requested");
                             self.write_message(&BeMessage::Negotiate)?;
                         }
                         StartupRequestCode::Normal => {
+                            // NB: startup() may change self.auth_type -- we are using that in proxy code
+                            // to bypass auth for new users.
+                            handler.startup(self, &m)?;
+
                             if self.auth_type == AuthType::Trust {
                                 self.write_message_noflush(&BeMessage::AuthenticationOk)?;
                                 // psycopg2 will not connect if client_encoding is not