totp-rs 6.0.0

use crate::Algorithm;
use core::fmt;

/// Represents a token generated by [`Totp`](crate::Totp).
/// This can be thought of as a string with a few notable difference:
///
/// * It is _only_ stack allocated
/// * Comparison is done in constant time regardless of the number of digits in the token
/// * Formatting is lazily evaluated
///
/// Since it implements [`Display`](core::fmt::Display), it can be directly
/// used in formatting just like a string:
///
/// ```
/// # use totp_rs::Token;
/// # extern crate alloc;
/// # use alloc::{format, string::String};
/// # #[cfg(feature = "alloc")]
/// # fn _foo(token: Token) {
/// let text: String = format!("{}", token);
/// # }
/// ```
///
/// While [`Token`] implements [`Eq`], it's strongly recommended that you rely on
/// [`check`](crate::Totp::check) instead, as that will take skew into account.
///
/// Note that while [`Token`] is stack-allocated and only contains plain data, it
/// does _not_ implement [`Copy`].
/// This is to preserve compatibility with the `zeroize` feature which will add
/// a non-trivial [`Drop`] implementation.
#[derive(Clone, Eq)]
#[cfg_attr(feature = "zeroize", derive(zeroize::Zeroize, zeroize::ZeroizeOnDrop))]
pub struct Token {
    #[cfg_attr(feature = "zeroize", zeroize(skip))]
    algorithm: Algorithm,
    #[cfg_attr(feature = "zeroize", zeroize(skip))]
    digits: u8,
    /// # Invariants
    ///
    /// * `self.value == self.value & 0x7FFF_FFFF;`
    /// * `self.value == self.value % Self::modulo(self.algorithm, self.digits);`
    value: u32,
}

impl Token {
    pub(crate) const fn new(algorithm: Algorithm, digits: u8, value: u32) -> Self {
        Self {
            algorithm,
            digits,
            value: (value & 0x7FFF_FFFF) % Self::modulo(algorithm, digits),
        }
    }

    pub(crate) const fn from_signature(algorithm: Algorithm, digits: u8, signature: &[u8]) -> Self {
        let last = *signature.last().unwrap();
        let offset = (last & 0xF) as usize;
        let value = u32::from_be_bytes([
            signature[offset],
            signature[offset + 1],
            signature[offset + 2],
            signature[offset + 3],
        ]);

        Self::new(algorithm, digits, value)
    }

    pub(crate) const fn try_from_formatted_string(
        algorithm: Algorithm,
        digits: u8,
        string: &str,
    ) -> Option<Self> {
        if string.len() != digits as usize {
            return None;
        }

        let value = match algorithm {
            Algorithm::SHA1 | Algorithm::SHA256 | Algorithm::SHA512 => {
                // `from_str_radix` accepts a leading `+`, which is not a valid
                // token character. Reject anything but ASCII digits first, then
                // rely on the parse only to read the value and catch overflow.
                let bytes = string.as_bytes();
                let mut i = 0;
                while i < bytes.len() {
                    if !bytes[i].is_ascii_digit() {
                        return None;
                    }
                    i += 1;
                }

                match u32::from_str_radix(string, 10) {
                    Ok(value) => value,
                    Err(_) => return None,
                }
            }
            #[cfg(feature = "steam")]
            Algorithm::Steam => {
                let radix = STEAM_CHARS.len();
                let mut value = 0;
                let mut place = 1;
                let mut bytes = string.as_bytes();

                while let [byte, rest @ ..] = bytes {
                    let mut i = 0;
                    let mut digits = STEAM_CHARS;
                    let index = loop {
                        match digits {
                            [x, _rest @ ..] if *x == *byte => break i,
                            [_, rest @ ..] => {
                                i += 1;
                                digits = rest;
                            }
                            [] => return None,
                        }
                    };

                    value += index * place;
                    bytes = rest;
                    place *= radix;
                }

                value as u32
            }
        };

        Some(Self::new(algorithm, digits, value))
    }

