sqlmodel_postgres/auth/
scram.rs

1//! SCRAM-SHA-256 Authentication implementation.
2
3use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
4use hmac::{Hmac, KeyInit, Mac};
5use rand::rand_core::UnwrapErr;
6use rand::{RngExt, distr::Alphanumeric, rngs::SysRng};
7use sha2::{Digest, Sha256};
8use sqlmodel_core::Error;
9use sqlmodel_core::error::{ConnectionError, ConnectionErrorKind, ProtocolError};
10use subtle::ConstantTimeEq;
11
12type HmacSha256 = Hmac<Sha256>;
13
14pub struct ScramClient {
15    username: String,
16    password: String,
17    client_nonce: String,
18
19    // State from server
20    server_nonce: Option<String>,
21    salt: Option<Vec<u8>>,
22    iterations: Option<u32>,
23
24    // Derived keys
25    salted_password: Option<[u8; 32]>,
26    auth_message: Option<String>,
27}
28
29impl ScramClient {
30    pub fn new(username: &str, password: &str) -> Self {
31        // Use SysRng (OS entropy; renamed from OsRng in rand 0.10) for
32        // cryptographically secure nonce generation.
33        // 32 characters of alphanumeric provides ~190 bits of entropy.
34        // SysRng is fallible-only (TryRng) in rand 0.10; UnwrapErr adapts it to
35        // the infallible Rng surface `sample_iter` needs.
36        let client_nonce: String = UnwrapErr(SysRng)
37            .sample_iter(&Alphanumeric)
38            .take(32)
39            .map(char::from)
40            .collect();
41
42        Self {
43            username: username.to_string(),
44            password: password.to_string(),
45            client_nonce,
46            server_nonce: None,
47            salt: None,
48            iterations: None,
49            salted_password: None,
50            auth_message: None,
51        }
52    }
53
54    /// Generate client-first message
55    pub fn client_first(&self) -> Vec<u8> {
56        // gs2-header: "n,," (no channel binding, no authzid)
57        // client-first-message-bare: "n=<user>,r=<nonce>"
58        // Note: SCRAM requires strict handling of "," in usernames but Postgres usually forbids it or requires escaping.
59        // For now we assume standard username.
60        format!("n,,n={},r={}", self.username, self.client_nonce).into_bytes()
61    }
62
63    /// Process server-first message and generate client-final
64    #[allow(clippy::result_large_err)]
65    pub fn process_server_first(&mut self, data: &[u8]) -> Result<Vec<u8>, Error> {
66        let msg = std::str::from_utf8(data)
67            .map_err(|e| protocol_error(format!("Invalid UTF-8 in SASL continue: {}", e)))?;
68
69        // Parse server-first: r=<nonce>,s=<salt>,i=<iterations>
70        let mut combined_nonce = None;
71        let mut salt = None;
72        let mut iterations = None;
73
74        for part in msg.split(',') {
75            if let Some(value) = part.strip_prefix("r=") {
76                combined_nonce = Some(value.to_string());
77            } else if let Some(value) = part.strip_prefix("s=") {
78                salt = Some(
79                    BASE64
80                        .decode(value)
81                        .map_err(|e| protocol_error(format!("Invalid base64 salt: {}", e)))?,
82                );
83            } else if let Some(value) = part.strip_prefix("i=") {
84                iterations = Some(
85                    value
86                        .parse()
87                        .map_err(|e| protocol_error(format!("Invalid iterations: {}", e)))?,
88                );
89            }
90        }
91
92        let combined_nonce = combined_nonce.ok_or_else(|| protocol_error("Missing nonce"))?;
93        let salt = salt.ok_or_else(|| protocol_error("Missing salt"))?;
94        let iterations = iterations.ok_or_else(|| protocol_error("Missing iterations"))?;
95
96        // Verify nonce starts with our client nonce
97        if !combined_nonce.starts_with(&self.client_nonce) {
98            return Err(protocol_error("Invalid server nonce"));
99        }
100
101        // Derive salted password using PBKDF2
102        let mut salted_password = [0u8; 32];
103        pbkdf2::pbkdf2::<HmacSha256>(
104            self.password.as_bytes(),
105            &salt,
106            iterations,
107            &mut salted_password,
108        )
109        .map_err(|e| protocol_error(format!("PBKDF2 failed: {}", e)))?;
110
111        // Build auth message
112        let client_first_bare = format!("n={},r={}", self.username, self.client_nonce);
113        let client_final_without_proof = format!("c=biws,r={}", combined_nonce); // biws = base64("n,,")
114        let auth_message = format!(
115            "{},{},{}",
116            client_first_bare, msg, client_final_without_proof
117        );
118
119        // Calculate client proof
120        let client_key = hmac_sha256(&salted_password, b"Client Key")?;
121        let stored_key = sha256(&client_key);
122        let client_signature = hmac_sha256(&stored_key, auth_message.as_bytes())?;
123
124        let client_proof: Vec<u8> = client_key
125            .iter()
126            .zip(client_signature.iter())
127            .map(|(a, b)| a ^ b)
128            .collect();
129
130        // Store for verification
131        self.server_nonce = Some(combined_nonce.clone());
132        self.salted_password = Some(salted_password);
133        self.auth_message = Some(auth_message);
134        self.salt = Some(salt);
135        self.iterations = Some(iterations);
136
137        // Build client-final message
138        let client_final = format!(
139            "c=biws,r={},p={}",
140            combined_nonce,
141            BASE64.encode(&client_proof)
142        );
143
144        Ok(client_final.into_bytes())
145    }
146
147    /// Verify server-final message
148    #[allow(clippy::result_large_err)]
149    pub fn verify_server_final(&self, data: &[u8]) -> Result<(), Error> {
150        let msg = std::str::from_utf8(data)
151            .map_err(|e| protocol_error(format!("Invalid UTF-8 in SASL final: {}", e)))?;
152
153        let server_signature_b64 = msg
154            .strip_prefix("v=")
155            .ok_or_else(|| protocol_error("Invalid server-final format"))?;
156
157        let server_signature = BASE64
158            .decode(server_signature_b64)
159            .map_err(|e| protocol_error(format!("Invalid base64 server signature: {}", e)))?;
160
161        // Calculate expected server signature
162        let salted_password = self
163            .salted_password
164            .as_ref()
165            .ok_or_else(|| protocol_error("Missing salted password state"))?;
166        let auth_message = self
167            .auth_message
168            .as_ref()
169            .ok_or_else(|| protocol_error("Missing auth message state"))?;
170
171        let server_key = hmac_sha256(salted_password, b"Server Key")?;
172        let expected_signature = hmac_sha256(&server_key, auth_message.as_bytes())?;
173
174        // Use constant-time comparison to prevent timing attacks.
175        // An attacker observing response times could otherwise recover
176        // the expected signature byte-by-byte.
177        if server_signature.ct_eq(&expected_signature).into() {
178            Ok(())
179        } else {
180            Err(auth_error("Server signature mismatch"))
181        }
182    }
183}
184
185// Helpers
186
187fn protocol_error(msg: impl Into<String>) -> Error {
188    Error::Protocol(ProtocolError {
189        message: msg.into(),
190        raw_data: None,
191        source: None,
192    })
193}
194
195fn auth_error(msg: impl Into<String>) -> Error {
196    Error::Connection(ConnectionError {
197        kind: ConnectionErrorKind::Authentication,
198        message: msg.into(),
199        source: None,
200    })
201}
202
203#[allow(clippy::result_large_err)]
204fn hmac_sha256(key: &[u8], data: &[u8]) -> Result<[u8; 32], Error> {
205    let mut mac = HmacSha256::new_from_slice(key)
206        .map_err(|e| protocol_error(format!("HMAC init failed: {}", e)))?;
207    mac.update(data);
208    let result = mac.finalize();
209    let bytes = result.into_bytes();
210    Ok(bytes.into())
211}
212
213fn sha256(data: &[u8]) -> [u8; 32] {
214    let mut hasher = Sha256::new();
215    hasher.update(data);
216    hasher.finalize().into()
217}
sqlmodel_postgres/auth/scram.rs

sqlmodel_postgres/auth/
scram.rs
