From 0267383803cc735c734b85fb05dfe49281be909b Mon Sep 17 00:00:00 2001
From: Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
Date: Wed, 25 Sep 2024 21:53:50 +0200
Subject: [PATCH] sasl: Improve frame parsing

All keys only have a single char, so avoid allocating a whole String for
them.  Also move to a BTreeMap as this is more efficient on such keys.
---
 sasl/src/client/mechanisms/scram.rs |  8 ++++----
 sasl/src/common/mod.rs              | 10 ++++++----
 sasl/src/server/mechanisms/scram.rs |  6 +++---
 3 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/sasl/src/client/mechanisms/scram.rs b/sasl/src/client/mechanisms/scram.rs
index 5c29188c76fae79b95a6ac2ad13b6d687541bc6d..05f7874478ecec6cae4a6ed4e705fe6e665e06d9 100644
--- a/sasl/src/client/mechanisms/scram.rs
+++ b/sasl/src/client/mechanisms/scram.rs
@@ -153,9 +153,9 @@ impl<S: ScramProvider> Mechanism for Scram<S> {
             } => {
                 let frame =
                     parse_frame(challenge).map_err(|_| MechanismError::CannotDecodeChallenge)?;
-                let server_nonce = frame.get("r");
-                let salt = frame.get("s").and_then(|v| Base64.decode(v).ok());
-                let iterations = frame.get("i").and_then(|v| v.parse().ok());
+                let server_nonce = frame.get(&'r');
+                let salt = frame.get(&'s').and_then(|v| Base64.decode(v).ok());
+                let iterations = frame.get(&'i').and_then(|v| v.parse().ok());
                 let server_nonce = server_nonce.ok_or(MechanismError::NoServerNonce)?;
                 let salt = salt.ok_or(MechanismError::NoServerSalt)?;
                 let iterations = iterations.ok_or(MechanismError::NoServerIterations)?;
@@ -206,7 +206,7 @@ impl<S: ScramProvider> Mechanism for Scram<S> {
             ScramState::GotServerData {
                 ref server_signature,
             } => {
-                if let Some(sig) = frame.get("v").and_then(|v| Base64.decode(v).ok()) {
+                if let Some(sig) = frame.get(&'v').and_then(|v| Base64.decode(v).ok()) {
                     if sig == *server_signature {
                         Ok(())
                     } else {
diff --git a/sasl/src/common/mod.rs b/sasl/src/common/mod.rs
index 7447b4ba05d2f420d54a38be02c431c395e50757..bcf6543087a28838121fb17ea6e577201ee84bef 100644
--- a/sasl/src/common/mod.rs
+++ b/sasl/src/common/mod.rs
@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::BTreeMap;
 use std::string::FromUtf8Error;
 
 #[cfg(feature = "scram")]
@@ -141,15 +141,17 @@ pub fn xor(a: &[u8], b: &[u8]) -> Vec<u8> {
 }
 
 #[doc(hidden)]
-pub fn parse_frame(frame: &[u8]) -> Result<HashMap<String, String>, FromUtf8Error> {
+pub fn parse_frame(frame: &[u8]) -> Result<BTreeMap<char, String>, FromUtf8Error> {
     let inner = String::from_utf8(frame.to_owned())?;
-    let mut ret = HashMap::new();
+    let mut ret = BTreeMap::new();
     for s in inner.split(',') {
         let mut tmp = s.splitn(2, '=');
         let key = tmp.next();
         let val = tmp.next();
         if let (Some(k), Some(v)) = (key, val) {
-            ret.insert(k.to_owned(), v.to_owned());
+            if let Some(k) = k.chars().next() {
+                ret.insert(k, v.to_owned());
+            }
         }
     }
     Ok(ret)
diff --git a/sasl/src/server/mechanisms/scram.rs b/sasl/src/server/mechanisms/scram.rs
index af0e410b82f6653542c874ad597f78100ee6d601..7bc48edf81c9aefcd4307fdd65b73f8318fc1c8d 100644
--- a/sasl/src/server/mechanisms/scram.rs
+++ b/sasl/src/server/mechanisms/scram.rs
@@ -108,9 +108,9 @@ where
                 }
                 let frame =
                     parse_frame(&rest).map_err(|_| MechanismError::CannotDecodeInitialMessage)?;
-                let username = frame.get("n").ok_or(MechanismError::NoUsername)?;
+                let username = frame.get(&'n').ok_or(MechanismError::NoUsername)?;
                 let identity = Identity::Username(username.to_owned());
-                let client_nonce = frame.get("r").ok_or(MechanismError::NoNonce)?;
+                let client_nonce = frame.get(&'r').ok_or(MechanismError::NoNonce)?;
                 let mut server_nonce = String::new();
                 server_nonce += client_nonce;
                 server_nonce +=
@@ -162,7 +162,7 @@ where
                 let stored_key = S::hash(&client_key);
                 let client_signature = S::hmac(&auth_message, &stored_key)?;
                 let client_proof = xor(&client_key, &client_signature);
-                let sent_proof = frame.get("p").ok_or(MechanismError::NoProof)?;
+                let sent_proof = frame.get(&'p').ok_or(MechanismError::NoProof)?;
                 let sent_proof = Base64
                     .decode(sent_proof)
                     .map_err(|_| MechanismError::CannotDecodeProof)?;