Prototype a local test mocking cplane

Cargo.lock

@@ -1334,7 +1334,6 @@ dependencies = [
 [[package]]
 name = "postgres-protocol"
 version = "0.6.1"
-source = "git+https://github.com/zenithdb/rust-postgres.git?rev=9eb0dbfbeb6a6c1b79099b9f7ae4a8c021877858#9eb0dbfbeb6a6c1b79099b9f7ae4a8c021877858"
 dependencies = [
  "base64 0.13.0",
  "byteorder",
@@ -1352,7 +1351,7 @@ dependencies = [
 [[package]]
 name = "postgres-protocol"
 version = "0.6.1"
-source = "git+https://github.com/zenithdb/rust-postgres.git?rev=f1f16657aaebe2b9b4b16ef7abf6dc42301bad5d#f1f16657aaebe2b9b4b16ef7abf6dc42301bad5d"
+source = "git+https://github.com/zenithdb/rust-postgres.git?rev=9eb0dbfbeb6a6c1b79099b9f7ae4a8c021877858#9eb0dbfbeb6a6c1b79099b9f7ae4a8c021877858"
 dependencies = [
  "base64 0.13.0",
  "byteorder",
@@ -1370,21 +1369,20 @@ dependencies = [
 [[package]]
 name = "postgres-types"
 version = "0.2.1"
-source = "git+https://github.com/zenithdb/rust-postgres.git?rev=9eb0dbfbeb6a6c1b79099b9f7ae4a8c021877858#9eb0dbfbeb6a6c1b79099b9f7ae4a8c021877858"
 dependencies = [
  "bytes",
  "fallible-iterator",
- "postgres-protocol 0.6.1 (git+https://github.com/zenithdb/rust-postgres.git?rev=9eb0dbfbeb6a6c1b79099b9f7ae4a8c021877858)",
+ "postgres-protocol 0.6.1",
 ]
 
 [[package]]
 name = "postgres-types"
 version = "0.2.1"
-source = "git+https://github.com/zenithdb/rust-postgres.git?rev=f1f16657aaebe2b9b4b16ef7abf6dc42301bad5d#f1f16657aaebe2b9b4b16ef7abf6dc42301bad5d"
+source = "git+https://github.com/zenithdb/rust-postgres.git?rev=9eb0dbfbeb6a6c1b79099b9f7ae4a8c021877858#9eb0dbfbeb6a6c1b79099b9f7ae4a8c021877858"
 dependencies = [
  "bytes",
  "fallible-iterator",
- "postgres-protocol 0.6.1 (git+https://github.com/zenithdb/rust-postgres.git?rev=f1f16657aaebe2b9b4b16ef7abf6dc42301bad5d)",
+ "postgres-protocol 0.6.1 (git+https://github.com/zenithdb/rust-postgres.git?rev=9eb0dbfbeb6a6c1b79099b9f7ae4a8c021877858)",
 ]
 
 [[package]]
@@ -1463,9 +1461,10 @@ dependencies = [
  "serde",
  "serde_json",
  "sha2",
+ "stringprep",
  "thiserror",
  "tokio",
- "tokio-postgres 0.7.1 (git+https://github.com/zenithdb/rust-postgres.git?rev=f1f16657aaebe2b9b4b16ef7abf6dc42301bad5d)",
+ "tokio-postgres 0.7.1",
  "zenith_utils",
 ]
 
@@ -2111,6 +2110,27 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "tokio-postgres"
+version = "0.7.1"
+dependencies = [
+ "async-trait",
+ "byteorder",
+ "bytes",
+ "fallible-iterator",
+ "futures",
+ "log",
+ "parking_lot",
+ "percent-encoding",
+ "phf",
+ "pin-project-lite",
+ "postgres-protocol 0.6.1",
+ "postgres-types 0.2.1",
+ "socket2",
+ "tokio",
+ "tokio-util",
+]
+
 [[package]]
 name = "tokio-postgres"
 version = "0.7.1"
@@ -2133,28 +2153,6 @@ dependencies = [
  "tokio-util",
 ]
 
-[[package]]
-name = "tokio-postgres"
-version = "0.7.1"
-source = "git+https://github.com/zenithdb/rust-postgres.git?rev=f1f16657aaebe2b9b4b16ef7abf6dc42301bad5d#f1f16657aaebe2b9b4b16ef7abf6dc42301bad5d"
-dependencies = [
- "async-trait",
- "byteorder",
- "bytes",
- "fallible-iterator",
- "futures",
- "log",
- "parking_lot",
- "percent-encoding",
- "phf",
- "pin-project-lite",
- "postgres-protocol 0.6.1 (git+https://github.com/zenithdb/rust-postgres.git?rev=f1f16657aaebe2b9b4b16ef7abf6dc42301bad5d)",
- "postgres-types 0.2.1 (git+https://github.com/zenithdb/rust-postgres.git?rev=f1f16657aaebe2b9b4b16ef7abf6dc42301bad5d)",
- "socket2",
- "tokio",
- "tokio-util",
-]
-
 [[package]]
 name = "tokio-rustls"
 version = "0.22.0"

proxy/Cargo.toml

@@ -22,9 +22,10 @@ rustls = "0.19.1"
 serde = "1"
 serde_json = "1"
 sha2 = "0.9.8"
+stringprep = "0.1.2"
 thiserror = "1.0.30"
 tokio = { version = "1.11", features = ['macros'] }
-# tokio-postgres = { path = "../../rust-postgres/tokio-postgres" } TODO remove this
-tokio-postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev = "f1f16657aaebe2b9b4b16ef7abf6dc42301bad5d"}
+tokio-postgres = { path = "../../rust-postgres/tokio-postgres" }
+# tokio-postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev = "f1f16657aaebe2b9b4b16ef7abf6dc42301bad5d"}
 
 zenith_utils = { path = "../zenith_utils" }

proxy/src/proxy.rs

@@ -4,12 +4,16 @@ use crate::ProxyState;
 use crate::db::{AuthSecret, DatabaseAuthInfo, ScramAuthSecret};
 use crate::scram::key::{SCRAM_KEY_LEN, ScramKey};
 use anyhow::{anyhow, bail};
+use hmac::digest::{FixedOutput};
+use hmac::{Hmac, Mac, NewMac};
 use lazy_static::lazy_static;
 use parking_lot::Mutex;
 use rand::prelude::StdRng;
-use rand::{Rng, SeedableRng};
+use rand::{Fill, Rng, SeedableRng};
+use sha2::{Digest, Sha256};
 use std::cell::Cell;
 use std::collections::HashMap;
+use std::convert::TryInto;
 use std::net::{SocketAddr, TcpStream};
 use std::{io, thread};
 use tokio_postgres::NoTls;
@@ -111,6 +115,39 @@ pub fn proxy_conn_main(state: &'static ProxyState, socket: TcpStream) -> anyhow:
     proxy(client.split(), server.split())
 }
 
+// HACK copied from tokio-postgres
+// since postgres passwords are not required to exclude saslprep-prohibited
+// characters or even be valid UTF8, we run saslprep if possible and otherwise
+// return the raw password.
+fn normalize(pass: &str) -> Vec<u8> {
+    match stringprep::saslprep(pass) {
+        Ok(pass) => pass.into_owned().into_bytes(),
+        Err(_) => pass.as_bytes().to_vec(),
+    }
+}
+
+// HACK copied from tokio-postgres
+fn hi(str: &[u8], salt: &[u8], i: u32) -> [u8; 32] {
+    let mut hmac = Hmac::<Sha256>::new_varkey(str).expect("HMAC is able to accept all key sizes");
+    hmac.update(salt);
+    hmac.update(&[0, 0, 0, 1]);
+    let mut prev = hmac.finalize().into_bytes();
+
+    let mut hi = prev;
+
+    for _ in 1..i {
+        let mut hmac = Hmac::<Sha256>::new_varkey(str).expect("already checked above");
+        hmac.update(&prev);
+        prev = hmac.finalize().into_bytes();
+
+        for (hi, prev) in hi.iter_mut().zip(prev) {
+            *hi ^= prev;
+        }
+    }
+
+    hi.into()
+}
+
 impl ProxyConnection {
     /// Returns Ok(None) when connection was successfully closed.
     fn handle_client(mut self) -> anyhow::Result<Option<(Stream, TcpStream)>> {
@@ -121,7 +158,7 @@ impl ProxyConnection {
             };
 
             // Both scenarios here should end up producing database credentials
-            if username.ends_with("@zenith") {
+            if true || username.ends_with("@zenith") {
                 self.handle_existing_user(&username, &dbname).map(Some)
             } else {
                 self.handle_new_user().map(Some)
@@ -217,21 +254,51 @@ impl ProxyConnection {
     fn handle_existing_user(&mut self, user: &str, db: &str) -> anyhow::Result<DatabaseAuthInfo> {
         let _cplane = CPlaneApi::new(&self.state.conf.auth_endpoint, &self.state.waiters);
 
+        // TODO read from postres instance
+        // I got these values from the proxy error log on key cache miss
+        let salt = [180, 76, 114, 155, 212, 214, 236, 192, 101, 236, 235, 4, 212, 87, 25, 85];
+        let iterations = 4096;
+
+        // TODO read from CLI
+        let password = "postgres";
+
+        let normalized = normalize(password);
+        let salted = hi(&normalized, &salt, iterations);
+
+        let mut mac = Hmac::<Sha256>::new_varkey(&salted).unwrap();
+        mac.update(b"Client Key");
+        let client_key = mac.finalize().into_bytes();
+
+        let mut mac = Hmac::<Sha256>::new_varkey(&salted).unwrap();
+        mac.update(b"Server Key");
+        let server_key = mac.finalize().into_bytes();
+
+        let mut hash = Sha256::default();
+        hash.update(client_key);
+        let stored_key = hash.finalize_fixed();
+
+        let secret = crate::scram::ScramSecret {
+            iterations,
+            salt_base64: base64::encode(salt),
+            stored_key: ScramKey { bytes: stored_key.try_into()? },
+            server_key: ScramKey { bytes: server_key.try_into()? },
+        };
+
         // TODO: fetch secret from console
         // user='user' password='password'
-        let secret = crate::scram::ScramSecret::parse(
-            &[
-                "SCRAM-SHA-256",
-                "4096:XiWzgkfGNyY3ipsz08PY+A==",
-                &[
-                    "YMmirZHYtTB6erVDCxL4Zjn66Kn7RCfS+aV3qROV4o8=",
-                    "aCSKHnugk1l9Ut6VhO5VeeWsB8xhVdPk/NyEgjOJ3nk=",
-                ]
-                .join(":"),
-            ]
-            .join("$"),
-        )
-        .unwrap();
+        // let secret = crate::scram::ScramSecret::parse(
+        //     &[
+        //         "SCRAM-SHA-256",
+        //         "4096:XiWzgkfGNyY3ipsz08PY+A==",
+        //         &[
+        //             "YMmirZHYtTB6erVDCxL4Zjn66Kn7RCfS+aV3qROV4o8=",
+        //             "aCSKHnugk1l9Ut6VhO5VeeWsB8xhVdPk/NyEgjOJ3nk=",
+        //         ]
+        //         .join(":"),
+        //     ]
+        //     .join("$"),
+        // )
+        // .unwrap();
 
         AuthStream::new(&mut self.pgb)
             .begin(auth::Scram(&secret))?
@@ -242,11 +309,12 @@ impl ProxyConnection {
             .write_message_noflush(&BeParameterStatusMessage::encoding())?;
 
         // TODO get this info from console and tell it to start the db
-        let host = "";
-        let port = 0;
+        let host = "127.0.0.1";
+        let port = 5432;
 
         // TODO fish out the real client_key from the authenticate() call above
-        let client_key: ScramKey = [0; SCRAM_KEY_LEN].into();
+        // let client_key: ScramKey = [0; SCRAM_KEY_LEN].into();
+        let client_key = ScramKey { bytes: client_key.as_slice().try_into().unwrap() };
         let scram_auth_secret = ScramAuthSecret {
             iterations: secret.iterations,
             salt_base64: secret.salt_base64,