osrand/

lib.rs

#![warn(clippy::all, clippy::pedantic)]
#![doc = include_str!("../README.md")]
use std::{
    error, fmt,
    fs::File,
    io::{self, BufReader, Read},
    num::TryFromIntError,
};

#[cfg(not(feature = "urandom"))]
static RAND_DEV: &str = "/dev/random";
#[cfg(feature = "urandom")]
static RAND_DEV: &str = "/dev/urandom";

static ALPHA_LOWER: [char; 26] = [
    'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's',
    't', 'u', 'v', 'w', 'x', 'y', 'z',
];

static ALPHA_UPPER: [char; 26] = [
    'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S',
    'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
];

static NUMERIC: [char; 10] = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'];

static SYMBOLS: [char; 20] = [
    '~', '!', '@', '#', '$', '%', '^', '&', '*', '-', '_', '=', '+', ':', ';', '<', '>', ',', '.',
    '?',
];

#[derive(Debug)]
pub enum Error {
    Io(io::Error),
    TryFromInt,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Io(e) => write!(f, "{e}"),
            Self::TryFromInt => write!(f, "TryFromIntError"),
        }
    }
}

impl error::Error for Error {
    fn source(&self) -> Option<&(dyn error::Error + 'static)> {
        match self {
            Self::Io(e) => Some(e),
            Self::TryFromInt => None,
        }
    }
}

impl From<io::Error> for Error {
    fn from(value: io::Error) -> Self {
        Self::Io(value)
    }
}

impl From<TryFromIntError> for Error {
    fn from(_value: TryFromIntError) -> Self {
        Self::TryFromInt
    }
}

/// Pulls random integers from the OS RNG device.
/// This object can be reused to create multiple random
/// numbers and uses an internal `BufReader` around the
/// rng device file in order to keep from doing multiple
/// small reads.
pub struct BufRng {
    reader: BufReader<File>,
}

impl BufRng {
    /// Creates a new instance
    /// # Errors
    /// Returns an io error if there is a problem opening the RNG device
    pub fn new() -> Result<Self, io::Error> {
        let fd = File::open(RAND_DEV)?;
        Ok(Self {
            reader: BufReader::new(fd),
        })
    }

    /// Gets a `u16`
    /// # Errors
    /// Returns an io error if there is a problem reading from the RNG device
    pub fn get_u16(&mut self) -> Result<u16, io::Error> {
        let mut buf = [0; 2];
        self.reader.read_exact(&mut buf)?;
        Ok(u16::from_ne_bytes(buf))
    }

    /// Gets a `u32`
    /// # Errors
    /// Returns an io error if there is a problem reading from the RNG device
    pub fn get_u32(&mut self) -> Result<u32, io::Error> {
        let mut buf = [0; 4];
        self.reader.read_exact(&mut buf)?;
        Ok(u32::from_ne_bytes(buf))
    }

    /// Gets a `u64`
    /// # Errors
    /// Returns an io error if there is a problem reading from the RNG device
    pub fn get_u64(&mut self) -> Result<u64, io::Error> {
        let mut buf = [0; 8];
        self.reader.read_exact(&mut buf)?;
        Ok(u64::from_ne_bytes(buf))
    }
}

#[repr(u8)]
#[derive(Clone, Copy)]
/// These flags specify the contents of the dictionary used to
/// create random strings.
pub enum Flags {
    Lowercase = 0o1,
    Uppercase = 0o2,
    Numeric = 0o4,
    Special = 0o10,
}

impl Flags {
    #[must_use]
    /// Creates a dictionary using all of the available characters
    pub fn all() -> Vec<char> {
        let mut dict = Vec::with_capacity(82);
        dict.extend_from_slice(&ALPHA_LOWER);
        dict.extend_from_slice(&ALPHA_UPPER);
        dict.extend_from_slice(&NUMERIC);
        dict.extend_from_slice(&SYMBOLS);
        dict
    }

    #[must_use]
    /// Creates a dictionary using only alphanumeric characters
    pub fn alphanumeric() -> Vec<char> {
        let mut dict = Vec::with_capacity(62);
        dict.extend_from_slice(&ALPHA_LOWER);
        dict.extend_from_slice(&ALPHA_UPPER);
        dict.extend_from_slice(&NUMERIC);
        dict
    }

    #[must_use]
    /// Creates a dictionary using only alphabet characters
    pub fn alphabetical() -> Vec<char> {
        let mut dict = Vec::with_capacity(52);
        dict.extend_from_slice(&ALPHA_LOWER);
        dict.extend_from_slice(&ALPHA_UPPER);
        dict
    }
}

/// Gets a `u16` from the RNG device. Do not use this function if
/// you require multiple random numbers, as each use will be a single
/// read.
/// # Errors
/// Returns an io error if there is a problem reading from the RNG device
pub fn random_u16() -> Result<u16, io::Error> {
    let mut buf = [0; 2];
    let mut fd = File::open(RAND_DEV)?;
    fd.read_exact(&mut buf)?;
    Ok(u16::from_ne_bytes(buf))
}

