From 8610f30f42e487f79dfc2a2bf31ad97bcdf829ff Mon Sep 17 00:00:00 2001
From: Hinton <hinton@users.noreply.github.com>
Date: Mon, 4 Dec 2023 10:47:17 +0100
Subject: [PATCH 1/9] Initial TOTP implementation

---
 Cargo.lock                                    |   7 +
 crates/bitwarden-uniffi/src/vault/mod.rs      |  13 +-
 crates/bitwarden/Cargo.toml                   |   1 +
 .../bitwarden/src/mobile/vault/client_totp.rs |   7 +-
 crates/bitwarden/src/vault/totp.rs            | 238 +++++++++++++++++-
 5 files changed, 257 insertions(+), 9 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index cc507ac33..2600d5f35 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -332,6 +332,7 @@ dependencies = [
  "bitwarden-api-identity",
  "cbc",
  "chrono",
+ "data-encoding",
  "getrandom 0.2.11",
  "hkdf",
  "hmac",
@@ -929,6 +930,12 @@ dependencies = [
  "syn 2.0.39",
 ]
 
+[[package]]
+name = "data-encoding"
+version = "2.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7e962a19be5cfc3f3bf6dd8f61eb50107f356ad6270fbb3ed41476571db78be5"
+
 [[package]]
 name = "deadpool"
 version = "0.9.5"
diff --git a/crates/bitwarden-uniffi/src/vault/mod.rs b/crates/bitwarden-uniffi/src/vault/mod.rs
index 79b3e5835..472ad2b53 100644
--- a/crates/bitwarden-uniffi/src/vault/mod.rs
+++ b/crates/bitwarden-uniffi/src/vault/mod.rs
@@ -3,7 +3,7 @@ use std::sync::Arc;
 use bitwarden::vault::TotpResponse;
 use chrono::{DateTime, Utc};
 
-use crate::Client;
+use crate::{error::Result, Client};
 
 pub mod ciphers;
 pub mod collections;
@@ -47,13 +47,18 @@ impl ClientVault {
     /// - A base32 encoded string
     /// - OTP Auth URI
     /// - Steam URI
-    pub async fn generate_totp(&self, key: String, time: Option<DateTime<Utc>>) -> TotpResponse {
-        self.0
+    pub async fn generate_totp(
+        &self,
+        key: String,
+        time: Option<DateTime<Utc>>,
+    ) -> Result<TotpResponse> {
+        Ok(self
+            .0
              .0
             .read()
             .await
             .vault()
             .generate_totp(key, time)
-            .await
+            .await?)
     }
 }
diff --git a/crates/bitwarden/Cargo.toml b/crates/bitwarden/Cargo.toml
index 308f98850..318e97737 100644
--- a/crates/bitwarden/Cargo.toml
+++ b/crates/bitwarden/Cargo.toml
@@ -33,6 +33,7 @@ chrono = { version = ">=0.4.26, <0.5", features = [
     "serde",
     "std",
 ], default-features = false }
+data-encoding = ">=2.5.0, <3.0"
 # We don't use this directly (it's used by rand), but we need it here to enable WASM support
 getrandom = { version = ">=0.2.9, <0.3", features = ["js"] }
 hkdf = ">=0.12.3, <0.13"
diff --git a/crates/bitwarden/src/mobile/vault/client_totp.rs b/crates/bitwarden/src/mobile/vault/client_totp.rs
index 97eb243af..903015e3e 100644
--- a/crates/bitwarden/src/mobile/vault/client_totp.rs
+++ b/crates/bitwarden/src/mobile/vault/client_totp.rs
@@ -1,5 +1,6 @@
 use chrono::{DateTime, Utc};
 
+use crate::error::Result;
 use crate::vault::{generate_totp, TotpResponse};
 
 use super::client_vault::ClientVault;
@@ -12,7 +13,11 @@ impl<'a> ClientVault<'a> {
     /// - OTP Auth URI
     /// - Steam URI
     ///
-    pub async fn generate_totp(&'a self, key: String, time: Option<DateTime<Utc>>) -> TotpResponse {
+    pub async fn generate_totp(
+        &'a self,
+        key: String,
+        time: Option<DateTime<Utc>>,
+    ) -> Result<TotpResponse> {
         generate_totp(key, time).await
     }
 }
diff --git a/crates/bitwarden/src/vault/totp.rs b/crates/bitwarden/src/vault/totp.rs
index 7e701f92e..f4b78dd7d 100644
--- a/crates/bitwarden/src/vault/totp.rs
+++ b/crates/bitwarden/src/vault/totp.rs
@@ -1,7 +1,19 @@
+use std::collections::HashMap;
+
+use crate::error::{Error, Result};
 use chrono::{DateTime, Utc};
+use data_encoding::BASE32;
+use hmac::{Hmac, Mac};
+use reqwest::Url;
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
+type HmacSha1 = Hmac<sha1::Sha1>;
+type HmacSha256 = Hmac<sha2::Sha256>;
+type HmacSha512 = Hmac<sha2::Sha512>;
+
+const STEAM_CHARS: &str = "23456789BCDFGHJKMNPQRTVWXY";
+
 #[derive(Serialize, Deserialize, Debug, JsonSchema)]
 #[serde(rename_all = "camelCase", deny_unknown_fields)]
 #[cfg_attr(feature = "mobile", derive(uniffi::Record))]
@@ -12,6 +24,33 @@ pub struct TotpResponse {
     pub period: u32,
 }
 
+#[derive(Clone, Copy, Debug)]
+enum TotpAlgorithm {
+    Sha1,
+    Sha256,
+    Sha512,
+    Steam,
+}
+
+#[derive(Debug)]
+struct TotpParams {
+    algorithm: TotpAlgorithm,
+    digits: u32,
+    period: u32,
+    secret: String,
+}
+
+impl Default for TotpParams {
+    fn default() -> Self {
+        Self {
+            algorithm: TotpAlgorithm::Sha1,
+            digits: 6,
+            period: 30,
+            secret: "".to_string(),
+        }
+    }
+}
+
 /// Generate a OATH or RFC 6238 TOTP code from a provided key.
 ///
 /// <https://datatracker.ietf.org/doc/html/rfc6238>
@@ -22,9 +61,200 @@ pub struct TotpResponse {
 /// - Steam URI
 ///
 /// Supports providing an optional time, and defaults to current system time if none is provided.
-pub async fn generate_totp(_key: String, _time: Option<DateTime<Utc>>) -> TotpResponse {
-    TotpResponse {
-        code: "000 000".to_string(),
-        period: 30,
+pub async fn generate_totp(key: String, time: Option<DateTime<Utc>>) -> Result<TotpResponse> {
+    let params = get_params(key)?;
+
+    // TODO: Should we swap the expected time to timestamp?
+    let time = time.unwrap_or_else(Utc::now);
+    print!("{:?}", params);
+
+    let t = time.timestamp() / params.period as i64;
+    let secret = BASE32.decode(params.secret.as_ref()).map_err(|e| {
+        println!("{:?}", e);
+        Error::Internal("Unable to decode secret")
+    })?;
+
+    let hash = derive_hash(params.algorithm, &secret, t.to_be_bytes().as_ref())?;
+    let binary = derive_binary(hash);
+
+    let otp = if let TotpAlgorithm::Steam = params.algorithm {
+        derive_steam_otp(binary, params.digits)
+    } else {
+        let otp = binary % 10_u32.pow(params.digits);
+        format!("{1:00$}", params.digits as usize, otp)
+    };
+
+    Ok(TotpResponse {
+        code: otp,
+        period: params.period,
+    })
+}
+
+/// Derive the Steam OTP from the hash with the given number of digits.
+fn derive_steam_otp(binary: u32, digits: u32) -> String {
+    let mut otp = String::new();
+
+    let mut full_code = binary & 0x7fffffff;
+    for _ in 0..digits {
+        otp.push(
+            STEAM_CHARS
+                .chars()
+                .nth(full_code as usize % STEAM_CHARS.len())
+                .expect("Should always be within range"),
+        );
+        full_code /= STEAM_CHARS.len() as u32;
+    }
+
+    otp
+}
+
+/// Parses the provided key and returns the corresponding `TotpParams`.
+///
+/// Key can be either:
+/// - A base32 encoded string
+/// - OTP Auth URI
+/// - Steam URI
+fn get_params(key: String) -> Result<TotpParams> {
+    let params = if key.starts_with("otpauth://") {
+        let url = Url::parse(&key).map_err(|_| Error::Internal("Unable to parse URL"))?;
+        let parts: HashMap<_, _> = url.query_pairs().collect();
+
+        let defaults = TotpParams::default();
+
+        TotpParams {
+            algorithm: parts
+                .get("algorithm")
+                .and_then(|v| match v.to_uppercase().as_ref() {
+                    "SHA1" => Some(TotpAlgorithm::Sha1),
+                    "SHA256" => Some(TotpAlgorithm::Sha256),
+                    "SHA512" => Some(TotpAlgorithm::Sha512),
+                    _ => None,
+                })
+                .unwrap_or(defaults.algorithm),
+            digits: parts
+                .get("digits")
+                .and_then(|v| v.parse().ok())
+                .map(|v: u32| v.clamp(0, 10))
+                .unwrap_or(defaults.digits),
+            period: parts
+                .get("period")
+                .and_then(|v| v.parse().ok())
+                .map(|v: u32| v.max(1))
+                .unwrap_or(defaults.period),
+            secret: parts
+                .get("secret")
+                .map(|v| v.to_string())
+                .unwrap_or(defaults.secret),
+        }
+    } else if key.starts_with("steam://") {
+        TotpParams {
+            algorithm: TotpAlgorithm::Steam,
+            digits: 5,
+            secret: key
+                .strip_prefix("steam://")
+                .expect("Prefix is defined")
+                .to_string(),
+            ..TotpParams::default()
+        }
+    } else {
+        TotpParams {
+            secret: key,
+            ..TotpParams::default()
+        }
+    };
+
+    Ok(params)
+}
+
+/// Derive the OTP from the hash with the given number of digits.
+fn derive_binary(hash: Vec<u8>) -> u32 {
+    let offset = (hash.last().unwrap_or(&0) & 15) as usize;
+
+    ((hash[offset] & 127) as u32) << 24
+        | (hash[offset + 1] as u32) << 16
+        | (hash[offset + 2] as u32) << 8
+        | hash[offset + 3] as u32
+}
+
+impl From<aes::cipher::InvalidLength> for Error {
+    fn from(_: aes::cipher::InvalidLength) -> Self {
+        Error::Internal("Invalid length")
+    }
+}
+
+// Derive the HMAC hash for the given algorithm
+fn derive_hash(algorithm: TotpAlgorithm, key: &[u8], time: &[u8]) -> Result<Vec<u8>> {
+    fn compute_digest<D: Mac>(mut digest: D, time: &[u8]) -> Vec<u8> {
+        digest.update(time);
+        digest.finalize().into_bytes().to_vec()
+    }
+
+    Ok(match algorithm {
+        TotpAlgorithm::Sha1 => compute_digest(HmacSha1::new_from_slice(key)?, time),
+        TotpAlgorithm::Sha256 => compute_digest(HmacSha256::new_from_slice(key)?, time),
+        TotpAlgorithm::Sha512 => compute_digest(HmacSha512::new_from_slice(key)?, time),
+        TotpAlgorithm::Steam => compute_digest(HmacSha1::new_from_slice(key)?, time),
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use chrono::Utc;
+
+    #[tokio::test]
+    async fn test_generate_totp() {
+        let key = "WQIQ25BRKZYCJVYP".to_string();
+        let time = Some(
+            DateTime::parse_from_rfc3339("2023-01-01T00:00:00.000Z")
+                .unwrap()
+                .with_timezone(&Utc),
+        );
+        let response = generate_totp(key, time).await.unwrap();
+
+        assert_eq!(response.code, "194506".to_string());
+        assert_eq!(response.period, 30);
+    }
+
+    #[tokio::test]
+    async fn test_generate_otpauth() {
+        let key = "otpauth://totp/test-account?secret=WQIQ25BRKZYCJVYP".to_string();
+        let time = Some(
+            DateTime::parse_from_rfc3339("2023-01-01T00:00:00.000Z")
+                .unwrap()
+                .with_timezone(&Utc),
+        );
+        let response = generate_totp(key, time).await.unwrap();
+
+        assert_eq!(response.code, "194506".to_string());
+        assert_eq!(response.period, 30);
+    }
+
+    #[tokio::test]
+    async fn test_generate_otpauth_period() {
+        let key = "otpauth://totp/test-account?secret=WQIQ25BRKZYCJVYP&period=60".to_string();
+        let time = Some(
+            DateTime::parse_from_rfc3339("2023-01-01T00:00:00.000Z")
+                .unwrap()
+                .with_timezone(&Utc),
+        );
+        let response = generate_totp(key, time).await.unwrap();
+
+        assert_eq!(response.code, "730364".to_string());
+        assert_eq!(response.period, 60);
+    }
+
+    #[tokio::test]
+    async fn test_generate_steam() {
+        let key = "steam://HXDMVJECJJWSRB3HWIZR4IFUGFTMXBOZ".to_string();
+        let time = Some(
+            DateTime::parse_from_rfc3339("2023-01-01T00:00:00.000Z")
+                .unwrap()
+                .with_timezone(&Utc),
+        );
+        let response = generate_totp(key, time).await.unwrap();
+
+        assert_eq!(response.code, "7W6CJ".to_string());
+        assert_eq!(response.period, 30);
     }
 }

From 79edca4b9ef6589993375e1f01572bb0d8a8be63 Mon Sep 17 00:00:00 2001
From: Hinton <hinton@users.noreply.github.com>
Date: Mon, 4 Dec 2023 11:03:28 +0100
Subject: [PATCH 2/9] Cleanup the code

---
 crates/bitwarden/src/vault/totp.rs | 248 ++++++++++++++++-------------
 1 file changed, 133 insertions(+), 115 deletions(-)

diff --git a/crates/bitwarden/src/vault/totp.rs b/crates/bitwarden/src/vault/totp.rs
index f4b78dd7d..686caf1ff 100644
--- a/crates/bitwarden/src/vault/totp.rs
+++ b/crates/bitwarden/src/vault/totp.rs
@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, str::FromStr};
 
 use crate::error::{Error, Result};
 use chrono::{DateTime, Utc};
@@ -24,26 +24,70 @@ pub struct TotpResponse {
     pub period: u32,
 }
 
+/// Generate a OATH or RFC 6238 TOTP code from a provided key.
+///
+/// <https://datatracker.ietf.org/doc/html/rfc6238>
+///
+/// Key can be either:
+/// - A base32 encoded string
+/// - OTP Auth URI
+/// - Steam URI
+///
+/// Supports providing an optional time, and defaults to current system time if none is provided.
+///
+/// Arguments:
+/// - `key` - The key to generate the TOTP code from
+/// - `time` - The time in UTC to generate the TOTP code for, defaults to current system time
+pub async fn generate_totp(key: String, time: Option<DateTime<Utc>>) -> Result<TotpResponse> {
+    let params: Totp = key.parse()?;
+
+    let time = time.unwrap_or_else(Utc::now);
+
+    let otp = params.derive_otp(time.timestamp())?;
+
+    Ok(TotpResponse {
+        code: otp,
+        period: params.period,
+    })
+}
+
 #[derive(Clone, Copy, Debug)]
-enum TotpAlgorithm {
+enum Algorithm {
     Sha1,
     Sha256,
     Sha512,
     Steam,
 }
 
+impl Algorithm {
+    // Derive the HMAC hash for the given algorithm
+    fn derive_hash(&self, key: &[u8], time: &[u8]) -> Result<Vec<u8>> {
+        fn compute_digest<D: Mac>(mut digest: D, time: &[u8]) -> Vec<u8> {
+            digest.update(time);
+            digest.finalize().into_bytes().to_vec()
+        }
+
+        Ok(match self {
+            Algorithm::Sha1 => compute_digest(HmacSha1::new_from_slice(key)?, time),
+            Algorithm::Sha256 => compute_digest(HmacSha256::new_from_slice(key)?, time),
+            Algorithm::Sha512 => compute_digest(HmacSha512::new_from_slice(key)?, time),
+            Algorithm::Steam => compute_digest(HmacSha1::new_from_slice(key)?, time),
+        })
+    }
+}
+
 #[derive(Debug)]
-struct TotpParams {
-    algorithm: TotpAlgorithm,
+struct Totp {
+    algorithm: Algorithm,
     digits: u32,
     period: u32,
     secret: String,
 }
 
-impl Default for TotpParams {
+impl Default for Totp {
     fn default() -> Self {
         Self {
-            algorithm: TotpAlgorithm::Sha1,
+            algorithm: Algorithm::Sha1,
             digits: 6,
             period: 30,
             secret: "".to_string(),
@@ -51,43 +95,89 @@ impl Default for TotpParams {
     }
 }
 
-/// Generate a OATH or RFC 6238 TOTP code from a provided key.
-///
-/// <https://datatracker.ietf.org/doc/html/rfc6238>
-///
-/// Key can be either:
-/// - A base32 encoded string
-/// - OTP Auth URI
-/// - Steam URI
-///
-/// Supports providing an optional time, and defaults to current system time if none is provided.
-pub async fn generate_totp(key: String, time: Option<DateTime<Utc>>) -> Result<TotpResponse> {
-    let params = get_params(key)?;
-
-    // TODO: Should we swap the expected time to timestamp?
-    let time = time.unwrap_or_else(Utc::now);
-    print!("{:?}", params);
-
-    let t = time.timestamp() / params.period as i64;
-    let secret = BASE32.decode(params.secret.as_ref()).map_err(|e| {
-        println!("{:?}", e);
-        Error::Internal("Unable to decode secret")
-    })?;
-
-    let hash = derive_hash(params.algorithm, &secret, t.to_be_bytes().as_ref())?;
-    let binary = derive_binary(hash);
-
-    let otp = if let TotpAlgorithm::Steam = params.algorithm {
-        derive_steam_otp(binary, params.digits)
-    } else {
-        let otp = binary % 10_u32.pow(params.digits);
-        format!("{1:00$}", params.digits as usize, otp)
-    };
+impl Totp {
+    fn derive_otp(&self, time: i64) -> Result<String> {
+        let time = time / self.period as i64;
+
+        let secret = BASE32.decode(self.secret.as_ref()).map_err(|e| {
+            println!("{:?}", e);
+            Error::Internal("Unable to decode secret")
+        })?;
+
+        let hash = self
+            .algorithm
+            .derive_hash(&secret, time.to_be_bytes().as_ref())?;
+        let binary = derive_binary(hash);
+
+        Ok(if let Algorithm::Steam = self.algorithm {
+            derive_steam_otp(binary, self.digits)
+        } else {
+            let otp = binary % 10_u32.pow(self.digits);
+            format!("{1:00$}", self.digits as usize, otp)
+        })
+    }
+}
 
-    Ok(TotpResponse {
-        code: otp,
-        period: params.period,
-    })
+impl FromStr for Totp {
+    type Err = Error;
+
+    /// Parses the provided key and returns the corresponding `TotpParams`.
+    ///
+    /// Key can be either:
+    /// - A base32 encoded string
+    /// - OTP Auth URI
+    /// - Steam URI
+    fn from_str(key: &str) -> Result<Self> {
+        let params = if key.starts_with("otpauth://") {
+            let url = Url::parse(key).map_err(|_| Error::Internal("Unable to parse URL"))?;
+            let parts: HashMap<_, _> = url.query_pairs().collect();
+
+            let defaults = Totp::default();
+
+            Totp {
+                algorithm: parts
+                    .get("algorithm")
+                    .and_then(|v| match v.to_uppercase().as_ref() {
+                        "SHA1" => Some(Algorithm::Sha1),
+                        "SHA256" => Some(Algorithm::Sha256),
+                        "SHA512" => Some(Algorithm::Sha512),
+                        _ => None,
+                    })
+                    .unwrap_or(defaults.algorithm),
+                digits: parts
+                    .get("digits")
+                    .and_then(|v| v.parse().ok())
+                    .map(|v: u32| v.clamp(0, 10))
+                    .unwrap_or(defaults.digits),
+                period: parts
+                    .get("period")
+                    .and_then(|v| v.parse().ok())
+                    .map(|v: u32| v.max(1))
+                    .unwrap_or(defaults.period),
+                secret: parts
+                    .get("secret")
+                    .map(|v| v.to_string())
+                    .unwrap_or(defaults.secret),
+            }
+        } else if key.starts_with("steam://") {
+            Totp {
+                algorithm: Algorithm::Steam,
+                digits: 5,
+                secret: key
+                    .strip_prefix("steam://")
+                    .expect("Prefix is defined")
+                    .to_string(),
+                ..Totp::default()
+            }
+        } else {
+            Totp {
+                secret: key.to_string(),
+                ..Totp::default()
+            }
+        };
+
+        Ok(params)
+    }
 }
 
 /// Derive the Steam OTP from the hash with the given number of digits.
@@ -108,64 +198,6 @@ fn derive_steam_otp(binary: u32, digits: u32) -> String {
     otp
 }
 
-/// Parses the provided key and returns the corresponding `TotpParams`.
-///
-/// Key can be either:
-/// - A base32 encoded string
-/// - OTP Auth URI
-/// - Steam URI
-fn get_params(key: String) -> Result<TotpParams> {
-    let params = if key.starts_with("otpauth://") {
-        let url = Url::parse(&key).map_err(|_| Error::Internal("Unable to parse URL"))?;
-        let parts: HashMap<_, _> = url.query_pairs().collect();
-
-        let defaults = TotpParams::default();
-
-        TotpParams {
-            algorithm: parts
-                .get("algorithm")
-                .and_then(|v| match v.to_uppercase().as_ref() {
-                    "SHA1" => Some(TotpAlgorithm::Sha1),
-                    "SHA256" => Some(TotpAlgorithm::Sha256),
-                    "SHA512" => Some(TotpAlgorithm::Sha512),
-                    _ => None,
-                })
-                .unwrap_or(defaults.algorithm),
-            digits: parts
-                .get("digits")
-                .and_then(|v| v.parse().ok())
-                .map(|v: u32| v.clamp(0, 10))
-                .unwrap_or(defaults.digits),
-            period: parts
-                .get("period")
-                .and_then(|v| v.parse().ok())
-                .map(|v: u32| v.max(1))
-                .unwrap_or(defaults.period),
-            secret: parts
-                .get("secret")
-                .map(|v| v.to_string())
-                .unwrap_or(defaults.secret),
-        }
-    } else if key.starts_with("steam://") {
-        TotpParams {
-            algorithm: TotpAlgorithm::Steam,
-            digits: 5,
-            secret: key
-                .strip_prefix("steam://")
-                .expect("Prefix is defined")
-                .to_string(),
-            ..TotpParams::default()
-        }
-    } else {
-        TotpParams {
-            secret: key,
-            ..TotpParams::default()
-        }
-    };
-
-    Ok(params)
-}
-
 /// Derive the OTP from the hash with the given number of digits.
 fn derive_binary(hash: Vec<u8>) -> u32 {
     let offset = (hash.last().unwrap_or(&0) & 15) as usize;
@@ -176,27 +208,13 @@ fn derive_binary(hash: Vec<u8>) -> u32 {
         | hash[offset + 3] as u32
 }
 
+/// Convert from the dependency `InvalidLength` error into this crate's `Error`.
 impl From<aes::cipher::InvalidLength> for Error {
     fn from(_: aes::cipher::InvalidLength) -> Self {
         Error::Internal("Invalid length")
     }
 }
 
-// Derive the HMAC hash for the given algorithm
-fn derive_hash(algorithm: TotpAlgorithm, key: &[u8], time: &[u8]) -> Result<Vec<u8>> {
-    fn compute_digest<D: Mac>(mut digest: D, time: &[u8]) -> Vec<u8> {
-        digest.update(time);
-        digest.finalize().into_bytes().to_vec()
-    }
-
-    Ok(match algorithm {
-        TotpAlgorithm::Sha1 => compute_digest(HmacSha1::new_from_slice(key)?, time),
-        TotpAlgorithm::Sha256 => compute_digest(HmacSha256::new_from_slice(key)?, time),
-        TotpAlgorithm::Sha512 => compute_digest(HmacSha512::new_from_slice(key)?, time),
-        TotpAlgorithm::Steam => compute_digest(HmacSha1::new_from_slice(key)?, time),
-    })
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;

From c1588391f814d4d84274e88b567a0d9938c6876b Mon Sep 17 00:00:00 2001
From: Hinton <hinton@users.noreply.github.com>
Date: Mon, 4 Dec 2023 11:04:41 +0100
Subject: [PATCH 3/9] Fix doc

---
 crates/bitwarden/src/vault/totp.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/bitwarden/src/vault/totp.rs b/crates/bitwarden/src/vault/totp.rs
index 686caf1ff..f2137842e 100644
--- a/crates/bitwarden/src/vault/totp.rs
+++ b/crates/bitwarden/src/vault/totp.rs
@@ -121,7 +121,7 @@ impl Totp {
 impl FromStr for Totp {
     type Err = Error;
 
-    /// Parses the provided key and returns the corresponding `TotpParams`.
+    /// Parses the provided key and returns the corresponding `Totp`.
     ///
     /// Key can be either:
     /// - A base32 encoded string

From c7567ac41258a1df95d8d96ebd22c5d0670b2b7b Mon Sep 17 00:00:00 2001
From: Hinton <hinton@users.noreply.github.com>
Date: Mon, 4 Dec 2023 11:06:22 +0100
Subject: [PATCH 4/9] Remove debug print

---
 crates/bitwarden/src/vault/totp.rs | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/crates/bitwarden/src/vault/totp.rs b/crates/bitwarden/src/vault/totp.rs
index f2137842e..543443b34 100644
--- a/crates/bitwarden/src/vault/totp.rs
+++ b/crates/bitwarden/src/vault/totp.rs
@@ -99,10 +99,9 @@ impl Totp {
     fn derive_otp(&self, time: i64) -> Result<String> {
         let time = time / self.period as i64;
 
-        let secret = BASE32.decode(self.secret.as_ref()).map_err(|e| {
-            println!("{:?}", e);
-            Error::Internal("Unable to decode secret")
-        })?;
+        let secret = BASE32
+            .decode(self.secret.as_ref())
+            .map_err(|e| Error::Internal("Unable to decode secret"))?;
 
         let hash = self
             .algorithm

From 7ee88e96bca9b320a6a5de0ca0aabd91a64bb674 Mon Sep 17 00:00:00 2001
From: Hinton <hinton@users.noreply.github.com>
Date: Mon, 4 Dec 2023 11:46:37 +0100
Subject: [PATCH 5/9] Fix clippy

---
 crates/bitwarden/src/vault/totp.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/bitwarden/src/vault/totp.rs b/crates/bitwarden/src/vault/totp.rs
index 543443b34..7b6f0f864 100644
--- a/crates/bitwarden/src/vault/totp.rs
+++ b/crates/bitwarden/src/vault/totp.rs
@@ -101,7 +101,7 @@ impl Totp {
 
         let secret = BASE32
             .decode(self.secret.as_ref())
-            .map_err(|e| Error::Internal("Unable to decode secret"))?;
+            .map_err(|_| Error::Internal("Unable to decode secret"))?;
 
         let hash = self
             .algorithm

From 1a3c779aaf0cd722ce9253e41d2a80d5637162fe Mon Sep 17 00:00:00 2001
From: Hinton <hinton@users.noreply.github.com>
Date: Mon, 4 Dec 2023 12:08:12 +0100
Subject: [PATCH 6/9] Add chrono clock

---
 crates/bitwarden/Cargo.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/crates/bitwarden/Cargo.toml b/crates/bitwarden/Cargo.toml
index 318e97737..a9e432c27 100644
--- a/crates/bitwarden/Cargo.toml
+++ b/crates/bitwarden/Cargo.toml
@@ -30,6 +30,7 @@ bitwarden-api-api = { path = "../bitwarden-api-api", version = "=0.2.2" }
 bitwarden-api-identity = { path = "../bitwarden-api-identity", version = "=0.2.2" }
 cbc = { version = ">=0.1.2, <0.2", features = ["alloc"] }
 chrono = { version = ">=0.4.26, <0.5", features = [
+    "clock",
     "serde",
     "std",
 ], default-features = false }

From 3a2198e8b12afc68d28db34ed29ba8adf75d0c8c Mon Sep 17 00:00:00 2001
From: Hinton <hinton@users.noreply.github.com>
Date: Mon, 4 Dec 2023 17:27:57 +0100
Subject: [PATCH 7/9] Resolve review feedback

---
 crates/bitwarden/src/vault/mod.rs  |   3 +-
 crates/bitwarden/src/vault/totp.rs | 107 +++++++++++++++--------------
 2 files changed, 59 insertions(+), 51 deletions(-)

diff --git a/crates/bitwarden/src/vault/mod.rs b/crates/bitwarden/src/vault/mod.rs
index 12910283c..ca9b12c1c 100644
--- a/crates/bitwarden/src/vault/mod.rs
+++ b/crates/bitwarden/src/vault/mod.rs
@@ -10,4 +10,5 @@ pub use collection::{Collection, CollectionView};
 pub use folder::{Folder, FolderView};
 pub use password_history::{PasswordHistory, PasswordHistoryView};
 pub use send::{Send, SendListView, SendView};
-pub use totp::{generate_totp, TotpResponse};
+pub(crate) use totp::generate_totp;
+pub use totp::TotpResponse;
diff --git a/crates/bitwarden/src/vault/totp.rs b/crates/bitwarden/src/vault/totp.rs
index 7b6f0f864..454d9e19d 100644
--- a/crates/bitwarden/src/vault/totp.rs
+++ b/crates/bitwarden/src/vault/totp.rs
@@ -38,12 +38,15 @@ pub struct TotpResponse {
 /// Arguments:
 /// - `key` - The key to generate the TOTP code from
 /// - `time` - The time in UTC to generate the TOTP code for, defaults to current system time
-pub async fn generate_totp(key: String, time: Option<DateTime<Utc>>) -> Result<TotpResponse> {
+pub(crate) async fn generate_totp(
+    key: String,
+    time: Option<DateTime<Utc>>,
+) -> Result<TotpResponse> {
     let params: Totp = key.parse()?;
 
     let time = time.unwrap_or_else(Utc::now);
 
-    let otp = params.derive_otp(time.timestamp())?;
+    let otp = params.derive_otp(time.timestamp());
 
     Ok(TotpResponse {
         code: otp,
@@ -61,18 +64,29 @@ enum Algorithm {
 
 impl Algorithm {
     // Derive the HMAC hash for the given algorithm
-    fn derive_hash(&self, key: &[u8], time: &[u8]) -> Result<Vec<u8>> {
-        fn compute_digest<D: Mac>(mut digest: D, time: &[u8]) -> Vec<u8> {
-            digest.update(time);
-            digest.finalize().into_bytes().to_vec()
+    fn derive_hash(&self, key: &[u8], time: &[u8]) -> Vec<u8> {
+        fn compute_digest<D: Mac>(digest: D, time: &[u8]) -> Vec<u8> {
+            digest.chain_update(time).finalize().into_bytes().to_vec()
         }
 
-        Ok(match self {
-            Algorithm::Sha1 => compute_digest(HmacSha1::new_from_slice(key)?, time),
-            Algorithm::Sha256 => compute_digest(HmacSha256::new_from_slice(key)?, time),
-            Algorithm::Sha512 => compute_digest(HmacSha512::new_from_slice(key)?, time),
-            Algorithm::Steam => compute_digest(HmacSha1::new_from_slice(key)?, time),
-        })
+        match self {
+            Algorithm::Sha1 => compute_digest(
+                HmacSha1::new_from_slice(key).expect("hmac new_from_slice should not fail"),
+                time,
+            ),
+            Algorithm::Sha256 => compute_digest(
+                HmacSha256::new_from_slice(key).expect("hmac new_from_slice should not fail"),
+                time,
+            ),
+            Algorithm::Sha512 => compute_digest(
+                HmacSha512::new_from_slice(key).expect("hmac new_from_slice should not fail"),
+                time,
+            ),
+            Algorithm::Steam => compute_digest(
+                HmacSha1::new_from_slice(key).expect("hmac new_from_slice should not fail"),
+                time,
+            ),
+        }
     }
 }
 
@@ -81,7 +95,7 @@ struct Totp {
     algorithm: Algorithm,
     digits: u32,
     period: u32,
-    secret: String,
+    secret: Vec<u8>,
 }
 
 impl Default for Totp {
@@ -90,30 +104,26 @@ impl Default for Totp {
             algorithm: Algorithm::Sha1,
             digits: 6,
             period: 30,
-            secret: "".to_string(),
+            secret: vec![],
         }
     }
 }
 
 impl Totp {
-    fn derive_otp(&self, time: i64) -> Result<String> {
+    fn derive_otp(&self, time: i64) -> String {
         let time = time / self.period as i64;
 
-        let secret = BASE32
-            .decode(self.secret.as_ref())
-            .map_err(|_| Error::Internal("Unable to decode secret"))?;
-
         let hash = self
             .algorithm
-            .derive_hash(&secret, time.to_be_bytes().as_ref())?;
+            .derive_hash(&self.secret, time.to_be_bytes().as_ref());
         let binary = derive_binary(hash);
 
-        Ok(if let Algorithm::Steam = self.algorithm {
+        if let Algorithm::Steam = self.algorithm {
             derive_steam_otp(binary, self.digits)
         } else {
             let otp = binary % 10_u32.pow(self.digits);
             format!("{1:00$}", self.digits as usize, otp)
-        })
+        }
     }
 }
 
@@ -127,6 +137,12 @@ impl FromStr for Totp {
     /// - OTP Auth URI
     /// - Steam URI
     fn from_str(key: &str) -> Result<Self> {
+        fn decode_secret(secret: &str) -> Result<Vec<u8>> {
+            BASE32
+                .decode(secret.as_bytes())
+                .map_err(|_| Error::Internal("Unable to decode secret"))
+        }
+
         let params = if key.starts_with("otpauth://") {
             let url = Url::parse(key).map_err(|_| Error::Internal("Unable to parse URL"))?;
             let parts: HashMap<_, _> = url.query_pairs().collect();
@@ -153,24 +169,23 @@ impl FromStr for Totp {
                     .and_then(|v| v.parse().ok())
                     .map(|v: u32| v.max(1))
                     .unwrap_or(defaults.period),
-                secret: parts
-                    .get("secret")
-                    .map(|v| v.to_string())
-                    .unwrap_or(defaults.secret),
+                secret: decode_secret(
+                    &parts
+                        .get("secret")
+                        .map(|v| v.to_string())
+                        .ok_or(Error::Internal("Missing secret in otpauth URI"))?,
+                )?,
             }
         } else if key.starts_with("steam://") {
             Totp {
                 algorithm: Algorithm::Steam,
                 digits: 5,
-                secret: key
-                    .strip_prefix("steam://")
-                    .expect("Prefix is defined")
-                    .to_string(),
+                secret: decode_secret(key.strip_prefix("steam://").expect("Prefix is defined"))?,
                 ..Totp::default()
             }
         } else {
             Totp {
-                secret: key.to_string(),
+                secret: decode_secret(key)?,
                 ..Totp::default()
             }
         };
@@ -181,20 +196,19 @@ impl FromStr for Totp {
 
 /// Derive the Steam OTP from the hash with the given number of digits.
 fn derive_steam_otp(binary: u32, digits: u32) -> String {
-    let mut otp = String::new();
-
     let mut full_code = binary & 0x7fffffff;
-    for _ in 0..digits {
-        otp.push(
-            STEAM_CHARS
-                .chars()
-                .nth(full_code as usize % STEAM_CHARS.len())
-                .expect("Should always be within range"),
-        );
-        full_code /= STEAM_CHARS.len() as u32;
-    }
 
-    otp
+    (0..digits)
+        .map(|_| {
+            let index = full_code as usize % STEAM_CHARS.len();
+            let char = STEAM_CHARS
+                .chars()
+                .nth(index)
+                .expect("Should always be within range");
+            full_code /= STEAM_CHARS.len() as u32;
+            char
+        })
+        .collect()
 }
 
 /// Derive the OTP from the hash with the given number of digits.
@@ -207,13 +221,6 @@ fn derive_binary(hash: Vec<u8>) -> u32 {
         | hash[offset + 3] as u32
 }
 
-/// Convert from the dependency `InvalidLength` error into this crate's `Error`.
-impl From<aes::cipher::InvalidLength> for Error {
-    fn from(_: aes::cipher::InvalidLength) -> Self {
-        Error::Internal("Invalid length")
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;

From 432ac90413cb3951487e2372447231ef6104f343 Mon Sep 17 00:00:00 2001
From: Hinton <hinton@users.noreply.github.com>
Date: Mon, 4 Dec 2023 17:32:54 +0100
Subject: [PATCH 8/9] Remove async

---
 crates/bitwarden-uniffi/src/vault/mod.rs      |  9 +-----
 .../bitwarden/src/mobile/vault/client_totp.rs |  4 +--
 crates/bitwarden/src/vault/totp.rs            | 29 +++++++++----------
 3 files changed, 16 insertions(+), 26 deletions(-)

diff --git a/crates/bitwarden-uniffi/src/vault/mod.rs b/crates/bitwarden-uniffi/src/vault/mod.rs
index 472ad2b53..bd5f5332c 100644
--- a/crates/bitwarden-uniffi/src/vault/mod.rs
+++ b/crates/bitwarden-uniffi/src/vault/mod.rs
@@ -52,13 +52,6 @@ impl ClientVault {
         key: String,
         time: Option<DateTime<Utc>>,
     ) -> Result<TotpResponse> {
-        Ok(self
-            .0
-             .0
-            .read()
-            .await
-            .vault()
-            .generate_totp(key, time)
-            .await?)
+        Ok(self.0 .0.read().await.vault().generate_totp(key, time)?)
     }
 }
diff --git a/crates/bitwarden/src/mobile/vault/client_totp.rs b/crates/bitwarden/src/mobile/vault/client_totp.rs
index 903015e3e..75bfd204c 100644
--- a/crates/bitwarden/src/mobile/vault/client_totp.rs
+++ b/crates/bitwarden/src/mobile/vault/client_totp.rs
@@ -13,11 +13,11 @@ impl<'a> ClientVault<'a> {
     /// - OTP Auth URI
     /// - Steam URI
     ///
-    pub async fn generate_totp(
+    pub fn generate_totp(
         &'a self,
         key: String,
         time: Option<DateTime<Utc>>,
     ) -> Result<TotpResponse> {
-        generate_totp(key, time).await
+        generate_totp(key, time)
     }
 }
diff --git a/crates/bitwarden/src/vault/totp.rs b/crates/bitwarden/src/vault/totp.rs
index 454d9e19d..382104657 100644
--- a/crates/bitwarden/src/vault/totp.rs
+++ b/crates/bitwarden/src/vault/totp.rs
@@ -38,10 +38,7 @@ pub struct TotpResponse {
 /// Arguments:
 /// - `key` - The key to generate the TOTP code from
 /// - `time` - The time in UTC to generate the TOTP code for, defaults to current system time
-pub(crate) async fn generate_totp(
-    key: String,
-    time: Option<DateTime<Utc>>,
-) -> Result<TotpResponse> {
+pub(crate) fn generate_totp(key: String, time: Option<DateTime<Utc>>) -> Result<TotpResponse> {
     let params: Totp = key.parse()?;
 
     let time = time.unwrap_or_else(Utc::now);
@@ -226,57 +223,57 @@ mod tests {
     use super::*;
     use chrono::Utc;
 
-    #[tokio::test]
-    async fn test_generate_totp() {
+    #[test]
+    fn test_generate_totp() {
         let key = "WQIQ25BRKZYCJVYP".to_string();
         let time = Some(
             DateTime::parse_from_rfc3339("2023-01-01T00:00:00.000Z")
                 .unwrap()
                 .with_timezone(&Utc),
         );
-        let response = generate_totp(key, time).await.unwrap();
+        let response = generate_totp(key, time).unwrap();
 
         assert_eq!(response.code, "194506".to_string());
         assert_eq!(response.period, 30);
     }
 
-    #[tokio::test]
-    async fn test_generate_otpauth() {
+    #[test]
+    fn test_generate_otpauth() {
         let key = "otpauth://totp/test-account?secret=WQIQ25BRKZYCJVYP".to_string();
         let time = Some(
             DateTime::parse_from_rfc3339("2023-01-01T00:00:00.000Z")
                 .unwrap()
                 .with_timezone(&Utc),
         );
-        let response = generate_totp(key, time).await.unwrap();
+        let response = generate_totp(key, time).unwrap();
 
         assert_eq!(response.code, "194506".to_string());
         assert_eq!(response.period, 30);
     }
 
-    #[tokio::test]
-    async fn test_generate_otpauth_period() {
+    #[test]
+    fn test_generate_otpauth_period() {
         let key = "otpauth://totp/test-account?secret=WQIQ25BRKZYCJVYP&period=60".to_string();
         let time = Some(
             DateTime::parse_from_rfc3339("2023-01-01T00:00:00.000Z")
                 .unwrap()
                 .with_timezone(&Utc),
         );
-        let response = generate_totp(key, time).await.unwrap();
+        let response = generate_totp(key, time).unwrap();
 
         assert_eq!(response.code, "730364".to_string());
         assert_eq!(response.period, 60);
     }
 
-    #[tokio::test]
-    async fn test_generate_steam() {
+    #[test]
+    fn test_generate_steam() {
         let key = "steam://HXDMVJECJJWSRB3HWIZR4IFUGFTMXBOZ".to_string();
         let time = Some(
             DateTime::parse_from_rfc3339("2023-01-01T00:00:00.000Z")
                 .unwrap()
                 .with_timezone(&Utc),
         );
-        let response = generate_totp(key, time).await.unwrap();
+        let response = generate_totp(key, time).unwrap();
 
         assert_eq!(response.code, "7W6CJ".to_string());
         assert_eq!(response.period, 30);

From 3bfa29b43ee3a865b2be0e917595491f7a6200ea Mon Sep 17 00:00:00 2001
From: Hinton <hinton@users.noreply.github.com>
Date: Mon, 4 Dec 2023 17:49:59 +0100
Subject: [PATCH 9/9] Remove default totp

---
 crates/bitwarden/src/vault/totp.rs | 33 ++++++++++++------------------
 1 file changed, 13 insertions(+), 20 deletions(-)

diff --git a/crates/bitwarden/src/vault/totp.rs b/crates/bitwarden/src/vault/totp.rs
index 382104657..e080a85a9 100644
--- a/crates/bitwarden/src/vault/totp.rs
+++ b/crates/bitwarden/src/vault/totp.rs
@@ -14,6 +14,10 @@ type HmacSha512 = Hmac<sha2::Sha512>;
 
 const STEAM_CHARS: &str = "23456789BCDFGHJKMNPQRTVWXY";
 
+const DEFAULT_ALGORITHM: Algorithm = Algorithm::Sha1;
+const DEFAULT_DIGITS: u32 = 6;
+const DEFAULT_PERIOD: u32 = 30;
+
 #[derive(Serialize, Deserialize, Debug, JsonSchema)]
 #[serde(rename_all = "camelCase", deny_unknown_fields)]
 #[cfg_attr(feature = "mobile", derive(uniffi::Record))]
@@ -95,17 +99,6 @@ struct Totp {
     secret: Vec<u8>,
 }
 
-impl Default for Totp {
-    fn default() -> Self {
-        Self {
-            algorithm: Algorithm::Sha1,
-            digits: 6,
-            period: 30,
-            secret: vec![],
-        }
-    }
-}
-
 impl Totp {
     fn derive_otp(&self, time: i64) -> String {
         let time = time / self.period as i64;
@@ -144,8 +137,6 @@ impl FromStr for Totp {
             let url = Url::parse(key).map_err(|_| Error::Internal("Unable to parse URL"))?;
             let parts: HashMap<_, _> = url.query_pairs().collect();
 
-            let defaults = Totp::default();
-
             Totp {
                 algorithm: parts
                     .get("algorithm")
@@ -155,17 +146,17 @@ impl FromStr for Totp {
                         "SHA512" => Some(Algorithm::Sha512),
                         _ => None,
                     })
-                    .unwrap_or(defaults.algorithm),
+                    .unwrap_or(DEFAULT_ALGORITHM),
                 digits: parts
                     .get("digits")
                     .and_then(|v| v.parse().ok())
                     .map(|v: u32| v.clamp(0, 10))
-                    .unwrap_or(defaults.digits),
+                    .unwrap_or(DEFAULT_DIGITS),
                 period: parts
                     .get("period")
                     .and_then(|v| v.parse().ok())
                     .map(|v: u32| v.max(1))
-                    .unwrap_or(defaults.period),
+                    .unwrap_or(DEFAULT_PERIOD),
                 secret: decode_secret(
                     &parts
                         .get("secret")
@@ -173,17 +164,19 @@ impl FromStr for Totp {
                         .ok_or(Error::Internal("Missing secret in otpauth URI"))?,
                 )?,
             }
-        } else if key.starts_with("steam://") {
+        } else if let Some(secret) = key.strip_prefix("steam://") {
             Totp {
                 algorithm: Algorithm::Steam,
                 digits: 5,
-                secret: decode_secret(key.strip_prefix("steam://").expect("Prefix is defined"))?,
-                ..Totp::default()
+                period: DEFAULT_PERIOD,
+                secret: decode_secret(secret)?,
             }
         } else {
             Totp {
+                algorithm: DEFAULT_ALGORITHM,
+                digits: DEFAULT_DIGITS,
+                period: DEFAULT_PERIOD,
                 secret: decode_secret(key)?,
-                ..Totp::default()
             }
         };