diff --git a/proxy/src/cplane_api.rs b/proxy/src/cplane_api.rs new file mode 100644 index 0000000000..ae544e98dd --- /dev/null +++ b/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 { + Some("postgresql://localhost/stas".to_string()) + } + + fn create_database(_user: String, _database: String) -> Option { + Some("postgresql://localhost/stas".to_string()) + } +} diff --git a/proxy/src/proxy.rs b/proxy/src/proxy.rs index 169503b987..9a5c1e6b62 100644 --- a/proxy/src/proxy.rs +++ b/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) } } diff --git a/zenith_utils/src/postgres_backend.rs b/zenith_utils/src/postgres_backend.rs index 7565ee75ba..9467dcbb69 100644 --- a/zenith_utils/src/postgres_backend.rs +++ b/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