    const fn modulo(algorithm: Algorithm, digits: u8) -> u32 {
        match algorithm {
            Algorithm::SHA1 | Algorithm::SHA256 | Algorithm::SHA512 => 10_u32.checked_pow(digits as u32)
                .expect("a `digits` value over 9 is a guaranteed corruption as 10^10 is 10_000_000_000, which does not fit in an u32."),
            #[cfg(feature = "steam")]
            Algorithm::Steam => (STEAM_CHARS.len() as u32).checked_pow(digits as u32)
                .expect("a `digits` value over 6 is a guaranteed corruption as 26^7 is 8_031_810_176, which does not fit in an u32."),
        }
    }
}

impl PartialEq for Token {
    // [`algorithm`](Token::algorithm) and [`digits`](Token::digits) are not considered
    // secret, so their comparison need not be constant time.
    fn eq(&self, other: &Self) -> bool {
        constant_time_eq::constant_time_eq_n(&self.value.to_ne_bytes(), &other.value.to_ne_bytes())
            && self.algorithm == other.algorithm
            && self.digits == other.digits
    }
}

impl fmt::Debug for Token {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        <Self as fmt::Display>::fmt(self, f)
    }
}

impl fmt::Display for Token {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.algorithm {
            Algorithm::SHA1 | Algorithm::SHA256 | Algorithm::SHA512 => write!(
                f,
                "{1:00$}",
                self.digits.into(),
                self.value % 10_u32.checked_pow(self.digits.into())
                    .expect("a `digits` value over 9 is a guaranteed corruption as 10^10 is 10_000_000_000, which does not fit in an u32."),
            ),
            #[cfg(feature = "steam")]
            Algorithm::Steam => {
                use core::fmt::Write as _;

                if self.digits >= 7 {
                    panic!("a `digits` value over 6 is a guaranteed corruption as 26^7 is 8_031_810_176, which does not fit in an u32.")
                }

                let chars = (0..self.digits).scan(self.value, |value, _| {
                    let digit = *value as usize % STEAM_CHARS.len();
                    *value /= STEAM_CHARS.len() as u32;
                    Some(char::from(STEAM_CHARS[digit]))
                });

                for c in chars {
                    f.write_char(c)?;
                }

                Ok(())
            }
        }
    }
}

/// Alphabet for Steam tokens.
#[cfg(feature = "steam")]
const STEAM_CHARS: &[u8] = b"23456789BCDFGHJKMNPQRTVWXY";

#[cfg(test)]
mod tests {
    use super::Token;
    use crate::Algorithm;

    /// While 5..=8 is typical, we test with 3 digits for brevity.
    const DIGITS: u8 = 3;
    /// Tests are written to exhaustively search under the modular arithmetic of the token.
    /// This caps the number of tokens to test just to avoid tests running for extremely long periods of time.
    const MAX_TOKEN: u32 = 1_000_000;
    /// We exhaustively test against all algorithms.
    const ALL_ALGORITHMS: &[Algorithm] = &[
        Algorithm::SHA1,
        Algorithm::SHA256,
        Algorithm::SHA512,
        #[cfg(feature = "steam")]
        Algorithm::Steam,
    ];

    /// Tests that all 5 and 6 digit tokens for all algorithms can
    /// be formatted as a string and then retrieved as the same token through parsing.
    ///
    /// Also ensures [`Display`](core::fmt::Display) and [`Debug`](core::fmt::Debug)
    /// formatting are equivalent.
    #[test]
    fn formatting_round_trip() {
        for &alg in ALL_ALGORITHMS {
            let digits = DIGITS;
            let modulo = Token::modulo(alg, digits);
            for value in 0..modulo.min(MAX_TOKEN) {
                let token = Token::new(alg, digits, value);

                let formatted = format!("{token}");
                let re_parsed = Token::try_from_formatted_string(alg, digits, &formatted);
                assert_eq!(
                    Some(&token),
                    re_parsed.as_ref(),
                    "{formatted} could not be re-parsed!"
                );

                let debug_formatted = format!("{token:?}");
                assert_eq!(
                    formatted, debug_formatted,
                    "debug and display formatting should be equivalent!"
                );
            }
        }
    }

