From d9fc95973801f9950064b0451a211cf555bcf412 Mon Sep 17 00:00:00 2001 From: Avril Date: Mon, 13 Jul 2020 18:16:59 +0100 Subject: [PATCH] added temp file --- Cargo.toml | 4 +- src/error.rs | 13 ++++++ src/main.rs | 1 + src/tempfile.rs | 94 +++++++++++++++++++++++++++++++++++++++++++ src/tempfile/error.rs | 51 +++++++++++++++++++++++ src/work_async.rs | 33 ++++++++------- 6 files changed, 178 insertions(+), 18 deletions(-) create mode 100644 src/tempfile.rs create mode 100644 src/tempfile/error.rs diff --git a/Cargo.toml b/Cargo.toml index 8fe7ce1..e4bfc60 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,4 +13,6 @@ async = ["tokio"] termprogress = { path = "../termprogress" } lazy_static = "1.4" tokio = {version = "0.2", features= ["full"], optional=true} -reqwest = {version = "0.10", features= ["stream"]} \ No newline at end of file +reqwest = {version = "0.10", features= ["stream"]} +memmap = "0.7" +getrandom = "0.1" \ No newline at end of file diff --git a/src/error.rs b/src/error.rs index 62b8bcc..857aabe 100644 --- a/src/error.rs +++ b/src/error.rs @@ -3,6 +3,9 @@ use std::{ fmt, io, }; +use super::{ + tempfile, +}; #[derive(Debug)] pub enum Error @@ -11,6 +14,7 @@ pub enum Error IO(io::Error), HTTP(reqwest::Error), HTTPStatus(reqwest::StatusCode), + TempFile(tempfile::error::Error), } impl error::Error for Error @@ -29,6 +33,7 @@ impl fmt::Display for Error fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { + Error::TempFile(tf) => write!(f, "tf recovered: {}", tf), Error::IO(io) => write!(f, "io: {}", io), Error::HTTP(http) => write!(f, "http internal error: {}", http), Error::HTTPStatus(status) => write!(f, "response returned status code {}", status), @@ -55,3 +60,11 @@ impl From for Error } } } + +impl From for Error +{ + fn from(er: tempfile::error::Error) -> Self + { + Self::TempFile(er) + } +} diff --git a/src/main.rs b/src/main.rs index 1f8db96..84bae27 100644 --- a/src/main.rs +++ b/src/main.rs @@ -10,6 +10,7 @@ mod config; mod args; mod url; mod error; +mod tempfile; mod loli; #[cfg(feature="async")] diff --git a/src/tempfile.rs b/src/tempfile.rs new file mode 100644 index 0000000..254a894 --- /dev/null +++ b/src/tempfile.rs @@ -0,0 +1,94 @@ +use std::{ + ops, + path::{ + PathBuf, + Path + }, + fmt, +}; +use lazy_static::lazy_static; + +pub mod error; + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct TempFile(PathBuf); + +const FILENAME_LEN_BYTES: usize = 16; +const PREFIX: &'static str = "lolistealer-"; + +lazy_static! { + pub static ref LOCATION: &'static Path = { + Box::leak(std::env::temp_dir().into_boxed_path()) + }; +} + +/// Generate a new random byte string for a filename. +fn generate_filename(mut to: impl fmt::Write) -> Result<(), error::Error> +{ + let mut buffer = [0u8; FILENAME_LEN_BYTES]; + getrandom::getrandom(&mut buffer[..])?; + + write!(to, "{}", PREFIX)?; + for byte in buffer.iter() { + write!(to, "{:2x}", *byte)?; + } + + Ok(()) +} + +impl TempFile +{ + /// Create a new temp file path + pub fn new() -> Self + { + let mut buffer = String::with_capacity((FILENAME_LEN_BYTES*2) + PREFIX.len()); + generate_filename(&mut buffer).expect("tf fatal"); + Self(LOCATION.join(buffer)) + } + /// Try to create a new temp file path + pub fn try_new() -> Result + { + let mut buffer = String::with_capacity((FILENAME_LEN_BYTES*2) + PREFIX.len()); + generate_filename(&mut buffer)?; + Ok(Self(LOCATION.join(buffer))) + } + + /// Consume the `TempFile` and return the raw path, performing no cleanup. + pub fn into_pathbuf(mut self) -> PathBuf + { + let pb = std::mem::replace(&mut self.0, PathBuf::default()); + std::mem::forget(self); + pb + } +} + +impl ops::Drop for TempFile +{ + fn drop(&mut self) + { + if self.0.exists() { + if self.0.is_file() { + let _ = std::fs::remove_file(&self.0); + } else { + let _ = std::fs::remove_dir_all(&self.0); + } + } + } +} + +impl AsRef for TempFile +{ + fn as_ref(&self) -> &Path + { + self.0.as_ref() + } +} + +impl ops::Deref for TempFile +{ + type Target = Path; + fn deref(&self) -> &Self::Target + { + self.0.as_ref() + } +} diff --git a/src/tempfile/error.rs b/src/tempfile/error.rs new file mode 100644 index 0000000..207fb86 --- /dev/null +++ b/src/tempfile/error.rs @@ -0,0 +1,51 @@ +use std::{ + error, + fmt, +}; + +#[derive(Debug)] +pub enum Error +{ + RNG, + Format(fmt::Error), + Unknown, +} + +impl error::Error for Error +{ + fn source(&self) -> Option<&(dyn error::Error + 'static)> + { + match &self { + _ => None, + Self::Format(rng) => Some(rng), + } + } +} +impl fmt::Display for Error +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result + { + write!(f, "temp file error: ")?; + match self + { + Error::RNG => write!(f, "rng failed"), + Error::Format(fmt) => write!(f, "formatting: {}", fmt), + _ => write!(f, "unknown"), + } + } +} + +impl From for Error +{ + fn from(f: fmt::Error) -> Self + { + Self::Format(f) + } +} +impl From for Error +{ + fn from(_f: getrandom::Error) -> Self + { + Self::RNG + } +} diff --git a/src/work_async.rs b/src/work_async.rs index faf80b9..ee7e924 100644 --- a/src/work_async.rs +++ b/src/work_async.rs @@ -14,8 +14,6 @@ mod tasklist; #[macro_use] mod progress; -//TODO: Create a module for temp files, pass the temp file to `perform` and do the base64 fixing after `perform` - /// Download a loli async pub async fn perform(url: impl AsRef, path: impl AsRef, mut progress: progress::CommandSender) -> Result<(), error::Error> { @@ -55,11 +53,7 @@ pub async fn perform(url: impl AsRef, path: impl AsRef, mut progress: if len.is_none() { prog_send!(progress.bump(1)); } - - //TODO: Decode `file`. - prog_send!(progress.println(format!("done for {}", url))); - prog_send!(link progress.pop_task(task)); Ok(()) @@ -68,7 +62,7 @@ pub async fn perform(url: impl AsRef, path: impl AsRef, mut progress: pub async fn work(conf: config::Config) -> Result<(), Box> { let rating = conf.rating; - let mut children = Vec::new(); + let mut children = Vec::with_capacity(conf.output.len()); let prog = progress::AsyncProgressCounter::new("Initialising...", 1); let mut prog_writer = prog.writer(); @@ -83,21 +77,26 @@ pub async fn work(conf: config::Config) -> Result<(), Box //println!("Starting download ({})...", url); prog.println(format!("Starting download ({})...", url)).await.expect("fatal"); - let path = match path { - config::OutputType::File(file) => file, //TODO: Download to temp file, memmap to str slice and then base64 decode. Also, determine if we need .png or .jpg from header. - config::OutputType::Directory(dir) => { - //TODO: Implement downloading to temp and renaming to hash - unimplemented!(); + /*let path = match path { + config::OutputType::File(file) => file, //TODO: Download to temp file, memmap to str slice and then base64 decode. Also, determine if we need .png or .jpg from header. + config::OutputType::Directory(dir) => { + //TODO: Implement downloading to temp and renaming to hash + unimplemented!(); + }, + };*/ + let temp = tempfile::TempFile::new(); + match perform(&url, &temp, prog).await { + Err(e) => panic!("Failed downloading {} -> {:?}: {}", url, temp, e), //TODO: Make real error handler + Ok(_) => { + //TODO: memmap `temp` and decode base64 into new file `path`. also determine the encoding. }, - }; - match perform(&url, &path, prog).await { - Err(e) => panic!("Failed downloading {} -> {:?}: {}", url, path, e), //TODO: Make real error handler - Ok(v) => v, } + + })); } - prog_send!(link prog_writer.println("Waiting for children...")); + prog_send!(link prog_writer.println("Children working...")); for child in children.into_iter() { match child.await {