use std::collections::BTreeMap; use std::future::Future; use std::marker::PhantomData; use std::path::{Path, PathBuf}; use std::pin::Pin; use std::process::Stdio; use std::sync::Arc; use std::task::{Context, Poll}; use std::time::Duration; use async_trait::async_trait; use clap::ArgMatches; use color_eyre::eyre::{self, eyre}; use fd_lock_rs::FdLock; use helpers::canonicalize; pub use helpers::NonDetachingJoinHandle; use lazy_static::lazy_static; pub use models::Version; use pin_project::pin_project; use sha2::Digest; use tokio::fs::File; use tokio::sync::{Mutex, OwnedMutexGuard, RwLock}; use tracing::instrument; use crate::shutdown::Shutdown; use crate::{Error, ErrorKind, ResultExt as _}; pub mod config; pub mod cpupower; pub mod docker; pub mod http_reader; pub mod io; pub mod logger; pub mod lshw; pub mod serde; pub mod crypto; #[derive(Clone, Copy, Debug, ::serde::Deserialize, ::serde::Serialize)] pub enum Never {} impl Never {} impl Never { pub fn absurd(self) -> T { match self {} } } impl std::fmt::Display for Never { fn fmt(&self, _f: &mut std::fmt::Formatter) -> std::fmt::Result { self.absurd() } } impl std::error::Error for Never {} #[async_trait::async_trait] pub trait Invoke { async fn invoke(&mut self, error_kind: crate::ErrorKind) -> Result, Error>; async fn invoke_timeout( &mut self, error_kind: crate::ErrorKind, timeout: Option, ) -> Result, Error>; } #[async_trait::async_trait] impl Invoke for tokio::process::Command { async fn invoke(&mut self, error_kind: crate::ErrorKind) -> Result, Error> { self.invoke_timeout(error_kind, None).await } async fn invoke_timeout( &mut self, error_kind: crate::ErrorKind, timeout: Option, ) -> Result, Error> { self.kill_on_drop(true); self.stdout(Stdio::piped()); self.stderr(Stdio::piped()); let res = match timeout { None => self.output().await?, Some(t) => tokio::time::timeout(t, self.output()) .await .with_kind(ErrorKind::Timeout)??, }; crate::ensure_code!( res.status.success(), error_kind, "{}", Some(&res.stderr) .filter(|a| !a.is_empty()) .or(Some(&res.stdout)) .filter(|a| !a.is_empty()) .and_then(|a| std::str::from_utf8(a).ok()) .unwrap_or(&format!("Unknown Error ({})", res.status)) ); Ok(res.stdout) } } pub trait Apply: Sized { fn apply O>(self, func: F) -> O { func(self) } } pub trait ApplyRef { fn apply_ref O>(&self, func: F) -> O { func(&self) } fn apply_mut O>(&mut self, func: F) -> O { func(self) } } impl Apply for T {} impl ApplyRef for T {} pub async fn daemon Fut, Fut: Future + Send + 'static>( mut f: F, cooldown: std::time::Duration, mut shutdown: tokio::sync::broadcast::Receiver>, ) -> Result<(), eyre::Error> { loop { tokio::select! { _ = shutdown.recv() => return Ok(()), _ = tokio::time::sleep(cooldown) => (), } match tokio::spawn(f()).await { Err(e) if e.is_panic() => return Err(eyre!("daemon panicked!")), _ => (), } } } pub trait SOption {} pub struct SSome(T); impl SSome { pub fn into(self) -> T { self.0 } } impl From for SSome { fn from(t: T) -> Self { SSome(t) } } impl SOption for SSome {} pub struct SNone(PhantomData); impl SNone { pub fn new() -> Self { SNone(PhantomData) } } impl SOption for SNone {} #[async_trait] pub trait AsyncFileExt: Sized { async fn maybe_open + Send + Sync>(path: P) -> std::io::Result>; async fn delete + Send + Sync>(path: P) -> std::io::Result<()>; } #[async_trait] impl AsyncFileExt for File { async fn maybe_open + Send + Sync>(path: P) -> std::io::Result> { match File::open(path).await { Ok(f) => Ok(Some(f)), Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None), Err(e) => Err(e), } } async fn delete + Send + Sync>(path: P) -> std::io::Result<()> { if let Ok(m) = tokio::fs::metadata(path.as_ref()).await { if m.is_dir() { tokio::fs::remove_dir_all(path).await } else { tokio::fs::remove_file(path).await } } else { Ok(()) } } } pub struct FmtWriter(W); impl std::io::Write for FmtWriter { fn write(&mut self, buf: &[u8]) -> std::io::Result { self.0 .write_str( std::str::from_utf8(buf) .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?, ) .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?; Ok(buf.len()) } fn flush(&mut self) -> std::io::Result<()> { Ok(()) } } pub fn display_none(_: T, _: &ArgMatches) {} pub struct Container(RwLock>); impl Container { pub fn new(value: Option) -> Self { Container(RwLock::new(value)) } pub async fn set(&self, value: T) -> Option { std::mem::replace(&mut *self.0.write().await, Some(value)) } pub async fn take(&self) -> Option { self.0.write().await.take() } pub async fn is_empty(&self) -> bool { self.0.read().await.is_none() } pub async fn drop(&self) { *self.0.write().await = None; } } #[pin_project] pub struct HashWriter { hasher: H, #[pin] writer: W, } impl HashWriter { pub fn new(hasher: H, writer: W) -> Self { HashWriter { hasher, writer } } pub fn finish(self) -> (H, W) { (self.hasher, self.writer) } pub fn inner(&self) -> &W { &self.writer } pub fn inner_mut(&mut self) -> &mut W { &mut self.writer } } impl tokio::io::AsyncWrite for HashWriter { fn poll_write( self: Pin<&mut Self>, cx: &mut Context, buf: &[u8], ) -> Poll> { let this = self.project(); let written = tokio::io::AsyncWrite::poll_write(this.writer, cx, buf); match written { // only update the hasher once Poll::Ready(res) => { if let Ok(n) = res { this.hasher.update(&buf[..n]); } Poll::Ready(res) } Poll::Pending => Poll::Pending, } } fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll> { self.project().writer.poll_flush(cx) } fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context) -> Poll> { self.project().writer.poll_shutdown(cx) } } pub trait IntoDoubleEndedIterator: IntoIterator { type IntoIter: Iterator + DoubleEndedIterator; fn into_iter(self) -> >::IntoIter; } impl IntoDoubleEndedIterator for T where T: IntoIterator, ::IntoIter: DoubleEndedIterator, { type IntoIter = ::IntoIter; fn into_iter(self) -> >::IntoIter { IntoIterator::into_iter(self) } } pub struct GeneralBoxedGuard(Option>); impl GeneralBoxedGuard { pub fn new(f: impl FnOnce() + 'static + Send + Sync) -> Self { GeneralBoxedGuard(Some(Box::new(f))) } pub fn drop(mut self) { self.0.take().unwrap()() } pub fn drop_without_action(mut self) { self.0 = None; } } impl Drop for GeneralBoxedGuard { fn drop(&mut self) { if let Some(destroy) = self.0.take() { destroy(); } } } pub struct GeneralGuard T, T = ()>(Option); impl T, T> GeneralGuard { pub fn new(f: F) -> Self { GeneralGuard(Some(f)) } pub fn drop(mut self) -> T { self.0.take().unwrap()() } pub fn drop_without_action(mut self) { self.0 = None; } } impl T, T> Drop for GeneralGuard { fn drop(&mut self) { if let Some(destroy) = self.0.take() { destroy(); } } } pub struct FileLock(OwnedMutexGuard<()>, Option>); impl Drop for FileLock { fn drop(&mut self) { if let Some(fd_lock) = self.1.take() { tokio::task::spawn_blocking(|| fd_lock.unlock(true).map_err(|(_, e)| e).unwrap()); } } } impl FileLock { #[instrument(skip_all)] pub async fn new(path: impl AsRef + Send + Sync, blocking: bool) -> Result { lazy_static! { static ref INTERNAL_LOCKS: Mutex>>> = Mutex::new(BTreeMap::new()); } let path = canonicalize(path.as_ref(), true) .await .with_kind(ErrorKind::Filesystem)?; let mut internal_locks = INTERNAL_LOCKS.lock().await; if !internal_locks.contains_key(&path) { internal_locks.insert(path.clone(), Arc::new(Mutex::new(()))); } let tex = internal_locks.get(&path).unwrap().clone(); drop(internal_locks); let tex_guard = if blocking { tex.lock_owned().await } else { tex.try_lock_owned() .with_kind(crate::ErrorKind::Filesystem)? }; let parent = path.parent().unwrap_or(Path::new("/")); if tokio::fs::metadata(parent).await.is_err() { tokio::fs::create_dir_all(parent) .await .with_ctx(|_| (crate::ErrorKind::Filesystem, parent.display().to_string()))?; } let f = File::create(&path) .await .with_ctx(|_| (crate::ErrorKind::Filesystem, path.display().to_string()))?; let file_guard = tokio::task::spawn_blocking(move || { fd_lock_rs::FdLock::lock(f, fd_lock_rs::LockType::Exclusive, blocking) }) .await .with_kind(crate::ErrorKind::Unknown)? .with_kind(crate::ErrorKind::Filesystem)?; Ok(FileLock(tex_guard, Some(file_guard))) } pub async fn unlock(mut self) -> Result<(), Error> { if let Some(fd_lock) = self.1.take() { tokio::task::spawn_blocking(|| fd_lock.unlock(true).map_err(|(_, e)| e)) .await .with_kind(crate::ErrorKind::Unknown)? .with_kind(crate::ErrorKind::Filesystem)?; } Ok(()) } } pub fn assure_send(x: T) -> T { x }