    /// Exhaustively tests that the highest bit is irrelevant.
    #[test]
    fn highest_bit_irrelevant() {
        for &alg in ALL_ALGORITHMS {
            let digits = DIGITS;
            let modulo = Token::modulo(alg, digits);
            for value in 0..modulo.min(MAX_TOKEN) {
                let token = Token::new(alg, digits, value);
                let token_with_high_bit = Token::new(alg, digits, value | 0x8000_0000);
                let token_without_high_bit = Token::new(alg, digits, value & !0x8000_0000);

                assert_eq!(
                    token, token_with_high_bit,
                    "setting high-bit made a difference when it shouldn't!"
                );
                assert_eq!(
                    token, token_without_high_bit,
                    "resetting high-bit made a difference when it shouldn't!"
                );
            }
        }
    }

    /// Tests that the modularity of a token's value is respected.
    #[test]
    fn modular_arithmetic() {
        for &alg in ALL_ALGORITHMS {
            let digits = DIGITS;
            let modulo = Token::modulo(alg, digits);
            for value in 0..modulo.min(MAX_TOKEN) {
                let token = Token::new(alg, digits, value);
                for order in 1..5 {
                    let token_next_mod = Token::new(alg, digits, value + order * modulo);
                    assert_eq!(
                        token, token_next_mod,
                        "tokens should be equivalent under their modulo!"
                    );
                }
            }
        }
    }

    /// Tests that [`Token::from_signature`] works as expected:
    /// * Last byte used as an offset
    /// * offset..offset + 4 treated as a big-endian [`u32`]
    /// * Passed into [Token::new]
    #[test]
    fn from_signature() {
        for &alg in ALL_ALGORITHMS {
            let digits = DIGITS;
            let modulo = Token::modulo(alg, digits);
            for value in 0..modulo.min(MAX_TOKEN) {
                for offset in 0..4 {
                    let mut signature = [0; 8];
                    *signature.last_mut().unwrap() = offset as u8;
                    signature[offset..][..4].copy_from_slice(&value.to_be_bytes());

                    assert_eq!(
                        Token::new(alg, digits, value),
                        Token::from_signature(alg, digits, &signature),
                        "expected {signature:?} to be equivalent to {value}!"
                    );
                }
            }
        }
    }

    /// Ensure tokens with an invalid character-set fail to parse.
    #[test]
    fn parsing_failure() {
        let invalid_token_for_sha1 = "abc123";
        let token = Token::try_from_formatted_string(
            Algorithm::SHA1,
            invalid_token_for_sha1.len() as u8,
            invalid_token_for_sha1,
        );
        assert_eq!(token, None);
    }

    /// `u32::from_str_radix` accepts a leading sign and would otherwise treat
    /// e.g. "+8020" as equivalent to the token "08020". Non-digit characters,
    /// including a leading `+`, must be rejected even when the length matches.
    #[test]
    fn parsing_rejects_non_digits() {
        for non_digit in ["+8020", "-8020", " 8020", "80 20"] {
            assert_eq!(
                Token::try_from_formatted_string(Algorithm::SHA1, 5, non_digit),
                None,
                "expected \"{non_digit}\" to be rejected"
            );
        }

        // A canonical all-digit token of the same length still parses.
        assert!(Token::try_from_formatted_string(Algorithm::SHA1, 5, "08020").is_some());
    }

    /// "22222", causing [`check`](crate::Totp::check) to accept it.
    #[test]
    #[cfg(feature = "steam")]
    fn steam_parsing_rejects_chars_outside_alphabet() {
        for invalid in ["AAAAA", "2222A", "ZZZZZ", "2345I"] {
            assert_eq!(
                Token::try_from_formatted_string(Algorithm::Steam, 5, invalid),
                None,
                "expected \"{invalid}\" to be rejected"
            );
        }

        // A token made only of alphabet characters still parses.
        assert!(Token::try_from_formatted_string(Algorithm::Steam, 5, "22222").is_some());
    }
}