Merge pull request #10 from viz-rs/fix-security

bump: v0.4.0

Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-version = "0.3.3"
+version = "0.4.0"
 name = "nano-id"
 edition = "2021"
 authors = ["Fangdun Tsai <cfddream@gmail.com>"]

src/lib.rs

@@ -5,23 +5,30 @@ macro_rules! gen {
     ($mod:tt, $len:tt, $alphabet:tt) => {
         #[doc = concat!(" Nanoid with alphabet table `", stringify!($alphabet), "`")]
         mod $mod {
-            pub const MASK: usize = $len - 1;
+            pub const MASK: usize = ($len as usize).next_power_of_two() - 1;
             pub const ALPHABET: &'static [u8; $len] = $alphabet;
         }
 
         #[doc = concat!(" Nanoid with ", stringify!($mod))]
         #[must_use]
         pub fn $mod<const N: usize>() -> String {
-            let mut bytes = [0u8; N];
+            let mut bytes = vec![0u8; 8 * N / 5];
+            let mut id = String::with_capacity(N);
 
-            ::getrandom::getrandom(&mut bytes)
-                .unwrap_or_else(|err| panic!("could not retreive random bytes: {err}"));
+            loop {
+                ::getrandom::getrandom(&mut bytes)
+                    .unwrap_or_else(|err| panic!("could not retreive random bytes: {err}"));
 
-            bytes
-                .iter_mut()
-                .for_each(|b| *b = $mod::ALPHABET[*b as usize & $mod::MASK]);
-
-            String::from_utf8_lossy(&bytes).to_string()
+                for byte in &bytes {
+                    let idx = *byte as usize & $mod::MASK;
+                    if idx < $len {
+                        id.push($mod::ALPHABET[idx] as char)
+                    }
+                    if id.len() == N {
+                        return id;
+                    }
+                }
+            }
         }
     };
 }
@@ -87,4 +94,26 @@ mod tests {
         println!("{}", &id);
         assert_eq!(id.len(), 21);
     }
+
+    #[test]
+    #[cfg(feature = "base62")]
+    fn symbols() {
+        use std::collections::BTreeMap;
+
+        let mut counts = BTreeMap::new();
+
+        for _ in 0..1_000_000 {
+            let id = base62::<10>();
+            for c in id.chars() {
+                *counts.entry(c).or_insert(0) += 1;
+            }
+        }
+
+        println!("{} symbols generated", counts.len());
+        for (c, count) in &counts {
+            println!("{}: {}", c, count);
+        }
+
+        assert_eq!(counts.len(), 62);
+    }
 }