1 files changed, 38 insertions, 39 deletions

diff --git a/src/crypto.rs b/src/crypto.rs
index 4413663..d681e86 100644
--- a/src/crypto.rs
+++ b/src/crypto.rs
@@ -17,7 +17,7 @@ use sha2::Sha256;
 
 use crate::{
     serde::as_hex,
-    types::{AccountResetID, HawkKey, KeyFetchID, PasswordChangeID, SessionID},
+    types::{AccountResetID, HawkKey, KeyFetchID, PasswordChangeID, SecretKey, SessionID},
 };
 
 const NAMESPACE: &[u8] = b"identity.mozilla.com/picl/v1/";
@@ -28,6 +28,14 @@ pub fn random_bytes<const N: usize>() -> [u8; N] {
     result
 }
 
+fn xor<const N: usize>(l: &[u8; N], r: &[u8; N]) -> [u8; N] {
+    let mut result = *l;
+    for (a, b) in result.iter_mut().zip(r.iter()) {
+        *a ^= b;
+    }
+    result
+}
+
 #[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(try_from = "String", into = "String")]
 pub struct SecretBytes<const N: usize>(pub [u8; N]);
@@ -38,16 +46,6 @@ impl<const N: usize> Debug for SecretBytes<N> {
     }
 }
 
