vmm_sys_util/
rand.rs

1// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: BSD-3-Clause
3
4//! Miscellaneous functions related to getting (pseudo) random numbers and
5//! strings.
6//!
7//! NOTE! This should not be used when you do need __real__ random numbers such
8//! as for encryption but will probably be suitable when you want locally
9//! unique ID's that will not be shared over the network.
10
11use std::ffi::OsString;
12use std::str;
13
14/// Gets an ever increasing u64 (at least for this process).
15///
16/// The number retrieved will be based upon the time of the last reboot (x86_64)
17/// and something undefined for other architectures.
18pub fn timestamp_cycles() -> u64 {
19    #[cfg(target_arch = "x86_64")]
20    // SAFETY: Safe because there's nothing that can go wrong with this call.
21    unsafe {
22        std::arch::x86_64::_rdtsc()
23    }
24
25    #[cfg(not(target_arch = "x86_64"))]
26    {
27        const MONOTONIC_CLOCK_MULTPIPLIER: u64 = 1_000_000_000;
28
29        let mut ts = libc::timespec {
30            tv_sec: 0,
31            tv_nsec: 0,
32        };
33        // SAFETY: We initialized the parameters correctly and we trust the function.
34        unsafe {
35            libc::clock_gettime(libc::CLOCK_MONOTONIC, &mut ts);
36        }
37        (ts.tv_sec as u64) * MONOTONIC_CLOCK_MULTPIPLIER + (ts.tv_nsec as u64)
38    }
39}
40
41/// Generate pseudo random u32 numbers based on the current timestamp.
42pub fn xor_pseudo_rng_u32() -> u32 {
43    let mut t: u32 = timestamp_cycles() as u32;
44    // Taken from https://en.wikipedia.org/wiki/Xorshift
45    t ^= t << 13;
46    t ^= t >> 17;
47    t ^ (t << 5)
48}
49
50// This will get an array of numbers that can safely be converted to strings
51// because they will be in the range [a-zA-Z0-9]. The return vector could be any
52// size between 0 and 4.
53fn xor_pseudo_rng_u8_alphanumerics(rand_fn: &dyn Fn() -> u32) -> Vec<u8> {
54    rand_fn()
55        .to_ne_bytes()
56        .to_vec()
57        .drain(..)
58        .filter(|val| {
59            (48..=57).contains(val) || (65..=90).contains(val) || (97..=122).contains(val)
60        })
61        .collect()
62}
63
64fn xor_pseudo_rng_u8_bytes(rand_fn: &dyn Fn() -> u32) -> Vec<u8> {
65    rand_fn().to_ne_bytes().to_vec()
66}
67
68fn rand_alphanumerics_impl(rand_fn: &dyn Fn() -> u32, len: usize) -> OsString {
69    let mut buf = OsString::new();
70    let mut done = 0;
71    loop {
72        for n in xor_pseudo_rng_u8_alphanumerics(rand_fn) {
73            done += 1;
74            buf.push(str::from_utf8(&[n]).unwrap_or("_"));
75            if done >= len {
76                return buf;
77            }
78        }
79    }
80}
81
82fn rand_bytes_impl(rand_fn: &dyn Fn() -> u32, len: usize) -> Vec<u8> {
83    let mut buf: Vec<Vec<u8>> = Vec::new();
84    let mut num = if len % 4 == 0 { len / 4 } else { len / 4 + 1 };
85    while num > 0 {
86        buf.push(xor_pseudo_rng_u8_bytes(rand_fn));
87        num -= 1;
88    }
89    buf.into_iter().flatten().take(len).collect()
90}
91
92/// Gets a pseudo random OsString of length `len` with characters in the
93/// range [a-zA-Z0-9].
94pub fn rand_alphanumerics(len: usize) -> OsString {
95    rand_alphanumerics_impl(&xor_pseudo_rng_u32, len)
96}
97
98/// Get a pseudo random vector of `len` bytes.
99pub fn rand_bytes(len: usize) -> Vec<u8> {
100    rand_bytes_impl(&xor_pseudo_rng_u32, len)
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    #[test]
108    fn test_timestamp_cycles() {
109        for _ in 0..1000 {
110            assert!(timestamp_cycles() < timestamp_cycles());
111        }
112    }
113
114    #[test]
115    fn test_xor_pseudo_rng_u32() {
116        for _ in 0..1000 {
117            assert_ne!(xor_pseudo_rng_u32(), xor_pseudo_rng_u32());
118        }
119    }
120
121    #[test]
122    fn test_xor_pseudo_rng_u8_alphas() {
123        let i = 3612982; // 55 (shifted 16 places), 33 (shifted 8 places), 54...
124                         // The 33 will be discarded as it is not a valid letter
125                         // (upper or lower) or number.
126        let s = xor_pseudo_rng_u8_alphanumerics(&|| i);
127        if cfg!(target_endian = "big") {
128            assert_eq!(vec![55, 54], s);
129        } else {
130            assert_eq!(vec![54, 55], s);
131        }
132    }
133
134    #[test]
135    fn test_rand_alphanumerics_impl() {
136        let s = rand_alphanumerics_impl(&|| 14134, 5);
137        if cfg!(target_endian = "big") {
138            assert_eq!("76767", s);
139        } else {
140            assert_eq!("67676", s);
141        }
142    }
143
144    #[test]
145    fn test_rand_alphanumerics() {
146        let s = rand_alphanumerics(5);
147        assert_eq!(5, s.len());
148    }
149
150    #[test]
151    fn test_xor_pseudo_rng_u8_bytes() {
152        let i = 3612982; // 55 (shifted 16 places), 33 (shifted 8 places), 54...
153                         // The 33 will be discarded as it is not a valid letter
154                         // (upper or lower) or number.
155        let s = xor_pseudo_rng_u8_bytes(&|| i);
156        if cfg!(target_endian = "big") {
157            assert_eq!(vec![0, 55, 33, 54], s);
158        } else {
159            assert_eq!(vec![54, 33, 55, 0], s);
160        }
161    }
162
163    #[test]
164    fn test_rand_bytes_impl() {
165        let s = rand_bytes_impl(&|| 1234567, 4);
166        if cfg!(target_endian = "big") {
167            assert_eq!(vec![0, 18, 214, 135], s);
168        } else {
169            assert_eq!(vec![135, 214, 18, 0], s);
170        }
171    }
172
173    #[test]
174    fn test_rand_bytes() {
175        for i in 0..8 {
176            assert_eq!(i, rand_bytes(i).len());
177        }
178    }
179}