From 91df20948a300a677998587405ac3f55169e5f40 Mon Sep 17 00:00:00 2001
From: lumi <lumi@pew.im>
Date: Sat, 25 Feb 2017 00:31:43 +0100
Subject: [PATCH] support SCRAM-SHA-256

---
 src/sasl/mechanisms/mod.rs   |  2 +-
 src/sasl/mechanisms/scram.rs | 42 ++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/src/sasl/mechanisms/mod.rs b/src/sasl/mechanisms/mod.rs
index c1d4bc24..59463ac0 100644
--- a/src/sasl/mechanisms/mod.rs
+++ b/src/sasl/mechanisms/mod.rs
@@ -6,4 +6,4 @@ mod scram;
 
 pub use self::anonymous::Anonymous;
 pub use self::plain::Plain;
-pub use self::scram::{Scram, Sha1, ScramProvider};
+pub use self::scram::{Scram, Sha1, Sha256, ScramProvider};
diff --git a/src/sasl/mechanisms/scram.rs b/src/sasl/mechanisms/scram.rs
index 7b426c09..9ae04fcf 100644
--- a/src/sasl/mechanisms/scram.rs
+++ b/src/sasl/mechanisms/scram.rs
@@ -69,6 +69,29 @@ impl ScramProvider for Sha1 { // TODO: look at all these unwraps
     }
 }
 
+pub struct Sha256;
+
+impl ScramProvider for Sha256 { // TODO: look at all these unwraps
+    fn name() -> &'static str { "SCRAM-SHA-256" }
+
+    fn hash(data: &[u8]) -> Vec<u8> {
+        hash(MessageDigest::sha256(), data).unwrap()
+    }
+
+    fn hmac(data: &[u8], key: &[u8]) -> Vec<u8> {
+        let pkey = PKey::hmac(key).unwrap();
+        let mut signer = Signer::new(MessageDigest::sha256(), &pkey).unwrap();
+        signer.update(data).unwrap();
+        signer.finish().unwrap()
+    }
+
+    fn derive(data: &[u8], salt: &[u8], iterations: usize) -> Vec<u8> {
+        let mut result = vec![0; 32];
+        pbkdf2_hmac(data, salt, iterations, MessageDigest::sha256(), &mut result).unwrap();
+        result
+    }
+}
+
 enum ScramState {
     Init,
     SentInitialMessage { initial_message: Vec<u8> },
@@ -251,4 +274,23 @@ mod tests {
                   , String::from_utf8(client_final[..].to_owned()).unwrap() ); // again, depends on ordering…
         mechanism.success(&server_final[..]).unwrap();
     }
+
+    #[test]
+    fn scram_sha256_works() { // Source: RFC 7677
+        let username = "user";
+        let password = "pencil";
+        let client_nonce = "rOprNGfwEbeRWgbNEkqO";
+        let client_init = b"n,,n=user,r=rOprNGfwEbeRWgbNEkqO";
+        let server_init = b"r=rOprNGfwEbeRWgbNEkqO%hvYDpWUa2RaTCAfuxFIlj)hNlF$k0,s=W22ZaJ0SNY7soEsUEjb6gQ==,i=4096";
+        let client_final = b"c=biws,r=rOprNGfwEbeRWgbNEkqO%hvYDpWUa2RaTCAfuxFIlj)hNlF$k0,p=dHzbZapWIk4jUhN+Ute9ytag9zjfMHgsqmmiz7AndVQ=";
+        let server_final = b"v=6rriTRBi23WpRR/wtup+mMhUZUn/dB5nLTJRsjl95G4=";
+        let mut mechanism = Scram::<Sha256>::new_with_nonce(username, password, client_nonce.to_owned());
+        let init = mechanism.initial().unwrap();
+        assert_eq!( String::from_utf8(init.clone()).unwrap()
+                  , String::from_utf8(client_init[..].to_owned()).unwrap() ); // depends on ordering…
+        let resp = mechanism.response(&server_init[..]).unwrap();
+        assert_eq!( String::from_utf8(resp.clone()).unwrap()
+                  , String::from_utf8(client_final[..].to_owned()).unwrap() ); // again, depends on ordering…
+        mechanism.success(&server_final[..]).unwrap();
+    }
 }