1use crypto::{Hasher, hmac::Hmac, sha2::Sha256};
2
3use crate::{
4 error::{PgError, Result},
5 protocol::{base64_decode, base64_encode},
6};
7
8fn hi(password: &str, salt: &[u8], iterations: u32) -> [u8; 32] {
9 let pw = password.as_bytes();
10 let mut u = hmac_sha256(pw, &[salt, &[0, 0, 0, 1]].concat());
11 let result = [0u8; 32];
12 let mut result = [0u8; 32];
13 result.copy_from_slice(u.as_ref());
14
15 for _ in 1..iterations {
16 u = hmac_sha256(pw, u.as_ref());
17 for (a, b) in result.iter_mut().zip(u.as_ref()) {
18 *a ^= b;
19 }
20 }
21
22 result
23}
24
25fn hmac_sha256(key: &[u8], data: &[u8]) -> crypto::Hash {
26 let mut mac = Hmac::<Sha256>::new(key);
27 mac.update(data);
28 mac.finalize()
29}
30
31pub(crate) struct ScramClient {
32 client_first_message_bare: String,
33 client_nonce: String,
34 password: String,
35 server_first_message: Option<String>,
36 client_final_without_proof: Option<String>,
37 salted_password: Option<[u8; 32]>,
38}
39
40impl ScramClient {
41 pub fn new(username: &str, password: &str) -> Self {
42 let raw: [u8; 24] = rand::random();
43 let client_nonce = hex::encode(&raw);
44
45 ScramClient {
46 client_first_message_bare: format!("n={},r={}", username, client_nonce),
47 client_nonce,
48 password: password.to_string(),
49 server_first_message: None,
50 client_final_without_proof: None,
51 salted_password: None,
52 }
53 }
54
55 pub fn client_first_message(&self) -> &str {
56 &self.client_first_message_bare
57 }
58
59 pub fn parse_server_first_message(&mut self, data: &[u8]) -> Result<()> {
60 let msg = std::str::from_utf8(data).map_err(|_| PgError::Auth("invalid utf-8 in server-first".into()))?;
61 self.server_first_message = Some(msg.to_string());
62
63 let mut combined_nonce = None;
64 let mut salt_b64 = None;
65 let mut iterations = None;
66
67 for part in msg.split(',') {
68 if let Some(val) = part.strip_prefix("r=") {
69 combined_nonce = Some(val.to_string());
70 } else if let Some(val) = part.strip_prefix("s=") {
71 salt_b64 = Some(val.to_string());
72 } else if let Some(val) = part.strip_prefix("i=") {
73 iterations = Some(
74 val.parse::<u32>()
75 .map_err(|_| PgError::Auth("invalid iteration count".into()))?,
76 );
77 }
78 }
79
80 let combined_nonce = combined_nonce.ok_or_else(|| PgError::Auth("missing nonce in server-first".into()))?;
81 let salt_b64 = salt_b64.ok_or_else(|| PgError::Auth("missing salt in server-first".into()))?;
82 let iterations = iterations.ok_or_else(|| PgError::Auth("missing iterations in server-first".into()))?;
83
84 if !combined_nonce.starts_with(&self.client_nonce) {
85 return Err(PgError::Auth("server nonce doesn't start with client nonce".into()));
86 }
87
88 let salt = base64_decode(&salt_b64).map_err(|e| PgError::Auth(format!("invalid base64 salt: {}", e)))?;
89
90 let salted_password = hi(&self.password, &salt, iterations);
91 self.salted_password = Some(salted_password);
92 self.client_final_without_proof = Some(format!("c=biws,r={}", combined_nonce));
93
94 Ok(())
95 }
96
97 pub fn build_client_final_message(&self) -> Vec<u8> {
98 let sp = self.salted_password.as_ref().expect("salted password not computed");
99
100 let client_key = hmac_sha256(sp, b"Client Key");
101 let client_key_bytes: &[u8] = client_key.as_ref();
102
103 let mut hasher = crypto::sha2::Sha256::new();
104 hasher.update(client_key_bytes);
105 let stored_key = hasher.sum();
106 let stored_key_bytes: &[u8] = stored_key.as_ref();
107
108 let server_first = self.server_first_message.as_ref().expect("no server-first message");
109 let cfnop = self
110 .client_final_without_proof
111 .as_ref()
112 .expect("no client-final-without-proof");
113
114 let auth_message = format!("{},{},{}", self.client_first_message_bare, server_first, cfnop);
115
116 let client_signature = hmac_sha256(stored_key_bytes, auth_message.as_bytes());
117
118 let mut client_proof = [0u8; 32];
119 client_proof.copy_from_slice(client_key_bytes);
120 for (a, b) in client_proof.iter_mut().zip(client_signature.as_ref()) {
121 *a ^= b;
122 }
123
124 let client_proof_b64 = base64_encode(&client_proof);
125 let client_final = format!("{},p={}", cfnop, client_proof_b64);
126 client_final.into_bytes()
127 }
128
129 pub fn parse_server_final_message(&self, data: &[u8]) -> Result<()> {
130 let msg = std::str::from_utf8(data).map_err(|_| PgError::Auth("invalid utf-8 in server-final".into()))?;
131
132 let mut server_sig_b64 = None;
133 for part in msg.split(',') {
134 if let Some(val) = part.strip_prefix("v=") {
135 server_sig_b64 = Some(val.to_string());
136 } else if let Some(val) = part.strip_prefix("e=") {
137 return Err(PgError::Auth(format!("server returned auth error: {}", val)));
138 }
139 }
140
141 let server_sig_b64 = server_sig_b64.ok_or_else(|| PgError::Auth("missing server signature".into()))?;
142
143 let sp = self.salted_password.as_ref().expect("salted password not computed");
144 let server_key = hmac_sha256(sp, b"Server Key");
145
146 let server_first = self.server_first_message.as_ref().expect("no server-first message");
147 let cfnop = self
148 .client_final_without_proof
149 .as_ref()
150 .expect("no client-final-without-proof");
151
152 let auth_message = format!("{},{},{}", self.client_first_message_bare, server_first, cfnop);
153
154 let expected_signature = hmac_sha256(server_key.as_ref(), auth_message.as_bytes());
155 let expected_b64 = base64_encode(expected_signature.as_ref());
156
157 if expected_b64 != server_sig_b64 {
158 return Err(PgError::Auth("server signature mismatch".into()));
159 }
160
161 Ok(())
162 }
163}