-impl<const N: usize> SecretBytes<N> {
-    fn xor(&self, other: &Self) -> Self {
-        let mut result = self.clone();
-        for (a, b) in result.0.iter_mut().zip(other.0.iter()) {
-            *a ^= b;
-        }
-        result
-    }
-}
-
 impl<const N: usize> From<SecretBytes<N>> for String {
     fn from(sb: SecretBytes<N>) -> Self {
         hex::encode(&sb.0)
@@ -136,11 +134,11 @@ impl<const N: usize> SecretBytes<N> {
 
 #[derive(Debug, Deserialize, Serialize)]
 #[serde(transparent)]
-pub struct AuthPW {
+pub(crate) struct AuthPW {
     pw: SecretBytes<32>,
 }
 
-pub struct StretchedPW {
+pub(crate) struct StretchedPW {
     pw: Output,
 }
 
@@ -161,16 +159,16 @@ impl StretchedPW {
         raw.into()
     }
 
-    fn wrap_wrap_key(&self) -> SecretBytes<32> {
+    fn wrap_wrap_key(&self) -> [u8; 32] {
         from_hkdf(self.pw.as_bytes(), &[NAMESPACE, b"wrapwrapKey"])
     }
 
-    pub fn decrypt_wwkb(&self, wwkb: &SecretBytes<32>) -> SecretBytes<32> {
-        wwkb.xor(&self.wrap_wrap_key())
+    pub fn decrypt_wwkb(&self, wwkb: &SecretKey) -> SecretKey {
+        SecretKey(xor(&wwkb.0, &self.wrap_wrap_key()))
     }
 
-    pub fn rewrap_wkb(&self, wkb: &SecretBytes<32>) -> SecretBytes<32> {
-        wkb.xor(&self.wrap_wrap_key())
+    pub fn rewrap_wkb(&self, wkb: &SecretKey) -> SecretKey {
+        SecretKey(xor(&wkb.0, &self.wrap_wrap_key()))
     }
 }
 
@@ -240,25 +238,25 @@ impl KeyFetchReq {
     }
 }
 
-pub struct KeyFetchResp {
+pub(crate) struct KeyFetchResp {
     resp_hmac_key: SecretBytes<32>,
     resp_xor_key: SecretBytes<64>,
 }
 
 impl KeyFetchResp {
     pub fn wrap_keys(&self, keys: &KeyBundle) -> WrappedKeyBundle {
-        let ciphertext = self.resp_xor_key.xor(&keys.to_bytes());
+        let ciphertext = xor(&self.resp_xor_key.0, &keys.to_bytes().0);
         #[allow(clippy::unwrap_used)]
         let mut hmac = Hmac::<Sha256>::new_from_slice(&self.resp_hmac_key.0).unwrap();
-        hmac.update(&ciphertext.0);
+        hmac.update(&ciphertext);
         let hmac = *hmac.finalize().into_bytes().as_ref();
         WrappedKeyBundle { ciphertext, hmac }
     }
 }
 
-pub struct KeyBundle {
-    pub ka: SecretBytes<32>,
-    pub wrap_kb: SecretBytes<32>,
+pub(crate) struct KeyBundle {
+    pub ka: SecretKey,
+    pub wrap_kb: SecretKey,
 }
 
 impl KeyBundle {
@@ -272,14 +270,14 @@ impl KeyBundle {
 
 #[derive(Debug)]
 pub struct WrappedKeyBundle {
-    ciphertext: SecretBytes<64>,
+    ciphertext: [u8; 64],
     hmac: [u8; 32],
 }
 
 impl WrappedKeyBundle {
     pub fn to_bytes(&self) -> [u8; 96] {
         let mut result = [0; 96];
-        result[0..64].copy_from_slice(&self.ciphertext.0);
+        result[0..64].copy_from_slice(&self.ciphertext);
         result[64..].copy_from_slice(&self.hmac);
         result
     }
@@ -349,19 +347,16 @@ mod test {
     use hex_literal::hex;
     use password_hash::{Output, SaltString};
 
-    use crate::crypto::{
-        AccountResetReq, AccountResetToken, KeyBundle, KeyFetchReq, KeyFetchToken,
-        PasswordChangeReq, PasswordChangeToken, SessionCredentials, SessionToken,
+    use crate::{
+        crypto::{
+            AccountResetReq, AccountResetToken, KeyBundle, KeyFetchReq, KeyFetchToken,
+            PasswordChangeReq, PasswordChangeToken, SessionCredentials, SessionToken,
+        },
+        types::SecretKey,
     };
 
     use super::{AuthPW, SecretBytes};
 
-    macro_rules! shex {
-        ( $s: literal ) => {
-            SecretBytes(hex!($s))
-        };
-    }
-
     #[test]
     fn test_derive_session() {
         let creds = SessionCredentials::derive_from(&SessionToken(hex!(
@@ -409,8 +404,12 @@ mod test {
         );
 
         let bundle = KeyBundle {
-            ka: shex!("2021222324252627 28292a2b2c2d2e2f 3031323334353637 38393a3b3c3d3e3f"),
-            wrap_kb: shex!("7effe354abecbcb2 34a8dfc2d7644b4a d339b525589738f2 d27341bb8622ecd8"),
+            ka: SecretKey(hex!(
+                "2021222324252627 28292a2b2c2d2e2f 3031323334353637 38393a3b3c3d3e3f"
+            )),
+            wrap_kb: SecretKey(hex!(
+                "7effe354abecbcb2 34a8dfc2d7644b4a d339b525589738f2 d27341bb8622ecd8"
+            )),
         };
         assert_eq!(
             bundle.to_bytes().0,
@@ -422,7 +421,7 @@ mod test {
 
         let wrapped = resp.wrap_keys(&bundle);
         assert_eq!(
-            wrapped.ciphertext.0,
+            wrapped.ciphertext,
             hex!(
                 "ee5c58845c7c9412 b11bbd20920c2fdd d83c33c9cd2c2de2 d66b222613364636
                  fc7e59d854d599f1 0e212801de3a47c3 4333f3b838ee3471 e0f285649c332bbb"
@@ -469,7 +468,7 @@ mod test {
         );
 
         assert_eq!(
-            stretched.wrap_wrap_key().0,
+            stretched.wrap_wrap_key(),
             hex!("3ebea117efa9faf5 7ce195899b290505 8368e7760cc26ea5 8a2a1be0da7fb287")
         );
         Ok(())