#[macro_use] extern crate cfg_if; #[cfg(feature="logging")] #[macro_use] extern crate tracing; #[cfg(feature="memfile")] #[macro_use] extern crate lazy_static; #[cfg(feature="memfile")] #[macro_use] extern crate stackalloc; /// Run this statement only if `tracing` is enabled macro_rules! if_trace { (? $expr:expr) => { cfg_if! { if #[cfg(all(feature="logging", debug_assertions))] { $expr; } } }; ($expr:expr) => { cfg_if! { if #[cfg(feature="logging")] { $expr; } } }; } #[cfg(feature="jemalloc")] extern crate jemallocator; #[cfg(feature="jemalloc")] const _:() = { #[global_allocator] static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc; }; use std::{ io, mem::MaybeUninit, os::unix::prelude::*, num::NonZeroUsize, }; #[allow(unused_imports)] use color_eyre::{ eyre::{ self, eyre, WrapErr, }, Section, SectionExt, Help, }; /// Get an `&'static str` of the current function name. macro_rules! function { () => {{ fn f() {} fn type_name_of(_: T) -> &'static str { ::std::any::type_name::() } let name = type_name_of(f); &name[..name.len() - 3] }} } mod buffers; use buffers::prelude::*; #[cfg(feature="memfile")] mod memfile; #[cfg(feature="bytes")] use bytes::{ Buf, BufMut, }; /* TODO: XXX: For colouring buffer::Perc #[derive(Debug)] struct StackStr(usize, std::mem::MaybeUninit<[u8; MAXLEN]>); impl StackStr { #[inline] pub const fn new() -> Self { Self(0, std::mem::MaybeUninit::uninit()) } #[inline(always)] pub const unsafe fn slice_mut(&mut self) -> &mut [u8] { &mut self.1[self.0..] } #[inline] pub const fn slice(&self) -> &[u8] { &self.1[self.0..] } #[inline] pub const unsafe fn as_str_unchecked(&self) -> &str { std::str::from_utf8_unchecked(&self.1[self.0..]) } #[inline] pub const unsafe fn as_mut_str_unchecked(&mut self) -> &mut str { std::str::from_utf8_unchecked_mut(&mut self.1[..self.0]) } #[inline] #[cfg_attr(feature="logging", instrument(level="debug"))] pub fn as_str(&self) -> &str { std::str::from_utf8(self.slice()).expect("Invalid string") } #[inline(always)] const fn left(&self) -> usize { SZ - self.0 } #[inline(always)] pub fn write_bytes(&mut self, s: &[u8]) -> usize { let b = &s[..std::cmp::min(match self.left() { 0 => return 0, x => x, }, s.len())]; unsafe { &mut self.slice_mut() [self.0..] }.copy_from_slice(b); let v = b.len(); self.0 += v; v } } impl std::fmt::Write for StackStr { #[inline] fn write_str(&mut self, s: &str) -> std::fmt::Result { self.write_bytes(s.as_bytes()); Ok(()) } #[inline] fn write_char(&mut self, c: char) -> std::fmt::Result { let l = c.len_utf8(); if l > self.left() { return Ok(()) } self.write_bytes(c.encode_utf8(unsafe { &mut self.slice_mut() [self.0..] })); self.0 += l; Ok(()) } } */ #[cfg_attr(feature="logging", instrument(level="info", skip(reader), fields(reader = ?std::any::type_name::())))] fn try_get_size(reader: &R) -> Option where R: AsRawFd { let fd = reader.as_raw_fd(); use libc::{ fstat64, stat64, }; if fd < 0 { return None; } let mut st: MaybeUninit = MaybeUninit::uninit(); unsafe { match fstat64(fd, st.as_mut_ptr()) { 0 => { NonZeroUsize::new(st.assume_init().st_size as usize) }, _ => None, } } } fn init() -> eyre::Result<()> { cfg_if!{ if #[cfg(feature="logging")] { fn install_tracing() { //! Install spantrace handling use tracing_error::ErrorLayer; use tracing_subscriber::prelude::*; use tracing_subscriber::{fmt, EnvFilter}; let fmt_layer = fmt::layer() .with_target(false) .with_writer(io::stderr); let filter_layer = EnvFilter::try_from_default_env() .or_else(|_| EnvFilter::try_new(if cfg!(debug_assertions) { "debug" } else { "info" })) .unwrap(); tracing_subscriber::registry() .with(fmt_layer) .with(filter_layer) .with(ErrorLayer::default()) .init(); } if !cfg!(feature="disable-logging") { install_tracing(); if_trace!(trace!("installed tracing")); } } } color_eyre::install() } #[cfg_attr(feature="logging", instrument(err))] #[inline] fn non_map_work() -> eyre::Result<()> { if_trace!(trace!("strategy: allocated buffer")); let (bytes, read) = { let stdin = io::stdin(); let mut bytes: buffers::DefaultMut = try_get_size(&stdin).create_buffer(); let read = io::copy(&mut stdin.lock(), &mut (&mut bytes).writer()) .with_section(|| bytes.len().header("Buffer size is")) .with_section(|| bytes.capacity().header("Buffer cap is")) .with_section(|| format!("{:?}", bytes).header("Buffer is")) .wrap_err("Failed to read into buffer")?; (bytes.freeze(), read as usize) }; if_trace!(info!("collected {read} from stdin. starting write.")); let written = io::copy(&mut (&bytes[..read]).reader() , &mut io::stdout().lock()) .with_section(|| read.header("Bytes read")) .with_section(|| bytes.len().header("Buffer length (frozen)")) .with_section(|| format!("{:?}", &bytes[..read]).header("Read Buffer")) .with_section(|| format!("{:?}", bytes).header("Full Buffer")) .wrap_err("Failed to write from buffer")?; if_trace!(info!("written {written} to stdout.")); if read != written as usize { return Err(io::Error::new(io::ErrorKind::BrokenPipe, format!("read {read} bytes, but only wrote {written}"))) .wrap_err("Writing failed: size mismatch"); } Ok(()) } #[cfg_attr(feature="logging", instrument(err))] #[inline] #[cfg(feature="memfile")] fn map_work() -> eyre::Result<()> { extern "C" { fn getpagesize() -> libc::c_int; } /// 8 pages const DEFAULT_BUFFER_SIZE: fn () -> Option = || { unsafe { std::num::NonZeroUsize::new((getpagesize() as usize) * 8) } }; if_trace!(trace!("strategy: mapped memory file")); use std::borrow::Borrow; #[inline(always)] fn tell_file(file: &mut T) -> io::Result where T: io::Seek + ?Sized { file.stream_position() } #[inline(always)] fn unwrap_int_string(i: impl Borrow>) -> String where T: std::fmt::Display, E: std::fmt::Display { i.borrow().as_ref().map(ToString::to_string) .unwrap_or_else(|e| format!("")) } let (mut file, read) = { let stdin = io::stdin(); let buffsz = try_get_size(&stdin).or_else(DEFAULT_BUFFER_SIZE); if_trace!(trace!("Attempted determining input size: {:?}", buffsz)); let mut file = memfile::create_memfile(Some("collect-buffer"), buffsz.map(|x| x.get()).unwrap_or(0)) .with_section(|| format!("{:?}", buffsz).header("Deduced input buffer size")) .wrap_err(eyre!("Failed to create in-memory buffer"))?; let read = io::copy(&mut stdin.lock(), &mut file) .with_section(|| format!("{:?}", file).header("Memory buffer file"))?; { use io::*; file.seek(SeekFrom::Start(0)) .with_section(|| read.header("Actual read bytes")) .wrap_err(eyre!("Failed to seek back to start of memory buffer file for output") .with_section(|| unwrap_int_string(file.stream_position()).header("Memfile position")) /*.with_section(|| file.stream_len().map(|x| x.to_string()) .unwrap_or_else(|e| format!("")).header("Memfile full length"))*/)?; } (file, usize::try_from(read) .wrap_err(eyre!("Failed to convert read bytes to `usize`") .with_section(|| read.header("Number of bytes was")) .with_section(|| u128::abs_diff(read.into(), usize::MAX as u128).header("Difference between `read` and `usize::MAX` is")) .with_suggestion(|| "It is likely you are running on a 32-bit ptr width machine and this input exceeds that of the maximum 32-bit unsigned integer value") .with_note(|| usize::MAX.header("Maximum value of `usize`")))?) }; if_trace!(info!("collected {read} from stdin. starting write.")); let written = io::copy(&mut file, &mut io::stdout().lock()) .with_section(|| read.header("Bytes read from stdin")) .with_section(|| unwrap_int_string(tell_file(&mut file)).header("Current buffer position")) .wrap_err("Failed to write buffer to stdout")?; if_trace!(info!("written {written} to stdout.")); if read != written as usize { return Err(io::Error::new(io::ErrorKind::BrokenPipe, format!("read {read} bytes, but only wrote {written}"))) .wrap_err("Writing failed: size mismatch"); } Ok(()) } #[cfg_attr(feature="logging", instrument(err))] fn main() -> eyre::Result<()> { init()?; if_trace!(debug!("initialised")); cfg_if!{ if #[cfg(feature="memfile")] { map_work() .wrap_err(eyre!("Operation failed").with_note(|| "With mapped memfd algorithm"))?; } else { non_map_work() .wrap_err(eyre!("Operation failed").with_note(|| "With alloc-buf (non-mapped) algorithm"))?; } } Ok(()) }