/// Gets a `u32` from the RNG device. Do not use this function if
/// you require multiple random numbers, as each use will be a single
/// read.
/// # Errors
/// Returns an io error if there is a problem reading from the RNG device
pub fn random_u32() -> Result<u32, io::Error> {
    let mut buf = [0; 4];
    let mut fd = File::open(RAND_DEV)?;
    fd.read_exact(&mut buf)?;
    Ok(u32::from_ne_bytes(buf))
}

/// Gets a `u64` from the RNG device. Do not use this function if
/// you require multiple random numbers, as each use will be a single
/// read.
/// # Errors
/// Returns an io error if there is a problem reading from the RNG device
pub fn random_u64() -> Result<u64, io::Error> {
    let mut buf = [0; 8];
    let mut fd = File::open(RAND_DEV)?;
    fd.read_exact(&mut buf)?;
    Ok(u64::from_ne_bytes(buf))
}

/// A random string generator which gets it's entropy from an internal
/// `BufReader` wrapping the OS RNG device. This generator may be re-used
/// as many times as required.
pub struct RandomString {
    dictionary: Vec<char>,
    rng: BufRng,
}

impl From<RandomString> for BufRng {
    fn from(value: RandomString) -> Self {
        value.rng
    }
}

impl From<BufRng> for RandomString {
    fn from(value: BufRng) -> Self {
        Self { dictionary: Flags::all(), rng: value }
    }
}

impl RandomString {
    /// Creates a new random string generator, which gets it's randomness
    /// from an internal `BufReader` around the OS RNG device. If `flags`
    /// is empty, the full dictionary will be used.
    /// # Errors
    /// Returns an io error if there is a problem reading from the RNG device
    pub fn new(flags: &[Flags]) -> Result<Self, io::Error> {
        let dictionary = if flags.is_empty() {
            Flags::all()
        } else {
            let mut dict = vec![];
            flags.iter().for_each(|f| match f {
                Flags::Lowercase => dict.extend_from_slice(&ALPHA_LOWER),
                Flags::Uppercase => dict.extend_from_slice(&ALPHA_UPPER),
                Flags::Numeric => dict.extend_from_slice(&NUMERIC),
                Flags::Special => dict.extend_from_slice(&SYMBOLS),
            });
            dict
        };
        Ok(Self {
            dictionary,
            rng: BufRng::new()?,
        })
    }

    /// Creates a new random string generator with the given dictionary, which
    /// gets it's randomness from an internal `BufReader` around the OS RNG device.
    /// If `dict` is empty, the full dictionary will be used.
    /// # Errors
    /// Returns an io error if there is a problem reading from the RNG device
    pub fn with_dict(dict: Vec<char>) -> Result<Self, io::Error> {
        let dictionary = if dict.is_empty() {
            Flags::all()
        } else {
            dict
        };
        Ok(Self {
            dictionary,
            rng: BufRng::new()?,
        })
    }

    /// Creates a new random string generator from the provided `BufRng` rng`
    /// and the `Dictionary` dict.
    /// # Errors
    /// Returns an io error if there is a problem reading from the RNG device
    pub fn from_parts(rng: BufRng, dict: Vec<char>) -> Self {
        let dictionary = if dict.is_empty() {
            Flags::all()
        } else {
            dict
        };
        Self {
            dictionary,
            rng,
        }
    }

    #[must_use]
    /// Gets the dictionary being used by the generator
    pub fn get_dictionary(&self) -> &[char] {
        &self.dictionary
    }

    /// Sets the dictionary to be used for new random strings
    pub fn set_dictionary(&mut self, dict: Vec<char>) {
        self.dictionary = dict;
    }

    /// Generates a random string of the given size
    /// # Errors
    /// Returns an io error if there is a problem reading from the RNG device
    pub fn gen(&mut self, len: usize) -> Result<String, Error> {
        let mut s = String::with_capacity(len);
        for _i in 0..len {
            let n = self.rng.get_u32()?;
            let idx = usize::try_from(n)? % self.dictionary.len();
            if let Some(c) = self.dictionary.get(idx) {
                s.push(*c);
            }
        }
        Ok(s)
    }

    /// Appends `len` random characters  to string `s` and returns the result
    /// # Errors
    /// Returns an io error if there is a problem reading from the RNG device
    pub fn append(&mut self, mut s: String, len: usize) -> Result<String, Error> {
        for _i in 0..len {
            let n = self.rng.get_u32()?;
            let idx = usize::try_from(n)? % self.dictionary.len();
            if let Some(c) = self.dictionary.get(idx) {
                s.push(*c);
            }
        }
        Ok(s)
    }
}

#[test]
fn random_string() {
    let mut rs = RandomString::new(&[
        Flags::Lowercase,
        Flags::Numeric,
        Flags::Uppercase,
        Flags::Special,
    ])
    .unwrap();
    let out = rs.gen(8).unwrap();
    assert_eq!(out.len(), 8);
}