zeros 17.3.2

/*
==--==--==--==--==--==--==--==--==--==--==--==--==--==--==--==--

Zeros

Copyright (C) 2019-2025  Anonymous

There are several releases over multiple years,
they are listed as ranges, such as: "2019-2025".

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.

::--::--::--::--::--::--::--::--::--::--::--::--::--::--::--::--
*/

//! # I/O
//!
//! ## Reading streams
//!
//! - A stream is [`Read`][trait:std/io/Read].
//! - Currently, read buffer is 16 KiB.
//! - Reading stops until the stream returns 0 bytes.
//! - [`ErrorKind::WouldBlock`][enum-variant:std/io/ErrorKind#WouldBlock] will be ignored. It then waits for 10 milliseconds, and retries.
//!
//! ## Reading `/dev/urandom`
//!
//! - `limit` must be larger than zero. Or an error is returned.
//! - If data read from that file is not equal to `limit`, an error is returned.
//!
//! [enum-variant:std/io/ErrorKind#WouldBlock]: https://doc.rust-lang.org/std/io/enum.ErrorKind.html#variant.WouldBlock
//! [trait:std/io/Read]: https://doc.rust-lang.org/std/io/trait.Read.html

use {
    core::mem,
    alloc::vec::Vec,
};

#[cfg(feature="std")]
use {
    core::time::Duration,
    std::{
        io::{ErrorKind, Read},
        thread,
    },
    crate::IoResult,
};

#[cfg(all(feature="std", unix))]
use std::{
    fs::File,
    io::{self, BufRead, BufReader},
};

#[cfg(not(feature="std"))]
#[doc(cfg(not(feature="std")))]
mod write;

#[cfg(not(feature="std"))]
#[doc(cfg(not(feature="std")))]
pub (crate) use self::write::*;

#[cfg(test)]
mod tests;

/// # Buffer size
#[cfg(feature="std")]
#[doc(cfg(feature="std"))]
pub (crate) const BUFFER_SIZE: usize = 1024 * 16;

/// # Reads a stream
#[cfg(feature="std")]
#[doc(cfg(feature="std"))]
pub (crate) fn read_stream<R, F>(stream: &mut R, mut processor: F) -> IoResult<()> where R: Read, F: FnMut(&[u8]) -> IoResult<()> {
    let mut buffer = [0; BUFFER_SIZE];
    loop {
        match stream.read(&mut buffer) {
            Ok(read) => match read {
                0 => return Ok(()),
                _ => if read <= buffer.len() {
                    processor(&buffer[..read])?;
                } else {
                    return Err(err!(
                        "read() returned invalid value: {read} (buffer size: {buffer_size})",
                        read=read, buffer_size=buffer.len(),
                    ).into());
                },
            },
            Err(err) => match err.kind() {
                ErrorKind::WouldBlock => {
                    thread::sleep(Duration::from_millis(10));
                    continue;
                },
                _ => return Err(err),
            },
        };
    }
}

/// # Fills buffer
///
/// - Required: `max - buffer.len()`.
/// - Extends buffer from the slice and reduces its size.
/// - If the slice has enough required data, returns `true`.
#[inline(always)]
pub (crate) fn fill_buffer(buffer: &mut Vec<u8>, max: usize, bytes: &mut &[u8]) -> bool {
    match max.checked_sub(buffer.len()) {
        Some(required) if required > 0 => {
            let available = bytes.len().min(required);
            buffer.extend(&bytes[..available]);
            *bytes = &bytes[available..];
            available == required
        },
        _ => true,
    }
}

/// # Copies data into a buffer
///
/// - The buffer index will be updated.
/// - Returns how many bytes were copied.
#[inline(always)]
pub (crate) fn copy_into_buffer<B>(bytes: B, buffer: &mut &mut [u8]) -> usize where B: AsRef<[u8]> {
    let bytes = bytes.as_ref();
    let count = bytes.len().min(buffer.len());
    if count > 0 {
        let (first, second) = mem::replace(buffer, &mut []).split_at_mut(count);
        first.copy_from_slice(&bytes[..count]);
        *buffer = second;
    }
    count
}

/// # Reads `/dev/urandom`
#[cfg(all(feature="std", unix))]
#[doc(cfg(all(feature="std", unix)))]
pub (crate) fn read_urandom<F>(limit: usize, mut proc: F) -> IoResult<()> where F: FnMut(&[u8]) {
    const PATH_DEV_URANDOM: &str = "/dev/urandom";

    if limit == 0 {
        return Err(err!("Limit must be larger than zero").into());
    }

    let mut reader = BufReader::new(
        File::open(PATH_DEV_URANDOM)?.take(limit.try_into().map_err(|_| io::Error::from(err!("Failed to convert {} into `u64`", limit)))?)
    );
    let mut total_read = 0;
    loop {
        let buffer = reader.fill_buf()?;
        match buffer.len() {
            0 => return if total_read == limit {
                Ok(())
            } else {
                Err(err!("Expected to read {limit} bytes, got {total_read}", limit=limit, total_read=total_read).into())
            },
            read => match total_read.checked_add(read) {
                Some(new) if new <= limit => {
                    proc(buffer);
                    reader.consume(read);
                    total_read = new;
                },
                // This is unlikely to happen, because we used File::open().take() above
                _ => return Err(e!().into()),
            },
        };
    }
}

/// # Uses [`CryptGenRandom()`](https://learn.microsoft.com/en-us/windows/win32/api/wincrypt/nf-wincrypt-cryptgenrandom)
#[cfg(all(feature="std", windows))]
#[doc(cfg(all(feature="std", windows)))]
pub (crate) fn read_urandom<F>(limit: usize, mut proc: F) -> IoResult<()> where F: FnMut(&[u8]) {
    use core::{
        ffi::c_char,
        ptr,
    };

    #[cfg(target_pointer_width="64")]
    type ULONG = u64;
    #[cfg(target_pointer_width="32")]
    type ULONG = u32;

    #[link(name="Advapi32")]
    unsafe extern "system" {
        unsafe fn CryptAcquireContextA(dev: *mut ULONG, container: *const c_char, provider: *const c_char, provider_type: u32, flags: u32)
        -> bool;
        unsafe fn CryptGenRandom(dev: ULONG, len: u32, buffer: *mut u8) -> bool;
        unsafe fn CryptReleaseContext(dev: ULONG, flags: u32) -> bool;
    }
    #[link(name="Kernel32")]
    unsafe extern "system" {
        safe fn GetLastError() -> u32;
    }

    macro_rules! buf_len { () => { 64 }}
    const PROV_RSA_FULL: u32 = 1;
    const CRYPT_VERIFYCONTEXT: u32 = 0xf0000000;

    if limit == 0 {
        return Err(err!("Limit must be larger than zero").into());
    }

    let mut dev = 0;
    if false == unsafe { CryptAcquireContextA(ptr::from_mut(&mut dev), ptr::null(), ptr::null(), PROV_RSA_FULL, CRYPT_VERIFYCONTEXT) } {
        return Err(e!("CryptAcquireContextA() failed").into());
    }

    let mut buffer = [u8::MIN; buf_len!()];
    let mut read = 0;
    while read < limit {
        if false == unsafe { CryptGenRandom(dev, buf_len!(), buffer.as_mut_ptr()) } {
            return Err(err!("CryptGenRandom() failed: {}", GetLastError()).into());
        }
        let count = buf_len!().min(limit - read);
        proc(&buffer[..count]);
        read += count;
    }
    if false == unsafe { CryptReleaseContext(dev, 0) } {
        return Err(e!("CryptReleaseContext() failed").into());
    }

    Ok(())
}