keyfork/crates/util/keyfork-prompt/src/validators.rs

//! Validator and parser types.

#![allow(clippy::type_complexity)]
use std::ops::RangeInclusive;

/// A trait to create validator functions.
pub trait Validator {
    /// The output of the validator function.
    type Output;

    /// The error type returned from the validator function.
    type Error;

    /// Create a validator function from the given parameters.
    fn to_fn(&self) -> Box<dyn Fn(String) -> Result<Self::Output, Box<dyn std::error::Error>>>;
}

/// A PIN could not be validated from the given input.
#[derive(thiserror::Error, Debug)]
pub enum PinError {
    /// The provided PIN was too short.
    #[error("PIN too short: {0} < {1}")]
    TooShort(usize, usize),

    /// The provided PIN was too long.
    #[error("PIN too long: {0} > {1}")]
    TooLong(usize, usize),

    /// The PIN contained invalid characters.
    #[error("PIN contained invalid characters (found {0} at position {1})")]
    InvalidCharacters(char, usize),

    /// The provided PIN had either too many repeated characters or too many sequential characters.
    #[error("PIN contained too many repeated or sequential characters")]
    InsecurePIN,
}

/// Validate that a PIN is of a certain length, matches a range of characters, and does not use
/// incrementing or decrementing sequences of characters.
///
/// The validator determines a score for a passphrase and, if the score is high enough, returns an
/// error.
///
/// Score is calculated based on:
/// * how many sequential characters are in the passphrase (ascending or descending)
/// * how many repeated characters are in the passphrase
#[derive(Default, Clone)]
pub struct SecurePinValidator {
    /// The minimum length of provided PINs.
    pub min_length: Option<usize>,

    /// The maximum length of provided PINs.
    pub max_length: Option<usize>,

    /// The characters allowed by the PIN parser.
    pub range: Option<RangeInclusive<char>>,

    /// Whether repeated characters count against the PIN.
    pub ignore_repeated_characters: bool,

    /// Whether sequential characters count against the PIN.
    pub ignore_sequential_characters: bool,
}

impl Validator for SecurePinValidator {
    type Output = String;
    type Error = PinError;

    fn to_fn(&self) -> Box<dyn Fn(String) -> Result<String, Box<dyn std::error::Error>>> {
        let min_len = self.min_length.unwrap_or(usize::MIN);
        let max_len = self.max_length.unwrap_or(usize::MAX);
        let range = self.range.clone().unwrap_or('0'..='9');
        let ignore_repeated_characters = self.ignore_repeated_characters;
        let ignore_sequential_characters = self.ignore_sequential_characters;
        Box::new(move |mut s: String| {
            s.truncate(s.trim_end().len());
            let len = s.len();
            if len < min_len {
                return Err(Box::new(PinError::TooShort(len, min_len)));
            }
            if len > max_len {
                return Err(Box::new(PinError::TooLong(len, max_len)));
            }
            let mut last_char = 0;
            let mut score = 0;
            for (index, ch) in s.chars().enumerate() {
                if !range.contains(&ch) {
                    return Err(Box::new(PinError::InvalidCharacters(ch, index)));
                }
                if [-1, 1].contains(&(ch as i32 - last_char))
                    && !ignore_sequential_characters
                {
                    score += 1;
                }
                last_char = ch as i32;
            }
            let mut chars = s.chars().collect::<Vec<_>>();
            chars.sort();
            chars.dedup();
            if !ignore_repeated_characters {
                // SAFETY: the amount of characters can't have _increased_ since deduping
                score += s.chars().count() - chars.len();
            }
            if score * 2 > s.chars().count() {
                return Err(Box::new(PinError::InsecurePIN))
            }
            Ok(s)
        })
    }
}

/// Validate that a PIN is of a certain length and matches a range of characters.
#[derive(Default, Clone)]
pub struct PinValidator {
    /// The minimum length of provided PINs.
    pub min_length: Option<usize>,

    /// The maximum length of provided PINs.
    pub max_length: Option<usize>,

    /// The characters allowed by the PIN parser.
    pub range: Option<RangeInclusive<char>>,
}

impl Validator for PinValidator {
    type Output = String;
    type Error = PinError;

    fn to_fn(&self) -> Box<dyn Fn(String) -> Result<String, Box<dyn std::error::Error>>> {
        let min_len = self.min_length.unwrap_or(usize::MIN);
        let max_len = self.max_length.unwrap_or(usize::MAX);
        let range = self.range.clone().unwrap_or('0'..='9');
        Box::new(move |mut s: String| {
            s.truncate(s.trim_end().len());
            let len = s.len();
            if len < min_len {
                return Err(Box::new(PinError::TooShort(len, min_len)));
            }
            if len > max_len {
                return Err(Box::new(PinError::TooLong(len, max_len)));
            }
            for (index, ch) in s.chars().enumerate() {
                if !range.contains(&ch) {
                    return Err(Box::new(PinError::InvalidCharacters(ch, index)));
                }
            }
            Ok(s)
        })
    }
}

#[cfg(feature = "mnemonic")]
pub mod mnemonic {
    //! Validators for mnemonics.

    use std::{ops::Range, str::FromStr};

    use super::Validator;

    use keyfork_bug::bug;
    use keyfork_mnemonic::{Mnemonic, MnemonicFromStrError};

    /// A mnemonic could not be validated from the given input.
    #[derive(thiserror::Error, Debug)]
    pub enum MnemonicValidationError {
        /// The provided mnemonic had an unexpected word length.
        #[error("Invalid word length: {0} does not match {1:?}")]
        InvalidLength(usize, WordLength),

