distinguish between new and old users

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)
     }
 }