aboutsummaryrefslogtreecommitdiff
path: root/src/rand.rs
blob: 097341f796788b3175b77c4a7ab6e699979fef7e (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: BSD-3-Clause

//! Miscellaneous functions related to getting (pseudo) random numbers and
//! strings.
//!
//! NOTE! This should not be used when you do need __real__ random numbers such
//! as for encryption but will probably be suitable when you want locally
//! unique ID's that will not be shared over the network.

use std::ffi::OsString;
use std::str;

/// Gets an ever increasing u64 (at least for this process).
///
/// The number retrieved will be based upon the time of the last reboot (x86_64)
/// and something undefined for other architectures.
pub fn timestamp_cycles() -> u64 {
    #[cfg(target_arch = "x86_64")]
    // SAFETY: Safe because there's nothing that can go wrong with this call.
    unsafe {
        std::arch::x86_64::_rdtsc() as u64
    }

    #[cfg(not(target_arch = "x86_64"))]
    {
        const MONOTONIC_CLOCK_MULTPIPLIER: u64 = 1_000_000_000;

        let mut ts = libc::timespec {
            tv_sec: 0,
            tv_nsec: 0,
        };
        // SAFETY: We initialized the parameters correctly and we trust the function.
        unsafe {
            libc::clock_gettime(libc::CLOCK_MONOTONIC, &mut ts);
        }
        (ts.tv_sec as u64) * MONOTONIC_CLOCK_MULTPIPLIER + (ts.tv_nsec as u64)
    }
}

/// Generate pseudo random u32 numbers based on the current timestamp.
pub fn xor_pseudo_rng_u32() -> u32 {
    let mut t: u32 = timestamp_cycles() as u32;
    // Taken from https://en.wikipedia.org/wiki/Xorshift
    t ^= t << 13;
    t ^= t >> 17;
    t ^ (t << 5)
}

// This will get an array of numbers that can safely be converted to strings
// because they will be in the range [a-zA-Z0-9]. The return vector could be any
// size between 0 and 4.
fn xor_pseudo_rng_u8_alphanumerics(rand_fn: &dyn Fn() -> u32) -> Vec<u8> {
    rand_fn()
        .to_ne_bytes()
        .to_vec()
        .drain(..)
        .filter(|val| {
            (48..=57).contains(val) || (65..=90).contains(val) || (97..=122).contains(val)
        })
        .collect()
}

fn xor_pseudo_rng_u8_bytes(rand_fn: &dyn Fn() -> u32) -> Vec<u8> {
    rand_fn().to_ne_bytes().to_vec()
}

fn rand_alphanumerics_impl(rand_fn: &dyn Fn() -> u32, len: usize) -> OsString {
    let mut buf = OsString::new();
    let mut done = 0;
    loop {
        for n in xor_pseudo_rng_u8_alphanumerics(rand_fn) {
            done += 1;
            buf.push(str::from_utf8(&[n]).unwrap_or("_"));
            if done >= len {
                return buf;
            }
        }
    }
}

fn rand_bytes_impl(rand_fn: &dyn Fn() -> u32, len: usize) -> Vec<u8> {
    let mut buf: Vec<Vec<u8>> = Vec::new();
    let mut num = if len % 4 == 0 { len / 4 } else { len / 4 + 1 };
    while num > 0 {
        buf.push(xor_pseudo_rng_u8_bytes(rand_fn));
        num -= 1;
    }
    buf.into_iter().flatten().take(len).collect()
}

/// Gets a pseudo random OsString of length `len` with characters in the
/// range [a-zA-Z0-9].
pub fn rand_alphanumerics(len: usize) -> OsString {
    rand_alphanumerics_impl(&xor_pseudo_rng_u32, len)
}

/// Get a pseudo random vector of `len` bytes.
pub fn rand_bytes(len: usize) -> Vec<u8> {
    rand_bytes_impl(&xor_pseudo_rng_u32, len)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_timestamp_cycles() {
        for _ in 0..1000 {
            assert!(timestamp_cycles() < timestamp_cycles());
        }
    }

    #[test]
    fn test_xor_pseudo_rng_u32() {
        for _ in 0..1000 {
            assert_ne!(xor_pseudo_rng_u32(), xor_pseudo_rng_u32());
        }
    }

    #[test]
    fn test_xor_pseudo_rng_u8_alphas() {
        let i = 3612982; // 55 (shifted 16 places), 33 (shifted 8 places), 54...
                         // The 33 will be discarded as it is not a valid letter
                         // (upper or lower) or number.
        let s = xor_pseudo_rng_u8_alphanumerics(&|| i);
        assert_eq!(vec![54, 55], s);
    }

    #[test]
    fn test_rand_alphanumerics_impl() {
        let s = rand_alphanumerics_impl(&|| 14134, 5);
        assert_eq!("67676", s);
    }

    #[test]
    fn test_rand_alphanumerics() {
        let s = rand_alphanumerics(5);
        assert_eq!(5, s.len());
    }

    #[test]
    fn test_xor_pseudo_rng_u8_bytes() {
        let i = 3612982; // 55 (shifted 16 places), 33 (shifted 8 places), 54...
                         // The 33 will be discarded as it is not a valid letter
                         // (upper or lower) or number.
        let s = xor_pseudo_rng_u8_bytes(&|| i);
        assert_eq!(vec![54, 33, 55, 0], s);
    }

    #[test]
    fn test_rand_bytes_impl() {
        let s = rand_bytes_impl(&|| 1234567, 4);
        assert_eq!(vec![135, 214, 18, 0], s);
    }

    #[test]
    fn test_rand_bytes() {
        for i in 0..8 {
            assert_eq!(i, rand_bytes(i).len());
        }
    }
}