        /// A mnemonic could not be parsed from the given mnemonic.
        #[error("{0}")]
        MnemonicFromStrError(#[from] MnemonicFromStrError),
    }

    /// The mnemonic had an unexpected word length.
    #[derive(Clone, Debug)]
    pub enum WordLength {
        /// The bounds of a mnemonic.
        Range(Range<usize>),

        /// The exact count of words.
        Count(usize),
    }

    impl WordLength {
        fn matches(&self, word_count: usize) -> bool {
            match self {
                WordLength::Range(r) => r.contains(&word_count),
                WordLength::Count(c) => c == &word_count,
            }
        }
    }

    /// Validate a mnemonic of a range of word lengths or a specific length.
    #[derive(Default, Clone)]
    pub struct MnemonicValidator {
        /// The allowed word length of provided mnemonics.
        pub word_length: Option<WordLength>,
    }

    impl Validator for MnemonicValidator {
        type Output = Mnemonic;
        type Error = MnemonicValidationError;

        fn to_fn(&self) -> Box<dyn Fn(String) -> Result<Mnemonic, Box<dyn std::error::Error>>> {
            let word_length = self.word_length.clone();
            Box::new(move |s: String| match word_length.as_ref() {
                Some(wl) => {
                    let count = s.split_whitespace().count();
                    if !wl.matches(count) {
                        return Err(Box::new(Self::Error::InvalidLength(count, wl.clone())));
                    }
                    let m = Mnemonic::from_str(&s)?;
                    Ok(m)
                }
                None => {
                    let m = Mnemonic::from_str(&s)?;
                    Ok(m)
                }
            })
        }
    }

    /// A mnemonic of a given choice of lengths. For example, a 128-bit or 256-bit BIP-0032 seed.
    #[derive(thiserror::Error, Debug)]
    pub enum MnemonicChoiceValidationError {
        /// The provided mnemonic did not match any of the valid ranges.
        #[error("Invalid word length: {0} was not in any {1:?}")]
        InvalidLength(usize, Vec<WordLength>),

        /// A mnemonic could not be parsed from the provided mnemonic.
        #[error("{0}")]
        MnemonicFromStrError(#[from] MnemonicFromStrError),
    }

    /// Validate a single mnemonic against a set of possible word lengths.
    #[derive(Clone)]
    pub struct MnemonicChoiceValidator<const N: usize> {
        /// The accepted [`WordLength`] of the mnemonic.
        pub word_lengths: [WordLength; N],
    }

    impl<const N: usize> Validator for MnemonicChoiceValidator<N> {
        type Output = Mnemonic;
        type Error = MnemonicChoiceValidationError;

        fn to_fn(&self) -> Box<dyn Fn(String) -> Result<Self::Output, Box<dyn std::error::Error>>> {
            let word_lengths = self.word_lengths.clone();
            Box::new(move |s: String| {
                let count = s.split_whitespace().count();
                for word_length in &word_lengths {
                    if word_length.matches(count) {
                        let m = Mnemonic::from_str(&s)?;
                        return Ok(m);
                    }
                }
                Err(Box::new(MnemonicChoiceValidationError::InvalidLength(
                    count,
                    word_lengths.to_vec(),
                )))
            })
        }
    }

    /// A mnemonic in the set of mnemonics could not be validated from the given inputs.
    #[derive(thiserror::Error, Debug)]
    pub enum MnemonicSetValidationError {
        /// The provided mnemonic did not have the correct amount of words.
        #[error("Invalid word length in set {0}: {1} != expected {2}")]
        InvalidSetLength(usize, usize, usize),

        /// A mnemonic could not be parsed from the provided mnemonics.
        #[error("Error parsing mnemonic set {0}: {1}")]
        MnemonicFromStrError(usize, MnemonicFromStrError),
    }

    /// Validate a set of mnemonics of a specific word length.
    #[derive(Clone)]
    pub struct MnemonicSetValidator<const N: usize> {
        /// The exact word lengths of all mnemonics. Unlike [`MnemonicValidator`], ranges of words
        /// are not allowed.
        pub word_lengths: [usize; N],
    }

    impl<const N: usize> Validator for MnemonicSetValidator<N> {
        type Output = [Mnemonic; N];
        type Error = MnemonicSetValidationError;

        fn to_fn(&self) -> Box<dyn Fn(String) -> Result<Self::Output, Box<dyn std::error::Error>>> {
            let word_lengths = self.word_lengths;
            Box::new(move |s: String| {
                let mut counter: usize = 0;
                let mut output = Vec::with_capacity(N);
                for (word_set, word_length) in word_lengths.into_iter().enumerate() {
                    let words = s
                        .split_whitespace()
                        .skip(counter)
                        .take(word_length)
                        .collect::<Vec<_>>();
                    if words.len() != word_length {
                        return Err(Box::new(MnemonicSetValidationError::InvalidSetLength(
                            word_set,
                            words.len(),
                            word_length,
                        )));
                    }
                    let mnemonic = match Mnemonic::from_str(&words.join(" ")) {
                        Ok(m) => m,
                        Err(e) => {
                            return Err(Box::new(Self::Error::MnemonicFromStrError(word_set, e)))
                        }
                    };
                    output.push(mnemonic);
                    counter += word_length;
                }

                Ok(output
                    .try_into()
                    .expect(bug!("vec with capacity of const N was not filled")))
            })
        }
    }
}