rok-core 0.6.1

Core primitives for the rok ecosystem — errors, crypto, i18n, config, DI, and more
Documentation
use std::fmt;
use std::str::FromStr;

use rand::RngCore;
use serde::{Deserialize, Serialize};

use crate::crypto::ids::IdError;

pub const DEFAULT_ALPHABET: &[u8] =
    b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789_-";
pub const DEFAULT_SIZE: usize = 21;

pub fn generate_nanoid_custom(alphabet: &[u8], size: usize) -> String {
    assert!(!alphabet.is_empty(), "nanoid: alphabet must not be empty");
    assert!(
        alphabet.len() <= 256,
        "nanoid: alphabet too large (max 256)"
    );

    let mut rng = rand::thread_rng();
    let alen = alphabet.len();

    let mask = {
        let mut m = 1usize;
        while m < alen {
            m = (m << 1) | 1;
        }
        m
    };

    let mut result = Vec::with_capacity(size);
    let mut buf = [0u8; 256];

    while result.len() < size {
        rng.fill_bytes(&mut buf);
        for &b in &buf {
            let idx = (b as usize) & mask;
            if idx < alen {
                result.push(alphabet[idx]);
                if result.len() == size {
                    break;
                }
            }
        }
    }

    String::from_utf8(result).expect("nanoid alphabet must be ASCII")
}

#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct NanoId(String);

impl NanoId {
    pub fn generate() -> Self {
        Self(generate_nanoid_custom(DEFAULT_ALPHABET, DEFAULT_SIZE))
    }

    pub fn with_size(size: usize) -> Self {
        Self(generate_nanoid_custom(DEFAULT_ALPHABET, size))
    }

    pub fn custom(alphabet: &[u8], size: usize) -> Self {
        Self(generate_nanoid_custom(alphabet, size))
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl fmt::Display for NanoId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(&self.0)
    }
}

impl FromStr for NanoId {
    type Err = IdError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.is_empty() {
            return Err(IdError::InvalidFormat("nanoid", "must not be empty"));
        }
        Ok(Self(s.to_owned()))
    }
}

impl AsRef<str> for NanoId {
    fn as_ref(&self) -> &str {
        &self.0
    }
}

#[cfg(feature = "crypto-sqlx")]
mod sqlx_impl {
    use super::NanoId;
    use sqlx::{
        encode::IsNull,
        error::BoxDynError,
        postgres::{PgArgumentBuffer, PgTypeInfo, PgValueRef},
    };

    impl sqlx::Type<sqlx::Postgres> for NanoId {
        fn type_info() -> PgTypeInfo {
            <String as sqlx::Type<sqlx::Postgres>>::type_info()
        }
    }

    impl<'q> sqlx::Encode<'q, sqlx::Postgres> for NanoId {
        fn encode_by_ref(&self, buf: &mut PgArgumentBuffer) -> Result<IsNull, BoxDynError> {
            <String as sqlx::Encode<'q, sqlx::Postgres>>::encode_by_ref(&self.0, buf)
        }
    }

    impl<'r> sqlx::Decode<'r, sqlx::Postgres> for NanoId {
        fn decode(value: PgValueRef<'r>) -> Result<Self, BoxDynError> {
            let s = <String as sqlx::Decode<sqlx::Postgres>>::decode(value)?;
            Ok(Self(s))
        }
    }
}