You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

97 lines
2.4 KiB

use super::*;
use std::cmp;
use std::ptr;
use std::ffi::c_void;
/// Explicitly zero out a byte buffer.
///
/// Essentially `explicit_bzero()`.
#[inline(never)]
pub fn explicit_clear(buffer : &mut[u8]) {
unsafe {
std::ptr::write_bytes(buffer.as_mut_ptr() as * mut c_void, 0, buffer.len());
if cfg!(target_arch = "x86_64") || cfg!(target_arch = "x86") {
asm!("clflush [{}]", in(reg)buffer.as_mut_ptr());
} else {
asm!("")
}
}
}
/// Set all bytes of this buffer to a single value.
///
/// Essentially `memset()`.
#[inline] pub fn set(buffer: &mut [u8], to: u8)
{
unsafe {
std::ptr::write_bytes(buffer.as_mut_ptr() as *mut c_void, to, buffer.len());
}
}
/// Zero out this buffer
///
/// Essentially `bzero()`.
#[inline(always)] pub fn clear(buffer: &mut [u8])
{
set(buffer, 0);
}
/// Copy bytes from one slice to another.
///
/// # Notes
/// The slices can overlap.
#[inline] pub fn memmove(from: &[u8], to: &mut [u8]) -> usize
{
let len = cmp::min(from.len(), to.len());
unsafe {
ptr::copy(from.as_ptr(), to.as_mut_ptr(), len);
}
len
}
/// Copy bytes from one slice to another.
///
/// # Notes
/// The slices must *not* overlap.
#[inline] pub fn memcpy(from: &[u8], to: &mut [u8]) -> usize
{
let len = cmp::min(from.len(), to.len());
unsafe {
ptr::copy_nonoverlapping(from.as_ptr(), to.as_mut_ptr(), len);
}
len
}
/// Max size of bytes we'll allocate to the stack at runtime before using a heap allocated buffer.
pub const STACK_SIZE_LIMIT: usize = 4096;
/// Allocate `size` bytes. Allocates on the stack if size is lower than `STACK_SIZE_LIMIT`, otherwise allocates on the heap.
pub fn alloca_limit<F, T>(size: usize, f: F) -> T
where F: FnOnce(&mut [u8]) -> T
{
if size > STACK_SIZE_LIMIT {
thread_local! {
static BUFFER: RefCell<Vec<u8>> = RefCell::new(vec![0u8; STACK_SIZE_LIMIT*2]);
}
BUFFER.with(move |buf| {
// If the borrow fails then `f` has recursively called into this function, so for that we allocate a new buffer instead of reusing this static one.
if let Ok(mut buf) = buf.try_borrow_mut() {
if buf.len() < size {
buf.resize(size, 0);
}
let res = f(&mut buf[..size]);
bytes::clear(&mut buf[..size]);
res
} else {
f(&mut vec![0u8; size])
}
})
} else {
stackalloc::alloca_zeroed(size, f)
// I don't think this is okay to do.
//stackalloc::alloca(size, move |buf| f(unsafe { stackalloc::helpers::slice_assume_init_mut(buf) }))
}
}