diff --git a/Cargo.lock b/Cargo.lock index 84f18b5..945b90b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -63,15 +63,6 @@ version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af" -[[package]] -name = "block-buffer" -version = "0.10.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" -dependencies = [ - "generic-array", -] - [[package]] name = "block-buffer" version = "0.12.0" @@ -124,7 +115,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" dependencies = [ "cfg-if", - "cpufeatures 0.3.0", + "cpufeatures", "rand_core 0.10.0", ] @@ -149,15 +140,6 @@ version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" -[[package]] -name = "cpufeatures" -version = "0.2.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" -dependencies = [ - "libc", -] - [[package]] name = "cpufeatures" version = "0.3.0" @@ -167,16 +149,6 @@ dependencies = [ "libc", ] -[[package]] -name = "crypto-common" -version = "0.1.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" -dependencies = [ - "generic-array", - "typenum", -] - [[package]] name = "crypto-common" version = "0.2.1" @@ -195,26 +167,15 @@ dependencies = [ "cmov", ] -[[package]] -name = "digest" -version = "0.10.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" -dependencies = [ - "block-buffer 0.10.4", - "crypto-common 0.1.7", - "subtle", -] - [[package]] name = "digest" version = "0.11.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4850db49bf08e663084f7fb5c87d202ef91a3907271aff24a94eb97ff039153c" dependencies = [ - "block-buffer 0.12.0", + "block-buffer", "const-oid", - "crypto-common 0.2.1", + "crypto-common", "ctutils", ] @@ -304,16 +265,6 @@ dependencies = [ "pin-project-lite", ] -[[package]] -name = "generic-array" -version = "0.14.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" -dependencies = [ - "typenum", - "version_check", -] - [[package]] name = "getrandom" version = "0.2.17" @@ -372,22 +323,13 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" -[[package]] -name = "hmac" -version = "0.12.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e" -dependencies = [ - "digest 0.10.7", -] - [[package]] name = "hmac" version = "0.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6303bc9732ae41b04cb554b844a762b4115a61bfaa81e3e83050991eeb56863f" dependencies = [ - "digest 0.11.2", + "digest", ] [[package]] @@ -516,7 +458,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69b6441f590336821bb897fb28fc622898ccceb1d6cea3fde5ea86b090c4de98" dependencies = [ "cfg-if", - "digest 0.11.2", + "digest", ] [[package]] @@ -602,14 +544,12 @@ checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" name = "pgwire-supabase-proxy" version = "0.1.0" dependencies = [ - "base64", "bytes", - "hmac 0.12.1", "jsonwebtoken", + "postgres-protocol", "rand 0.8.5", "serde", "serde_json", - "sha2 0.10.9", "thiserror", "tokio", "tokio-postgres", @@ -653,11 +593,11 @@ dependencies = [ "byteorder", "bytes", "fallible-iterator", - "hmac 0.13.0", + "hmac", "md-5", "memchr", "rand 0.10.0", - "sha2 0.11.0", + "sha2", "stringprep", ] @@ -855,17 +795,6 @@ dependencies = [ "zmij", ] -[[package]] -name = "sha2" -version = "0.10.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" -dependencies = [ - "cfg-if", - "cpufeatures 0.2.17", - "digest 0.10.7", -] - [[package]] name = "sha2" version = "0.11.0" @@ -873,8 +802,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" dependencies = [ "cfg-if", - "cpufeatures 0.3.0", - "digest 0.11.2", + "cpufeatures", + "digest", ] [[package]] @@ -944,12 +873,6 @@ dependencies = [ "unicode-properties", ] -[[package]] -name = "subtle" -version = "2.6.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" - [[package]] name = "syn" version = "2.0.117" @@ -1205,12 +1128,6 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" -[[package]] -name = "version_check" -version = "0.9.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" - [[package]] name = "wasi" version = "0.11.1+wasi-snapshot-preview1" diff --git a/Cargo.toml b/Cargo.toml index 5ab69f5..ea85287 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,6 +17,7 @@ tokio = { version = "1", features = ["rt-multi-thread", "net", "macros", "sync", # Wire protocol bytes = "1" +postgres-protocol = "0.6.11" # Auth jsonwebtoken = { version = "10", default-features = false, features = ["aws_lc_rs"] } @@ -32,13 +33,6 @@ thiserror = "2" tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } -# SCRAM -hmac = "0.12" -sha2 = "0.10" - -# Base64 (SCRAM) -base64 = "0.22" - # Random rand = "0.8" diff --git a/src/lib.rs b/src/lib.rs index cc2d563..2e34454 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,7 +7,6 @@ mod auth; mod error; mod proxy; -mod scram; mod wire; // Re-export Config for use by integration tests and main binary. diff --git a/src/main.rs b/src/main.rs index 1a76627..6165e78 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,8 +6,7 @@ use tokio::net::TcpListener; #[tokio::main] async fn main() -> std::result::Result<(), Box> { - let backend_postgres_url = - std::env::var("DATABASE_URL").expect("DATABASE_URL must be set"); + let backend_postgres_url = std::env::var("DATABASE_URL").expect("DATABASE_URL must be set"); let jwt_secret = std::env::var("SUPABASE_JWT_SECRET").expect("SUPABASE_JWT_SECRET must be set"); let listen_addr: SocketAddr = std::env::var("LISTEN_ADDR") .unwrap_or_else(|_| "0.0.0.0:5432".to_string()) diff --git a/src/proxy.rs b/src/proxy.rs index 456f239..7e7544b 100644 --- a/src/proxy.rs +++ b/src/proxy.rs @@ -2,7 +2,6 @@ use crate::auth::{Claims, JwtAuthenticator}; use crate::error::ProxyError; -use crate::scram; use crate::wire; use std::collections::HashMap; use std::net::SocketAddr; @@ -10,8 +9,6 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; - - /// Start the byte-forward proxy server. pub async fn serve( config: crate::Config, @@ -136,14 +133,20 @@ async fn handle_connection( let (backend_host, backend_port, backend_user, backend_password, backend_db) = parse_backend_url(&config.backend_postgres_url)?; - let mut backend = match tokio::net::TcpStream::connect((backend_host.as_str(), backend_port)).await { - Ok(s) => s, - Err(e) => { - tracing::error!(peer_addr = %peer_addr, error = %e, "backend TCP connect failed"); - wire::write_error_response(&mut client, "08001", &format!("backend connection failed: {}", e)).await?; - return Ok(()); - } - }; + let mut backend = + match tokio::net::TcpStream::connect((backend_host.as_str(), backend_port)).await { + Ok(s) => s, + Err(e) => { + tracing::error!(peer_addr = %peer_addr, error = %e, "backend TCP connect failed"); + wire::write_error_response( + &mut client, + "08001", + &format!("backend connection failed: {}", e), + ) + .await?; + return Ok(()); + } + }; tracing::debug!(peer_addr = %peer_addr, backend = %backend_host, "backend TCP opened"); // Step 7: Send backend StartupMessage (plain TCP, no TLS) @@ -162,17 +165,40 @@ async fn handle_connection( } wire::AuthMethod::CleartextPassword => { wire::write_password_message(&mut backend, &backend_password).await?; + tracing::debug!(peer_addr = %peer_addr, "backend auth: cleartext password sent"); } wire::AuthMethod::Sasl { mechanisms } => { - if !mechanisms.contains(&"SCRAM-SHA-256".to_string()) { + use postgres_protocol::authentication::sasl::{ + ChannelBinding, ScramSha256, SCRAM_SHA_256, + }; + + if !mechanisms.iter().any(|m| m.as_str() == SCRAM_SHA_256) { return Err(Box::new(ProxyError::BackendAuth(format!( "unsupported SASL mechanisms: {:?}", mechanisms )))); } - scram::scram_sha_256_authenticate(&mut backend, &backend_user, &backend_password) - .await - .map_err(|e| Box::new(e) as Box)?; + + let mut scram = + ScramSha256::new(backend_password.as_bytes(), ChannelBinding::unsupported()); + + wire::write_sasl_initial_response(&mut backend, SCRAM_SHA_256, scram.message()).await?; + + let server_first = wire::read_sasl_auth_message(&mut backend, 11).await?; + scram + .update(&server_first) + .map_err(|e| ProxyError::BackendAuth(format!("scram update: {}", e)))?; + + wire::write_sasl_response(&mut backend, scram.message()).await?; + + let server_final = wire::read_sasl_auth_message(&mut backend, 12).await?; + scram + .finish(&server_final) + .map_err(|e| ProxyError::BackendAuth(format!("scram finish: {}", e)))?; + + let _ = wire::read_sasl_auth_message(&mut backend, 0).await?; + + tracing::debug!(peer_addr = %peer_addr, "backend SCRAM-SHA-256 auth complete"); } wire::AuthMethod::Md5Password { .. } => { return Err(Box::new(ProxyError::BackendAuth( diff --git a/src/scram.rs b/src/scram.rs deleted file mode 100644 index 814398d..0000000 --- a/src/scram.rs +++ /dev/null @@ -1,437 +0,0 @@ -//! SCRAM-SHA-256 client authentication helper. - -use crate::error::ProxyError; -use base64::{engine::general_purpose::STANDARD as BASE64, Engine}; -use hmac::{Hmac, Mac}; -use rand::Rng; -use sha2::Sha256; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; - -type HmacSha256 = Hmac; - -const SHA256_NAME: &str = "SCRAM-SHA-256"; - -/// Perform SCRAM-SHA-256 authentication with the backend. -pub async fn scram_sha_256_authenticate( - stream: &mut S, - username: &str, - password: &str, -) -> Result<(), ProxyError> -where - S: AsyncReadExt + AsyncWriteExt + Unpin, -{ - // Generate nonce before first .await (ThreadRng is !Send) - let client_nonce: String = { - let mut rng = rand::thread_rng(); - (0..18) - .map(|_| { - let b: u8 = rng.gen(); - "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789" - .chars() - .nth((b % 62) as usize) - .unwrap() - }) - .collect() - }; // rng dropped here, before first .await - - // Step 1: ClientFirst — GS2 header is "n,," (no channel binding) - let client_first_bare = format!("n={},r={}", username, client_nonce); - let client_first = format!("n,,{}", client_first_bare); - send_sasl_initial_response(stream, SHA256_NAME, client_first.as_bytes()).await?; - - // Step 2: ServerFirst - let server_first_raw = read_sasl_continue(stream).await?; - let server_first_str = std::str::from_utf8(&server_first_raw) - .map_err(|_| ProxyError::BackendAuth("invalid UTF-8 in server-first".into()))?; - - let sf = parse_server_first(server_first_str)?; - - if !sf.server_nonce.starts_with(&client_nonce) || sf.server_nonce.len() <= client_nonce.len() { - return Err(ProxyError::BackendAuth( - "server nonce doesn't start with client nonce".into(), - )); - } - - // Step 3: ClientFinal - let client_final_without_proof = format!("c=biws,r={}", sf.server_nonce); - - let client_proof = compute_client_proof( - password, - client_first_bare.as_bytes(), - server_first_str.as_bytes(), - client_final_without_proof.as_bytes(), - &sf.salt, - sf.iteration_count, - )?; - - let client_final_message = format!(",{}", client_proof); - let full_client_final = format!("{}{}", client_final_without_proof, client_final_message); - - send_sasl_response(stream, full_client_final.as_bytes()).await?; - - // Step 4: ServerSignature — server sends AuthenticationSASLFinal (type 12) - let server_final_raw = read_sasl_final(stream).await?; - let server_final_str = std::str::from_utf8(&server_final_raw) - .map_err(|_| ProxyError::BackendAuth("invalid UTF-8 in server-final".into()))?; - - if let Some(server_sig) = server_final_str.strip_prefix("v=") { - let expected_sig = compute_server_signature( - password, - client_first_bare.as_bytes(), - server_first_str.as_bytes(), - client_final_without_proof.as_bytes(), - &sf.salt, - sf.iteration_count, - )?; - - if server_sig != expected_sig { - return Err(ProxyError::BackendAuth("server signature mismatch".into())); - } - } else if let Some(err_msg) = server_final_str.strip_prefix("e=") { - return Err(ProxyError::BackendAuth(format!( - "server error: {}", - err_msg - ))); - } else { - return Err(ProxyError::BackendAuth(format!( - "unexpected server-final: {}", - server_final_str - ))); - } - - tracing::debug!(username = %username, "SCRAM authentication successful"); - Ok(()) -} - -// ─── SCRAM internals ───────────────────────────────────────────────────────── - -struct ServerFirst { - salt: Vec, - iteration_count: u32, - server_nonce: String, -} - -fn parse_server_first(s: &str) -> Result { - let mut salt = None; - let mut iter_count = None; - let mut server_nonce = None; - - // The nonce in r= may contain commas, so split from the right: after the final ,s= or ,i= - if let Some(r_pos) = s.find("r=") { - // Extract r= value: from "r=" up to the last ",s=" or ",i=" - let after_r = &s[r_pos + 2..]; - let end = after_r - .rfind(",s=") - .or(after_r.rfind(",i=")) - .unwrap_or(after_r.len()); - server_nonce = Some(after_r[..end].to_string()); - } - - // Remaining attrs: s= and i= (after the nonce) - if let Some(rest) = s.split(",s=").nth(1) { - let parts: Vec<&str> = rest.splitn(2, ",i=").collect(); - salt = Some(base64_decode(parts[0]).map_err(|e| ProxyError::BackendAuth(e.to_string()))?); - if parts.len() > 1 { - iter_count = Some( - parts[1] - .parse() - .map_err(|_| ProxyError::BackendAuth("invalid iteration count".into()))?, - ); - } - } - // Handle i= without s= (fallback) - if iter_count.is_none() { - if let Some(rest) = s.split(",i=").nth(1) { - iter_count = Some( - rest.parse() - .map_err(|_| ProxyError::BackendAuth("invalid iteration count".into()))?, - ); - } - } - - Ok(ServerFirst { - salt: salt.ok_or_else(|| ProxyError::BackendAuth("missing salt".into()))?, - iteration_count: iter_count - .ok_or_else(|| ProxyError::BackendAuth("missing iteration count".into()))?, - server_nonce: server_nonce - .ok_or_else(|| ProxyError::BackendAuth("missing server nonce".into()))?, - }) -} - -fn compute_client_proof( - password: &str, - client_first_bare: &[u8], - server_first: &[u8], - client_final_message_without_proof: &[u8], - salt: &[u8], - iteration_count: u32, -) -> Result { - let normalized_password = normalize_password(password); - let salted_password = hi(&normalized_password, salt, iteration_count)?; - - let client_key = hmac_sign(&salted_password, b"Client Key"); - let stored_key = sha256_hash(&client_key); - - let auth_message: Vec = join_bytes(&[ - client_first_bare, - server_first, - client_final_message_without_proof, - ]); - - let client_signature = hmac_sign(&stored_key, &auth_message); - - let mut client_proof = vec![0u8; client_key.len()]; - for i in 0..client_key.len() { - client_proof[i] = client_key[i] ^ client_signature[i]; - } - - Ok(base64_encode(&client_proof)) -} - -fn compute_server_signature( - password: &str, - client_first_bare: &[u8], - server_first: &[u8], - client_final_message_without_proof: &[u8], - salt: &[u8], - iteration_count: u32, -) -> Result { - let normalized_password = normalize_password(password); - let salted_password = hi(&normalized_password, salt, iteration_count)?; - - let server_key = hmac_sign(&salted_password, b"Server Key"); - - let auth_message: Vec = join_bytes(&[ - client_first_bare, - server_first, - client_final_message_without_proof, - ]); - - let server_signature = hmac_sign(&server_key, &auth_message); - - Ok(base64_encode(&server_signature)) -} - -/// PBKDF2-HMAC-SHA256 key derivation. -fn hi(password: &[u8], salt: &[u8], iterations: u32) -> Result, ProxyError> { - let mut result = vec![0u8; 32]; - let mut u = vec![0u8; 32]; - - let mut mac = - HmacSha256::new_from_slice(password).map_err(|e| ProxyError::BackendAuth(e.to_string()))?; - mac.update(salt); - mac.update(&1u32.to_be_bytes()); - u.copy_from_slice(&mac.finalize().into_bytes()); - result.copy_from_slice(&u); // XOR in U1 (result is zeroed, so copy = XOR) - - for _ in 2..=iterations { - let mut mac = HmacSha256::new_from_slice(password) - .map_err(|e| ProxyError::BackendAuth(e.to_string()))?; - mac.update(&u); - u.copy_from_slice(&mac.finalize().into_bytes()); - for i in 0..32 { - result[i] ^= u[i]; - } - } - - Ok(result) -} - -fn hmac_sign(key: &[u8], data: &[u8]) -> Vec { - let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts any key size"); - mac.update(data); - mac.finalize().into_bytes().to_vec() -} - -fn sha256_hash(data: &[u8]) -> Vec { - use sha2::Digest; - Sha256::new().chain_update(data).finalize().to_vec() -} - -fn normalize_password(password: &str) -> Vec { - // RFC 5802: normalize according to SASLprep profile - // For simplicity, we use the password as-is (Postgres SCRAM accepts this) - password.as_bytes().to_vec() -} - -/// Join byte slices with a separator. -fn join_bytes(parts: &[&[u8]]) -> Vec { - let sep = b','; - let total: usize = parts.iter().map(|p| p.len()).sum::() + parts.len().saturating_sub(1); - let mut result = Vec::with_capacity(total); - for (i, part) in parts.iter().enumerate() { - if i > 0 { - result.push(sep); - } - result.extend_from_slice(part); - } - result -} - -fn base64_encode(data: &[u8]) -> String { - BASE64.encode(data) -} - -fn base64_decode(s: &str) -> Result, &'static str> { - BASE64 - .decode(s) - .map_err(|_| "invalid base64") -} - -// ─── Wire framing helpers ──────────────────────────────────────────────────── - -async fn send_sasl_initial_response( - stream: &mut S, - mechanism: &str, - initial_response: &[u8], -) -> Result<(), ProxyError> -where - S: AsyncWriteExt + Unpin, -{ - let mut buf = Vec::new(); - buf.push(b'p'); - // length = 4 (self) + mechanism + null + Int32(initial_response_len) + initial_response - let response_len = 4 + mechanism.len() + 1 + 4 + initial_response.len(); - buf.extend_from_slice(&(response_len as u32).to_be_bytes()); - buf.extend_from_slice(mechanism.as_bytes()); - buf.push(0); - if initial_response.is_empty() { - buf.extend_from_slice(&(-1i32).to_be_bytes()); - } else { - buf.extend_from_slice(&(initial_response.len() as i32).to_be_bytes()); - buf.extend_from_slice(initial_response); - } - stream.write_all(&buf).await?; - stream.flush().await?; - Ok(()) -} - -async fn send_sasl_response(stream: &mut S, data: &[u8]) -> Result<(), ProxyError> -where - S: AsyncWriteExt + Unpin, -{ - let mut buf = Vec::new(); - buf.push(b'p'); - let len: u32 = 4 + data.len() as u32; - buf.extend_from_slice(&len.to_be_bytes()); - buf.extend_from_slice(data); - stream.write_all(&buf).await?; - stream.flush().await?; - Ok(()) -} - -async fn read_sasl_auth_message( - stream: &mut S, - expected_type: u32, -) -> Result, ProxyError> -where - S: AsyncReadExt + Unpin, -{ - let mut type_buf = [0u8; 1]; - stream.read_exact(&mut type_buf).await?; - if type_buf[0] != b'R' { - return Err(ProxyError::ProtocolViolation(format!( - "expected SASL auth message (R), got {:02x}", - type_buf[0] - ))); - } - let mut len_buf = [0u8; 4]; - stream.read_exact(&mut len_buf).await?; - let len = u32::from_be_bytes(len_buf); - let mut body = vec![0u8; (len - 4) as usize]; - stream.read_exact(&mut body).await?; - let auth_type = u32::from_be_bytes([body[0], body[1], body[2], body[3]]); - if auth_type != expected_type { - return Err(ProxyError::ProtocolViolation(format!( - "expected SASL auth type {}, got {}", - expected_type, auth_type - ))); - } - Ok(body[4..].to_vec()) -} - -/// Read AuthenticationSASLContinue (type 11) from the backend. -async fn read_sasl_continue(stream: &mut S) -> Result, ProxyError> -where - S: AsyncReadExt + Unpin, -{ - read_sasl_auth_message(stream, 11).await -} - -/// Read AuthenticationSASLFinal (type 12) from the backend. -async fn read_sasl_final(stream: &mut S) -> Result, ProxyError> -where - S: AsyncReadExt + Unpin, -{ - read_sasl_auth_message(stream, 12).await -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_base64_roundtrip() { - let cases: &[&[u8]] = &[b"a", b"ab", b"abc", b"Hello, World!", b"\x00\xff\xfe\xfd"]; - for case in cases { - let encoded = base64_encode(case); - let decoded = base64_decode(&encoded).unwrap(); - assert_eq!(decoded.as_slice(), *case, "roundtrip failed for {:?}", case); - } - // Also test empty explicitly - let encoded = base64_encode(b""); - assert_eq!(base64_decode(&encoded).unwrap(), b""); - } - - #[test] - fn test_parse_server_first_valid() { - let s = - "r=fyko+d2lbbFgONe9WqKkE2qtVdgo,+5qdLY9Rw=,s=QSXCRQD6Yt6AS+kWSMEpqhGkg5e/klE+,i=4096"; - let sf = parse_server_first(s).unwrap(); - assert_eq!(sf.server_nonce, "fyko+d2lbbFgONe9WqKkE2qtVdgo,+5qdLY9Rw="); - assert_eq!(sf.iteration_count, 4096); - } - - #[test] - fn test_parse_server_first_missing_fields() { - assert!(parse_server_first("r=nonce").is_err()); - assert!(parse_server_first("s=salt").is_err()); - assert!(parse_server_first("i=4096").is_err()); - } - - #[test] - fn test_sha256_hash_known() { - use sha2::Digest; - // Known SHA256 of "test" - let result = sha256_hash(b"test"); - let expected = Sha256::digest(b"test"); - assert_eq!(result, expected.to_vec()); - } - - #[test] - fn test_hmac_sign_deterministic() { - let sig1 = hmac_sign(b"key", b"data"); - let sig2 = hmac_sign(b"key", b"data"); - assert_eq!(sig1, sig2); - assert_eq!(sig1.len(), 32); // SHA256 output = 32 bytes - } - - #[test] - fn test_hi_includes_first_iteration() { - // Regression test: hi() must XOR in U1 (the first HMAC iteration). - // Without this, only iterations 2..n are XORed, giving a wrong result. - // Test vector from RFC 6070 / test vectors for PBKDF2-SHA256: - // password="password", salt="salt", c=4096, DK=120fb06c... - let password = b"password"; - let salt = b"salt"; - let iterations = 4096; - let result = hi(password, salt, iterations).expect("hi should succeed"); - // RFC 6070 test vector: PBKDF2-SHA256("password", "salt", 4096) - assert_eq!( - &result[..4], - &[0xc5, 0xe4, 0x78, 0xd5], - "hi() must XOR in U1 (first iteration); without it the result is wrong" - ); - } - -} diff --git a/src/wire.rs b/src/wire.rs index 82e72d3..2c902ed 100644 --- a/src/wire.rs +++ b/src/wire.rs @@ -253,6 +253,92 @@ where Ok(()) } +/// Write a SASL InitialResponse ('p' message) carrying the chosen mechanism +/// name and the client-first message body. +/// +/// Wire format: `'p' | u32 length | mechanism bytes | 0u8 | i32 initial_len | initial bytes` +/// (initial_len = -1 if `initial` is empty.) +pub async fn write_sasl_initial_response( + stream: &mut S, + mechanism: &str, + initial: &[u8], +) -> Result<(), ProxyError> +where + S: AsyncWriteExt + Unpin, +{ + let mut buf = Vec::with_capacity(1 + 4 + mechanism.len() + 1 + 4 + initial.len()); + buf.push(b'p'); + let body_len: u32 = (4 + mechanism.len() + 1 + 4 + initial.len()) as u32; + buf.extend_from_slice(&body_len.to_be_bytes()); + buf.extend_from_slice(mechanism.as_bytes()); + buf.push(0); + if initial.is_empty() { + buf.extend_from_slice(&(-1i32).to_be_bytes()); + } else { + buf.extend_from_slice(&(initial.len() as i32).to_be_bytes()); + buf.extend_from_slice(initial); + } + stream.write_all(&buf).await?; + stream.flush().await?; + Ok(()) +} + +/// Write a SASL Response ('p' message) carrying the client-final message body. +/// +/// Wire format: `'p' | u32 length | data bytes` +pub async fn write_sasl_response(stream: &mut S, data: &[u8]) -> Result<(), ProxyError> +where + S: AsyncWriteExt + Unpin, +{ + let mut buf = Vec::with_capacity(1 + 4 + data.len()); + buf.push(b'p'); + let body_len: u32 = (4 + data.len()) as u32; + buf.extend_from_slice(&body_len.to_be_bytes()); + buf.extend_from_slice(data); + stream.write_all(&buf).await?; + stream.flush().await?; + Ok(()) +} + +/// Read a backend AuthenticationRequest ('R') message and assert its auth_type +/// matches `expected_type`. Returns the body bytes after the auth_type field. +/// +/// - type=11 (AuthenticationSASLContinue) +/// - type=12 (AuthenticationSASLFinal) +/// - type=0 (AuthenticationOk) +pub async fn read_sasl_auth_message( + stream: &mut S, + expected_type: u32, +) -> Result, ProxyError> +where + S: AsyncReadExt + Unpin, +{ + let mut tag = [0u8; 1]; + stream.read_exact(&mut tag).await?; + if tag[0] != b'R' { + return Err(ProxyError::ProtocolViolation(format!( + "expected AuthenticationRequest (R), got {:02x}", + tag[0] + ))); + } + let len = read_message_length(stream).await?; + let mut body = vec![0u8; (len - 4) as usize]; + stream.read_exact(&mut body).await?; + if body.len() < 4 { + return Err(ProxyError::ProtocolViolation( + "AuthenticationRequest body too short".into(), + )); + } + let auth_type = u32::from_be_bytes([body[0], body[1], body[2], body[3]]); + if auth_type != expected_type { + return Err(ProxyError::ProtocolViolation(format!( + "expected SASL auth type {}, got {}", + expected_type, auth_type + ))); + } + Ok(body[4..].to_vec()) +} + /// Write a Query message ('Q'). pub async fn write_query(stream: &mut S, sql: &str) -> Result<(), ProxyError> where @@ -437,3 +523,90 @@ fn split_null(slice: &[u8]) -> (&[u8], &[u8]) { None => (slice, &[][..]), } } + +#[cfg(test)] +mod tests { + use super::*; + + // Helper for tests — wraps a byte slice and implements Unpin + AsyncReadExt + struct SliceReader<'a>(&'a [u8]); + impl tokio::io::AsyncRead for SliceReader<'_> { + fn poll_read( + mut self: std::pin::Pin<&mut Self>, + _: &mut std::task::Context<'_>, + buf: &mut tokio::io::ReadBuf<'_>, + ) -> std::task::Poll> { + let remaining = self.0.len(); + if remaining == 0 { + return std::task::Poll::Ready(Ok(())); + } + let chunk = std::cmp::min(buf.remaining(), remaining); + buf.put_slice(&self.0[..chunk]); + self.0 = &self.0[chunk..]; + std::task::Poll::Ready(Ok(())) + } + } + + #[tokio::test] + async fn write_sasl_initial_response_layout() { + let mut buf = Vec::new(); + write_sasl_initial_response(&mut buf, "SCRAM-SHA-256", b"hello") + .await + .unwrap(); + // Expected: 'p' | u32(4 + 13 + 1 + 4 + 5) = 27 | "SCRAM-SHA-256" | 0 | i32(5) | "hello" + assert_eq!(buf[0], b'p'); + assert_eq!(&buf[1..5], &27u32.to_be_bytes()); + assert_eq!(&buf[5..18], b"SCRAM-SHA-256"); + assert_eq!(buf[18], 0); + assert_eq!(&buf[19..23], &5i32.to_be_bytes()); + assert_eq!(&buf[23..28], b"hello"); + } + + #[tokio::test] + async fn write_sasl_response_layout() { + let mut buf = Vec::new(); + write_sasl_response(&mut buf, b"final-msg").await.unwrap(); + // Expected: 'p' | u32(4 + 9) = 13 | "final-msg" + assert_eq!(buf[0], b'p'); + assert_eq!(&buf[1..5], &13u32.to_be_bytes()); + assert_eq!(&buf[5..14], b"final-msg"); + } + + #[tokio::test] + async fn write_sasl_initial_response_empty_initial() { + let mut buf = Vec::new(); + write_sasl_initial_response(&mut buf, "SCRAM-SHA-256", b"") + .await + .unwrap(); + // Expected: 'p' | u32(4 + 13 + 1 + 4) = 22 | "SCRAM-SHA-256" | 0 | i32(-1) + assert_eq!(buf[0], b'p'); + assert_eq!(&buf[1..5], &22u32.to_be_bytes()); + assert_eq!(&buf[5..18], b"SCRAM-SHA-256"); + assert_eq!(buf[18], 0); + assert_eq!(&buf[19..23], &(-1i32).to_be_bytes()); + assert_eq!(buf.len(), 23); + } + + #[tokio::test] + async fn read_sasl_auth_message_wrong_tag() { + // SliceReader wraps a byte slice and implements Unpin + AsyncReadExt + let result = read_sasl_auth_message(&mut SliceReader(&[b'E'][..]), 11).await; + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("expected AuthenticationRequest (R)")); + } + + #[tokio::test] + async fn read_sasl_auth_message_type_mismatch() { + // 'R' | u32(len=8) | u32(auth_type=5) — we ask for type 11, should mismatch + let result = + read_sasl_auth_message(&mut SliceReader(&[b'R', 0, 0, 0, 8, 0, 0, 0, 5][..]), 11).await; + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("expected SASL auth type 11, got 5")); + } +}