use crate::{StrErr, Utf8CStr, ffi}; use argh::EarlyExit; use libc::c_char; use std::fmt::Arguments; use std::io::Write; use std::mem::ManuallyDrop; use std::process::exit; use std::sync::Arc; use std::sync::atomic::{AtomicPtr, Ordering}; use std::{fmt, slice, str}; pub fn errno() -> &'static mut i32 { unsafe { &mut *libc::__errno() } } // When len is 0, don't care whether buf is null or not #[inline] pub unsafe fn slice_from_ptr<'a, T>(buf: *const T, len: usize) -> &'a [T] { unsafe { if len == 0 { &[] } else { slice::from_raw_parts(buf, len) } } } // When len is 0, don't care whether buf is null or not #[inline] pub unsafe fn slice_from_ptr_mut<'a, T>(buf: *mut T, len: usize) -> &'a mut [T] { unsafe { if len == 0 { &mut [] } else { slice::from_raw_parts_mut(buf, len) } } } pub trait BytesExt { fn find(&self, needle: &[u8]) -> Option; fn contains(&self, needle: &[u8]) -> bool { self.find(needle).is_some() } } impl + ?Sized> BytesExt for T { fn find(&self, needle: &[u8]) -> Option { fn inner(haystack: &[u8], needle: &[u8]) -> Option { unsafe { let ptr: *const u8 = libc::memmem( haystack.as_ptr().cast(), haystack.len(), needle.as_ptr().cast(), needle.len(), ) .cast(); if ptr.is_null() { None } else { Some(ptr.offset_from(haystack.as_ptr()) as usize) } } } inner(self.as_ref(), needle) } } pub trait MutBytesExt { fn patch(&mut self, from: &[u8], to: &[u8]) -> Vec; } impl + ?Sized> MutBytesExt for T { fn patch(&mut self, from: &[u8], to: &[u8]) -> Vec { ffi::mut_u8_patch(self.as_mut(), from, to) } } // SAFETY: libc guarantees argc and argv are properly setup and are static #[allow(clippy::not_unsafe_ptr_arg_deref)] pub fn map_args(argc: i32, argv: *const *const c_char) -> Result, StrErr> { unsafe { slice::from_raw_parts(argv, argc as usize) } .iter() .map(|s| unsafe { Utf8CStr::from_ptr(*s) }.map(|s| s.as_str())) .collect() } pub trait EarlyExitExt { fn on_early_exit(self, print_help_msg: F) -> T; } impl EarlyExitExt for Result { fn on_early_exit(self, print_help_msg: F) -> T { match self { Ok(t) => t, Err(EarlyExit { output, status }) => match status { Ok(_) => { print_help_msg(); exit(0) } Err(_) => { eprintln!("{output}"); print_help_msg(); exit(1) } }, } } } pub struct FmtAdaptor<'a, T>(pub &'a mut T) where T: Write; impl fmt::Write for FmtAdaptor<'_, T> { fn write_str(&mut self, s: &str) -> fmt::Result { self.0.write_all(s.as_bytes()).map_err(|_| fmt::Error) } fn write_fmt(&mut self, args: Arguments<'_>) -> fmt::Result { self.0.write_fmt(args).map_err(|_| fmt::Error) } } pub struct AtomicArc { ptr: AtomicPtr, } impl AtomicArc { pub fn new(arc: Arc) -> AtomicArc { let raw = Arc::into_raw(arc); Self { ptr: AtomicPtr::new(raw as *mut _), } } pub fn load(&self) -> Arc { let raw = self.ptr.load(Ordering::Acquire); // SAFETY: the raw pointer is always created from Arc::into_raw let arc = ManuallyDrop::new(unsafe { Arc::from_raw(raw) }); ManuallyDrop::into_inner(arc.clone()) } fn swap_ptr(&self, raw: *const T) -> Arc { let prev = self.ptr.swap(raw as *mut _, Ordering::AcqRel); // SAFETY: the raw pointer is always created from Arc::into_raw unsafe { Arc::from_raw(prev) } } pub fn swap(&self, arc: Arc) -> Arc { let raw = Arc::into_raw(arc); self.swap_ptr(raw) } pub fn store(&self, arc: Arc) { // Drop the previous value let _ = self.swap(arc); } } impl Drop for AtomicArc { fn drop(&mut self) { // Drop the internal value let _ = self.swap_ptr(std::ptr::null()); } } impl Default for AtomicArc { fn default() -> Self { Self::new(Default::default()) } } pub struct Chunker { chunk: Box<[u8]>, chunk_size: usize, pos: usize, } impl Chunker { pub fn new(chunk_size: usize) -> Self { Chunker { // SAFETY: all bytes will be initialized before it is used, tracked by self.pos chunk: unsafe { Box::new_uninit_slice(chunk_size).assume_init() }, chunk_size, pos: 0, } } pub fn set_chunk_size(&mut self, chunk_size: usize) { self.chunk_size = chunk_size; self.pos = 0; if self.chunk.len() < chunk_size { self.chunk = unsafe { Box::new_uninit_slice(chunk_size).assume_init() }; } } // Returns (remaining buf, Option) pub fn add_data<'a, 'b: 'a>(&'a mut self, mut buf: &'b [u8]) -> (&'b [u8], Option<&'a [u8]>) { let mut chunk = None; if self.pos > 0 { // Try to fill the chunk let len = std::cmp::min(self.chunk_size - self.pos, buf.len()); self.chunk[self.pos..self.pos + len].copy_from_slice(&buf[..len]); self.pos += len; // If the chunk is filled, consume it if self.pos == self.chunk_size { chunk = Some(&self.chunk[..self.chunk_size]); self.pos = 0; } buf = &buf[len..]; } else if buf.len() >= self.chunk_size { // Directly consume a chunk from buf chunk = Some(&buf[..self.chunk_size]); buf = &buf[self.chunk_size..]; } else { // Copy buf into chunk self.chunk[self.pos..self.pos + buf.len()].copy_from_slice(buf); self.pos += buf.len(); return (&[], None); } (buf, chunk) } pub fn get_available(&mut self) -> &[u8] { let chunk = &self.chunk[..self.pos]; self.pos = 0; chunk } }