randgen 0.1.2

Generate pseudo-random bytes
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

// SPDX-FileCopyrightText: 2025 Shun Sakai
//
// SPDX-License-Identifier: Apache-2.0 OR MIT

use std::io::{self, BufWriter, Write};

use anyhow::Context;
use clap::Parser;

use crate::{
    cli::{Format, Opt},
    rng::Rng,
};

// 8 KiB.
const CHUNK_SIZE: usize = 1 << 13;

// 1 MiB.
const BUF_SIZE: usize = 1 << 20;

/// Runs the program and returns the result.
pub fn run() -> anyhow::Result<()> {
    let opt = Opt::parse();

    if let Some(shell) = opt.generate_completion {
        Opt::print_completion(shell);
        return Ok(());
    }

    let rng = opt.random_number_generator;
    let mut rng = if let Some(seed) = opt.seed {
        Rng::seed_from_u64(&rng, seed)
    } else {
        Rng::try_from_os_rng(&rng).context("could not create a new instance of the RNG")?
    };

    let stdout = io::stdout();
    let mut writer = BufWriter::with_capacity(BUF_SIZE, stdout.lock());

    let mut remaining = opt
        .length
        .expect("the number of bytes to generate should be provided")
        .try_into()?;

    let mut buf = [u8::default(); CHUNK_SIZE];

    match opt.format {
        #[cfg(feature = "base64")]
        format @ (Format::Base64 | Format::Base64Url) => {
            let engine = match format {
                Format::Base64 => base64::engine::general_purpose::STANDARD,
                Format::Base64Url => base64::engine::general_purpose::URL_SAFE,
                _ => unreachable!(),
            };
            let mut writer = base64::write::EncoderWriter::new(writer, &engine);
            while remaining > 0 {
                let chunk_size = CHUNK_SIZE.min(remaining);
                rng.fill_bytes(&mut buf[..chunk_size]);
                writer
                    .write_all(&buf[..chunk_size])
                    .context("could not write base64 encoded random bytes to standard output")?;
                remaining -= chunk_size;
            }
            let mut writer = writer.finish().context(
                "could not write remaining base64 encoded random bytes to standard output",
            )?;
            writer
                .flush()
                .context("could not flush base64 encoded random bytes to standard output")?;
        }
        format => {
            while remaining > 0 {
                let chunk_size = CHUNK_SIZE.min(remaining);
                rng.fill_bytes(&mut buf[..chunk_size]);
                match format {
                    Format::Raw => writer
                        .write_all(&buf[..chunk_size])
                        .context("could not write raw random bytes to standard output")?,
                    #[cfg(feature = "hex")]
                    Format::Hex => {
                        let s = faster_hex::hex_string(&buf[..chunk_size]);
                        write!(writer, "{s}").context(
                            "could not write hex encoded random bytes to standard output",
                        )?;
                    }
                    #[cfg(feature = "base64")]
                    _ => unreachable!(),
                }
                remaining -= chunk_size;
            }
            writer
                .flush()
                .context("could not flush random bytes to standard output")?;
        }
    }
    Ok(())
}

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

    #[test]
    fn chunk_size() {
        assert_eq!(CHUNK_SIZE, 8192);
    }

    #[test]
    fn buf_size() {
        assert_eq!(BUF_SIZE, 1_048_576);
    }
}