From 22926ce29e3f8e040ec349401aeb6a77f32eae72 Mon Sep 17 00:00:00 2001 From: 魏曹先生 <1992414357@qq.com> Date: Wed, 4 Mar 2026 21:26:04 +0800 Subject: Initialize Butchunker project structure and policy system --- src/bin/butck.rs | 109 ++++++ src/bin/butckrepo-guide.rs | 13 + src/bin/butckrepo-refresh.rs | 619 +++++++++++++++++++++++++++++++++ src/chunker.rs | 4 + src/chunker/constants.rs | 3 + src/chunker/context.rs | 226 ++++++++++++ src/chunker/entry.rs | 39 +++ src/chunker/rw.rs | 2 + src/chunker/rw/error.rs | 61 ++++ src/chunker/rw/storage.rs | 88 +++++ src/chunker/rw/storage/build.rs | 250 +++++++++++++ src/chunker/rw/storage/write.rs | 118 +++++++ src/chunker/rw/storage/write/simple.rs | 368 ++++++++++++++++++++ src/chunker/rw/storage/write/stream.rs | 12 + src/core.rs | 1 + src/core/hash.rs | 38 ++ src/lib.rs | 11 + src/log.rs | 33 ++ src/macros.rs | 47 +++ src/utils.rs | 2 + src/utils/file_input_solve.rs | 82 +++++ src/utils/size_display.rs | 14 + 22 files changed, 2140 insertions(+) create mode 100644 src/bin/butck.rs create mode 100644 src/bin/butckrepo-guide.rs create mode 100644 src/bin/butckrepo-refresh.rs create mode 100644 src/chunker.rs create mode 100644 src/chunker/constants.rs create mode 100644 src/chunker/context.rs create mode 100644 src/chunker/entry.rs create mode 100644 src/chunker/rw.rs create mode 100644 src/chunker/rw/error.rs create mode 100644 src/chunker/rw/storage.rs create mode 100644 src/chunker/rw/storage/build.rs create mode 100644 src/chunker/rw/storage/write.rs create mode 100644 src/chunker/rw/storage/write/simple.rs create mode 100644 src/chunker/rw/storage/write/stream.rs create mode 100644 src/core.rs create mode 100644 src/core/hash.rs create mode 100644 src/lib.rs create mode 100644 src/log.rs create mode 100644 src/macros.rs create mode 100644 src/utils.rs create mode 100644 src/utils/file_input_solve.rs create mode 100644 src/utils/size_display.rs (limited to 'src') diff --git a/src/bin/butck.rs b/src/bin/butck.rs new file mode 100644 index 0000000..6a81fbb --- /dev/null +++ b/src/bin/butck.rs @@ -0,0 +1,109 @@ +use std::process::exit; + +use butchunker::{ + chunker::{ + context::ButckContext, + entry::{entry, print_help, print_version}, + rw::error::{ButckRWError, ButckRWErrorKind}, + }, + log::init_logger, + special_argument, special_flag, +}; +use just_progress::{progress, renderer}; +use log::error; +use tokio::join; + +#[tokio::main] +async fn main() { + // Collect arguments + let mut args: Vec = std::env::args().skip(1).collect(); + + let version = special_flag!(args, "-v", "--version"); + let help = special_flag!(args, "-h", "--help"); + + if version { + print_version(); + exit(0) + } + + // Special arguments, early return + if help || args.is_empty() { + print_help(); + exit(0) + } + + // Init colored + #[cfg(windows)] + colored::control::set_virtual_terminal(true).unwrap(); + + // Output control flags + let quiet = special_flag!(args, "-q", "--quiet"); + let no_progress = special_flag!(args, "-np", "--no-progress"); + + // Logger + if !quiet { + let logger_level = match special_argument!(args, "-l", "--log-level") { + Some(level) => match level.trim().to_lowercase().as_str() { + "trace" => log::LevelFilter::Trace, + "debug" => log::LevelFilter::Debug, + "info" => log::LevelFilter::Info, + "warn" => log::LevelFilter::Warn, + "error" => log::LevelFilter::Error, + _ => log::LevelFilter::Info, + }, + None => log::LevelFilter::Info, + }; + init_logger(Some(logger_level)); + } + + let ctx = ButckContext::from_args(args.clone()); + + // When `--no-progress` or `--quiet` is enabled, + // the progress system will not be initialized + if no_progress || quiet { + handle_entry_result(entry(ctx, args).await); + } else { + let progress = progress::init(); + let renderer = renderer::ProgressSimpleRenderer::new().with_subprogress(true); + let bind = progress::bind(progress, move |name, state| renderer.update(name, state)); + join!( + async { + handle_entry_result(entry(ctx, args).await); + progress::close(); + }, + bind + ); + } +} + +fn handle_entry_result(r: Result<(), ButckRWError>) { + match r { + Ok(_) => {} + Err(e) => match e.kind() { + ButckRWErrorKind::NoButckStorageFound => { + error!("No butck storage found"); + error!("Use `--storage ` to specify or init butck storage"); + } + ButckRWErrorKind::ChunkingPolicyNotSpecified => { + error!("Chunking policy not specified"); + error!("Use `--policy ` to specify chunking policy"); + error!("or use `butck policies` to output the available policies"); + } + ButckRWErrorKind::ReadingMethodAmbiguous => error!("Reading method ambiguous"), + ButckRWErrorKind::OutputCountMismatch => { + error!("Output count mismatch"); + error!("When processing a single file, use `--output-file` to specify output path"); + error!( + "When processing multiple files, use `--output-dir` to specify output directory" + ); + } + ButckRWErrorKind::ChunkNotFound(chunk_id) => { + error!("Chunk not found in storage: {}", chunk_id) + } + ButckRWErrorKind::RebuildFailed(reason) => error!("Failed to rebuild file: {}", reason), + ButckRWErrorKind::ChunkFailed(_chunk_failed) => error!("Chunk failed"), + ButckRWErrorKind::IOError(error) => error!("IO error: {}", error), + ButckRWErrorKind::InvalidBidxFormat => error!("Invalid bidx format"), + }, + } +} diff --git a/src/bin/butckrepo-guide.rs b/src/bin/butckrepo-guide.rs new file mode 100644 index 0000000..d694ba5 --- /dev/null +++ b/src/bin/butckrepo-guide.rs @@ -0,0 +1,13 @@ +use colored::Colorize; + +fn main() { + println!("Welcome to Butchunker!"); + println!( + "Please add your policy crates to the `{}` directory", + "./policy/".bright_green() + ); + println!( + "Then run `{}` to update the policy registry", + "cargo run --bin butckrepo-refresh".bright_green() + ); +} diff --git a/src/bin/butckrepo-refresh.rs b/src/bin/butckrepo-refresh.rs new file mode 100644 index 0000000..9184efb --- /dev/null +++ b/src/bin/butckrepo-refresh.rs @@ -0,0 +1,619 @@ +use colored::Colorize; +use just_fmt::fmt_path::fmt_path_str; +use just_template::{Template, tmpl, tmpl_param}; +use std::{ + env::current_dir, + path::{Path, PathBuf}, +}; +use tokio::fs; + +const LIB_RS_TEMPLATE_PATH: &str = "policy/_policies/src/lib.rs.t"; +const CARGO_TOML_TEMPLATE_PATH: &str = "policy/_policies/Cargo.toml.t"; +const LIB_RS_PATH: &str = "./policy/_policies/src/lib.rs"; +const CARGO_TOML_PATH: &str = "./policy/_policies/Cargo.toml"; + +#[tokio::main] +async fn main() { + let current_dir = current_dir().unwrap(); + precheck(¤t_dir).await; + + println!("Updating policies ..."); + let (mut lib_rs_template, mut cargo_toml_template) = { + let lib_rs_template_path = current_dir.join("policy/_policies/src/lib.rs.t"); + let cargo_toml_template_path = current_dir.join("policy/_policies/Cargo.toml.t"); + + let lib_rs_content = fs::read_to_string(&lib_rs_template_path) + .await + .unwrap_or_else(|_| { + eprintln!( + "{}", + format!( + "Error: Failed to read template file: {}", + lib_rs_template_path.display() + ) + .red() + ); + std::process::exit(1); + }); + + let cargo_toml_content = fs::read_to_string(&cargo_toml_template_path) + .await + .unwrap_or_else(|_| { + eprintln!( + "{}", + format!( + "Error: Failed to read template file: {}", + cargo_toml_template_path.display() + ) + .red() + ); + std::process::exit(1); + }); + + ( + Template::from(lib_rs_content), + Template::from(cargo_toml_content), + ) + }; + + let cargo_toml_pathes = find_cargo_toml_dirs(¤t_dir.join("policy")).await; + println!( + "Found {} crates, register to `{}`", + cargo_toml_pathes.len(), + CARGO_TOML_PATH.bright_green() + ); + + tmpl_param!(lib_rs_template, policy_count = cargo_toml_pathes.len()); + + let collect_futures = cargo_toml_pathes.iter().map(collect).collect::>(); + + for policy in futures::future::join_all(collect_futures).await { + let Some(policy) = policy else { continue }; + tmpl!(cargo_toml_template += { + deps { (crate_name = policy.crate_name, path = policy.path) } + }); + // Determine which export template to use based on detected functions + if policy.matched_func_stream.is_some() { + let stream_struct_id = format!( + "{}::{}", + policy.crate_name, + policy.stream_struct_id.unwrap() + ); + if policy.matched_func.is_empty() { + // Only stream function + tmpl!(lib_rs_template += { + exports_stream { ( + crate_name = policy.crate_name, + matched_func_stream = policy.matched_func_stream.unwrap(), + has_await_stream = + if policy.matched_func_stream_has_await { ".await" } else { "" }, + stream_struct_id = stream_struct_id + ) }, + match_arms { ( + crate_name = policy.crate_name, + ) }, + match_arms_stream { ( + crate_name = policy.crate_name, + stream_struct_id = stream_struct_id + ) }, + policy_names { ( + name = policy.crate_name, + ) } + }); + } else { + // Both simple and stream functions + tmpl!(lib_rs_template += { + exports_both { ( + crate_name = policy.crate_name, + matched_func = policy.matched_func, + has_await = + if policy.matched_func_has_await { ".await" } else { "" }, + matched_func_stream = policy.matched_func_stream.unwrap(), + has_await_stream = + if policy.matched_func_stream_has_await { ".await" } else { "" }, + stream_struct_id = stream_struct_id + ) }, + match_arms { ( + crate_name = policy.crate_name, + ) }, + match_arms_stream { ( + crate_name = policy.crate_name, + stream_struct_id = stream_struct_id + ) }, + policy_names { ( + name = policy.crate_name, + ) } + }); + } + } else { + // Only simple function + tmpl!(lib_rs_template += { + exports_simple { ( + crate_name = policy.crate_name, + matched_func = policy.matched_func, + has_await = + if policy.matched_func_has_await { ".await" } else { "" } + ) }, + match_arms { ( + crate_name = policy.crate_name, + ) }, + policy_names { ( + name = policy.crate_name, + ) } + }); + } + } + + let (write_cargo, write_lib) = tokio::join!( + fs::write(CARGO_TOML_PATH, cargo_toml_template.expand().unwrap()), + fs::write(LIB_RS_PATH, lib_rs_template.expand().unwrap()) + ); + write_cargo.unwrap(); + write_lib.unwrap(); +} + +struct CollectedPolicy { + crate_name: String, + path: String, + matched_func: String, + matched_func_has_await: bool, + matched_func_stream: Option, + matched_func_stream_has_await: bool, + stream_struct_id: Option, +} + +async fn collect(policy_crate_path: &PathBuf) -> Option { + let lib_rs_path = policy_crate_path.join("src").join("lib.rs"); + let lib_rs_content = fs::read_to_string(&lib_rs_path).await.ok()?; + + let cargo_toml_content = fs::read_to_string(policy_crate_path.join("Cargo.toml")) + .await + .ok()?; + let cargo_toml: toml::Value = toml::from_str(&cargo_toml_content).ok()?; + let crate_name = cargo_toml + .get("package")? + .get("name")? + .as_str()? + .to_string(); + let crate_path = fmt_path_str( + policy_crate_path + .strip_prefix(current_dir().unwrap()) + .unwrap() + .to_string_lossy(), + ) + .ok()?; + + let ( + matched_func, + matched_func_has_await, + matched_func_stream, + matched_func_stream_has_await, + stream_struct_id, + ) = collect_matched_func(lib_rs_content.as_str())?; + + println!( + "{} {} (at: `{}`) with func `{}{}{}{}(..)`", + "Register:".bright_blue().bold(), + crate_name, + crate_path.bright_green(), + "pub ".bright_magenta(), + if matched_func_has_await { "async " } else { "" }.bright_magenta(), + "fn ".bright_magenta(), + matched_func.bright_blue(), + ); + if let Some(stream_func) = &matched_func_stream { + println!( + " and stream func `{}{}{}{}(..)`", + "pub ".bright_magenta(), + if matched_func_stream_has_await { + "async " + } else { + "" + } + .bright_magenta(), + "fn ".bright_magenta(), + stream_func.bright_blue() + ); + } + + Some(CollectedPolicy { + crate_name, + path: crate_path, + matched_func, + matched_func_has_await, + matched_func_stream, + matched_func_stream_has_await, + stream_struct_id, + }) +} + +fn collect_matched_func( + lib_rs_content: &str, +) -> Option<(String, bool, Option, bool, Option)> { + let syntax_tree = syn::parse_file(lib_rs_content).ok()?; + + let mut matched_func = None; + let mut matched_func_has_await = false; + let mut matched_func_stream = None; + let mut matched_func_stream_has_await = false; + let mut stream_struct_id = None; + + // Iterate over all items, looking for functions that match the criteria + for item in &syntax_tree.items { + let syn::Item::Fn(func) = item else { continue }; + + // Check if the function visibility is pub + if !matches!(func.vis, syn::Visibility::Public(_)) { + continue; + } + + let sig = &func.sig; + + // Check for simple chunk function (returns Vec) + if check_simple_chunk_function(sig) { + matched_func = Some(sig.ident.to_string()); + matched_func_has_await = sig.asyncness.is_some(); + } + // Check for stream chunk function (returns Option) + else if let Some(struct_id) = check_stream_chunk_function(sig, &syntax_tree) { + matched_func_stream = Some(sig.ident.to_string()); + matched_func_stream_has_await = sig.asyncness.is_some(); + stream_struct_id = Some(struct_id); + } + } + + if matched_func.is_some() || matched_func_stream.is_some() { + Some(( + matched_func.unwrap_or_default(), + matched_func_has_await, + matched_func_stream, + matched_func_stream_has_await, + stream_struct_id, + )) + } else { + None + } +} + +fn check_simple_chunk_function(sig: &syn::Signature) -> bool { + // Check if the return type is Vec + let return_type_matches = match &sig.output { + syn::ReturnType::Type(_, ty) => { + let syn::Type::Path(type_path) = &**ty else { + return false; + }; + let segments = &type_path.path.segments; + + segments.len() == 1 + && segments[0].ident == "Vec" + && matches!(&segments[0].arguments, syn::PathArguments::AngleBracketed(args) + if args.args.len() == 1 && + matches!(&args.args[0], syn::GenericArgument::Type(syn::Type::Path(inner_type)) + if inner_type.path.segments.len() == 1 && + inner_type.path.segments[0].ident == "u32" + ) + ) + } + _ => false, + }; + + if !return_type_matches { + return false; + } + + // Check that there are exactly 2 parameters + if sig.inputs.len() != 2 { + return false; + } + + // Check that the first parameter type is &[u8] + let first_param_matches = match &sig.inputs[0] { + syn::FnArg::Typed(pat_type) => { + let syn::Type::Reference(type_ref) = &*pat_type.ty else { + return false; + }; + let syn::Type::Slice(slice_type) = &*type_ref.elem else { + return false; + }; + let syn::Type::Path(type_path) = &*slice_type.elem else { + return false; + }; + + type_path.path.segments.len() == 1 && type_path.path.segments[0].ident == "u8" + } + _ => false, + }; + + // Check that the second parameter type is &HashMap<&str, &str> + let second_param_matches = match &sig.inputs[1] { + syn::FnArg::Typed(pat_type) => { + let syn::Type::Reference(type_ref) = &*pat_type.ty else { + return false; + }; + let syn::Type::Path(type_path) = &*type_ref.elem else { + return false; + }; + + type_path.path.segments.len() == 1 + && type_path.path.segments[0].ident == "HashMap" + && matches!(&type_path.path.segments[0].arguments, syn::PathArguments::AngleBracketed(args) + if args.args.len() == 2 && + matches!(&args.args[0], syn::GenericArgument::Type(syn::Type::Reference(first_ref)) + if matches!(&*first_ref.elem, syn::Type::Path(first_path) + if first_path.path.segments.len() == 1 && + first_path.path.segments[0].ident == "str" + ) + ) && + matches!(&args.args[1], syn::GenericArgument::Type(syn::Type::Reference(second_ref)) + if matches!(&*second_ref.elem, syn::Type::Path(second_path) + if second_path.path.segments.len() == 1 && + second_path.path.segments[0].ident == "str" + ) + ) + ) + } + _ => false, + }; + + first_param_matches && second_param_matches +} + +fn check_stream_chunk_function(sig: &syn::Signature, syntax_tree: &syn::File) -> Option { + // Check if the return type is Option + let return_type_matches = match &sig.output { + syn::ReturnType::Type(_, ty) => { + let syn::Type::Path(type_path) = &**ty else { + return None; + }; + let segments = &type_path.path.segments; + + segments.len() == 1 + && segments[0].ident == "Option" + && matches!(&segments[0].arguments, syn::PathArguments::AngleBracketed(args) + if args.args.len() == 1 && + matches!(&args.args[0], syn::GenericArgument::Type(syn::Type::Path(inner_type)) + if inner_type.path.segments.len() == 1 && + inner_type.path.segments[0].ident == "u32" + ) + ) + } + _ => false, + }; + + if !return_type_matches { + return None; + } + + // Check that there are exactly 4 parameters + if sig.inputs.len() != 4 { + return None; + } + + // Check that the first parameter type is &[u8] + let first_param_matches = match &sig.inputs[0] { + syn::FnArg::Typed(pat_type) => { + let syn::Type::Reference(type_ref) = &*pat_type.ty else { + return None; + }; + let syn::Type::Slice(slice_type) = &*type_ref.elem else { + return None; + }; + let syn::Type::Path(type_path) = &*slice_type.elem else { + return None; + }; + + // Check it's u8 + type_path.path.segments.len() == 1 && type_path.path.segments[0].ident == "u8" + } + _ => false, + }; + + // Check that the second parameter type is u32 + let second_param_matches = match &sig.inputs[1] { + syn::FnArg::Typed(pat_type) => { + let syn::Type::Path(type_path) = &*pat_type.ty else { + return None; + }; + type_path.path.segments.len() == 1 && type_path.path.segments[0].ident == "u32" + } + _ => false, + }; + + // Check that the third parameter type is &mut T where T is a struct defined in this crate + let third_param_info = match &sig.inputs[2] { + syn::FnArg::Typed(pat_type) => { + let syn::Type::Reference(type_ref) = &*pat_type.ty else { + return None; + }; + + // Check it's mutable reference + type_ref.mutability?; + + // Get the inner type + let syn::Type::Path(type_path) = &*type_ref.elem else { + return None; + }; + + // Get the struct identifier + if type_path.path.segments.len() != 1 { + return None; + } + + let struct_ident = type_path.path.segments[0].ident.to_string(); + + // Check if this struct is defined in the current crate and implements Default + if is_struct_defined_in_crate(&struct_ident, syntax_tree) { + Some(struct_ident) + } else { + None + } + } + _ => None, + }; + + let struct_ident = third_param_info?; + + // Check that the fourth parameter type is &HashMap<&str, &str> + let fourth_param_matches = match &sig.inputs[3] { + syn::FnArg::Typed(pat_type) => { + let syn::Type::Reference(type_ref) = &*pat_type.ty else { + return None; + }; + let syn::Type::Path(type_path) = &*type_ref.elem else { + return None; + }; + + type_path.path.segments.len() == 1 + && type_path.path.segments[0].ident == "HashMap" + && matches!(&type_path.path.segments[0].arguments, syn::PathArguments::AngleBracketed(args) + if args.args.len() == 2 && + matches!(&args.args[0], syn::GenericArgument::Type(syn::Type::Reference(first_ref)) + if matches!(&*first_ref.elem, syn::Type::Path(first_path) + if first_path.path.segments.len() == 1 && + first_path.path.segments[0].ident == "str" + ) + ) && + matches!(&args.args[1], syn::GenericArgument::Type(syn::Type::Reference(second_ref)) + if matches!(&*second_ref.elem, syn::Type::Path(second_path) + if second_path.path.segments.len() == 1 && + second_path.path.segments[0].ident == "str" + ) + ) + ) + } + _ => false, + }; + + if first_param_matches && second_param_matches && fourth_param_matches { + Some(struct_ident) + } else { + None + } +} + +fn is_struct_defined_in_crate(struct_ident: &str, syntax_tree: &syn::File) -> bool { + for item in &syntax_tree.items { + match item { + syn::Item::Struct(item_struct) => { + if item_struct.ident == struct_ident { + // Check if it implements Default via derive attribute + return has_default_derive(&item_struct.attrs) + || has_default_trait_bound(&item_struct.generics); + } + } + _ => continue, + } + } + false +} + +fn has_default_derive(attrs: &[syn::Attribute]) -> bool { + for attr in attrs { + if attr.path().is_ident("derive") { + // Parse the attribute meta to check for Default + if let syn::Meta::List(list) = &attr.meta { + // Convert tokens to string and check for Default + let tokens = list.tokens.to_string(); + if tokens.contains("Default") { + return true; + } + } + } + } + false +} + +fn has_default_trait_bound(generics: &syn::Generics) -> bool { + for param in &generics.params { + if let syn::GenericParam::Type(type_param) = param { + for bound in &type_param.bounds { + if let syn::TypeParamBound::Trait(trait_bound) = bound { + let path = &trait_bound.path; + if path.segments.len() == 1 && path.segments[0].ident == "Default" { + return true; + } + } + } + } + } + false +} + +async fn find_cargo_toml_dirs(root: &Path) -> Vec { + let mut result = Vec::new(); + let mut dirs_to_visit = vec![root.to_path_buf()]; + + while let Some(current_dir) = dirs_to_visit.pop() { + let cargo_toml_path = current_dir.join("Cargo.toml"); + if fs::metadata(&cargo_toml_path).await.is_ok() { + result.push(current_dir); + continue; + } + + let mut read_dir = match fs::read_dir(¤t_dir).await { + Ok(rd) => rd, + Err(_) => continue, + }; + + while let Ok(Some(entry)) = read_dir.next_entry().await { + if let Ok(file_type) = entry.file_type().await + && file_type.is_dir() + { + let path = entry.path(); + if let Some(file_name) = path.file_name() + && let Some(name_str) = file_name.to_str() + && name_str.starts_with('_') + { + continue; + } + dirs_to_visit.push(path); + } + } + } + + result +} + +async fn precheck(current_dir: &Path) { + let cargo_toml_path = current_dir.join("Cargo.toml"); + let cargo_toml_content = fs::read_to_string(&cargo_toml_path) + .await + .unwrap_or_else(|_| { + eprintln!( + "{}", + "Error: Cargo.toml not found in current directory".red() + ); + std::process::exit(1); + }); + let cargo_toml: toml::Value = toml::from_str(&cargo_toml_content).unwrap_or_else(|_| { + eprintln!("{}", "Error: Failed to parse Cargo.toml".red()); + std::process::exit(1); + }); + let package_name = cargo_toml + .get("package") + .unwrap_or_else(|| { + eprintln!("{}", "Error: No package section in Cargo.toml".red()); + std::process::exit(1); + }) + .get("name") + .unwrap_or_else(|| { + eprintln!("{}", "Error: No package.name in Cargo.toml".red()); + std::process::exit(1); + }) + .as_str() + .unwrap_or_else(|| { + eprintln!("{}", "Error: package.name is not a string".red()); + std::process::exit(1); + }); + if package_name != "butchunker" { + eprintln!( + "{}", + format!( + "Error: package.name must be 'butchunker', found '{}'", + package_name + ) + .red() + ); + std::process::exit(1); + } +} diff --git a/src/chunker.rs b/src/chunker.rs new file mode 100644 index 0000000..3143f68 --- /dev/null +++ b/src/chunker.rs @@ -0,0 +1,4 @@ +pub mod constants; +pub mod context; +pub mod entry; +pub mod rw; diff --git a/src/chunker/constants.rs b/src/chunker/constants.rs new file mode 100644 index 0000000..5e4870e --- /dev/null +++ b/src/chunker/constants.rs @@ -0,0 +1,3 @@ +pub const BUTCK_STORAGE_DIR_NAME: &str = ".butck"; +pub const BUTCK_INDEX_FILE_SUFFIX: &str = "bidx"; +pub const BUTCK_INDEX_MAGIC: [u8; 4] = [0xfe, 0xe1, 0xf0, 0x0d]; diff --git a/src/chunker/context.rs b/src/chunker/context.rs new file mode 100644 index 0000000..79254f5 --- /dev/null +++ b/src/chunker/context.rs @@ -0,0 +1,226 @@ +use std::{collections::HashMap, env::current_dir, path::PathBuf, process::exit, str::FromStr}; + +use log::{error, warn}; + +use crate::{ + chunker::constants::BUTCK_STORAGE_DIR_NAME, core::hash::ChunkWriteHash, special_argument, + special_flag, utils::file_input_solve::parse_path_input, +}; + +#[derive(Debug, Default)] +pub struct ButckContext { + /// All input files + pub file_paths: Vec, + + /// Path of Butck Storage + pub storage_path: Option, + + // Display chunk boundaries + pub display_boundaries: bool, + + /// Whether to read in stream mode + pub stream_read: Option, + + /// Whether to read files using memory mapping + pub memmap_read: bool, + + /// Register name + pub register_name: Option, + + /// Chunking policy name + pub policy_name: Option, + + /// Hash algorithm used for chunking + pub chunk_hash: ChunkWriteHash, + + /// Output directory + pub output_dir: PathBuf, + + /// Output file (not available for some commands) + pub output_file: Option, + + /// Override parameters + pub params: HashMap, +} + +impl ButckContext { + /// Apply the args of ChunkerContext to itself + pub fn from_args(mut args: Vec) -> Self { + let mut ctx = ButckContext::default(); + let recursive = ctx.read_recursive(&mut args); + ctx.apply_stream_read(&mut args); + ctx.apply_memmap_read(&mut args); + ctx.apply_register_name(&mut args); + ctx.apply_policy_name(&mut args); + ctx.apply_chunk_hash(&mut args); + ctx.apply_storage_dir(&mut args); + ctx.apply_output_paths(&mut args); + ctx.apply_params(&mut args); + ctx.apply_display_boundaries(&mut args); + + // Finally, parse path input + ctx.file_paths = parse_path_input(args, recursive, vec![BUTCK_STORAGE_DIR_NAME]); + ctx + } + + fn read_recursive(&mut self, args: &mut Vec) -> bool { + special_flag!(args, "-r", "--recursive") + } + + fn apply_stream_read(&mut self, args: &mut Vec) { + if let Some(size_str) = special_argument!(args, "-S", "--stream-read") + && let Ok(size) = size_str.parse::() { + self.stream_read = Some(size); + } + } + + fn apply_memmap_read(&mut self, args: &mut Vec) -> bool { + special_flag!(args, "-m", "--memmap-read") + } + + fn apply_register_name(&mut self, args: &mut Vec) { + self.register_name = special_argument!(args, "-R", "--register"); + } + + fn apply_policy_name(&mut self, args: &mut Vec) { + self.policy_name = special_argument!(args, "-p", "--policy"); + } + + fn apply_chunk_hash(&mut self, args: &mut Vec) { + let chunk_hash_str = special_argument!(args, "-H", "--chunk-hash"); + self.chunk_hash = match chunk_hash_str { + Some(ref s) => match s.as_str() { + "blake3" => ChunkWriteHash::Blake3, + "sha256" => ChunkWriteHash::Sha256, + _ => ChunkWriteHash::default(), + }, + None => ChunkWriteHash::default(), + }; + } + + fn apply_output_paths(&mut self, args: &mut Vec) { + let output_dir_str = special_argument!(args, "-o", "--output-dir"); + let output_file_str = special_argument!(args, "-O", "--output-file"); + + let current_dir = current_dir().unwrap(); + + let output_dir = if let Some(output_dir_str) = output_dir_str { + let path = PathBuf::from(output_dir_str); + if path.exists() { Some(path) } else { None } + } else { + None + }; + + self.output_dir = if let Some(output_dir) = output_dir { + output_dir + } else if let Some(storage_path) = &self.storage_path { + storage_path.clone() + } else { + current_dir + }; + + self.output_file = output_file_str.map(PathBuf::from) + } + + fn apply_params(&mut self, args: &mut Vec) { + while let Some(arg) = special_argument!(args, "+p", "+param") { + let split = arg.split('=').collect::>(); + if split.len() == 2 { + self.params + .insert(split[0].to_string(), split[1].to_string()); + } + } + } + + fn apply_storage_dir(&mut self, args: &mut Vec) { + self.storage_path = { + let storage_override = match special_argument!(args, "-s", "--storage") { + Some(o) => { + let path = PathBuf::from_str(o.as_str()); + if let Ok(p) = &path { + Self::init_butck_storage(p.clone()); + } + path.ok() + } + None => None, + }; + Self::find_butck_storage_dir(storage_override) + }; + } + + fn apply_display_boundaries(&mut self, args: &mut Vec) { + self.display_boundaries = special_flag!(args, "-D", "--display-boundaries"); + } + + fn init_butck_storage(path: PathBuf) -> Option { + if !path.exists() { + // If the path does not exist, create it and initialize Butck Storage here + if let Err(e) = std::fs::create_dir_all(&path) { + error!("Failed to create directory '{}': {}", path.display(), e); + exit(1); + } + let butck_dir = path.join(BUTCK_STORAGE_DIR_NAME); + if let Err(e) = std::fs::create_dir_all(&butck_dir) { + error!( + "Failed to create '{}' directory: {}", + BUTCK_STORAGE_DIR_NAME, e + ); + exit(1); + } + Some(path) + } else { + let butck_dir = path.join(BUTCK_STORAGE_DIR_NAME); + + // Check if Butck Storage already exists + if butck_dir.exists() { + // Butck Storage already exists, return the path + Some(path) + } else { + // Butck Storage doesn't exist, create it with a warning if directory is not empty + let is_empty = path + .read_dir() + .map(|mut entries| entries.next().is_none()) + .unwrap_or(false); + + if !is_empty { + // Warn about creating storage in non-empty directory + warn!( + "Creating '{}' storage in non-empty directory: {}", + BUTCK_STORAGE_DIR_NAME, + path.display() + ); + } + + // Create Butck Storage directory + if let Err(e) = std::fs::create_dir_all(&butck_dir) { + error!( + "Failed to create '{}' directory: {}", + BUTCK_STORAGE_DIR_NAME, e + ); + exit(1); + } + Some(path) + } + } + } + + // Get the ButckStorage directory based on context + fn find_butck_storage_dir(from: Option) -> Option { + let mut current_dir = match from { + Some(path) => path, + None => std::env::current_dir().ok()?, + }; + + loop { + let butck_dir = current_dir.join(BUTCK_STORAGE_DIR_NAME); + if butck_dir.is_dir() { + return Some(current_dir); + } + + if !current_dir.pop() { + break; + } + } + None + } +} diff --git a/src/chunker/entry.rs b/src/chunker/entry.rs new file mode 100644 index 0000000..4fdb1f8 --- /dev/null +++ b/src/chunker/entry.rs @@ -0,0 +1,39 @@ +use std::process::exit; + +use log::info; + +use crate::chunker::{ + context::ButckContext, + rw::{self, error::ButckRWError}, +}; + +pub async fn entry(ctx: ButckContext, args: Vec) -> Result<(), ButckRWError> { + if let Some(subcommand) = args.first() { + return match subcommand.as_str() { + "write" => rw::storage::write(ctx).await, + "build" => rw::storage::build(ctx).await, + "policies" => { + butck_policies::policies() + .iter() + .for_each(|p| info!("{}", p)); + return Ok(()); + } + _ => { + print_help(); + exit(1) + } + }; + } + Ok(()) +} + +pub fn print_help() { + println!("{}", include_str!("../../resources/helps/butck.txt").trim()); +} + +pub fn print_version() { + println!( + "{}", + include_str!("../../resources/version_info.txt").trim() + ); +} diff --git a/src/chunker/rw.rs b/src/chunker/rw.rs new file mode 100644 index 0000000..85e734e --- /dev/null +++ b/src/chunker/rw.rs @@ -0,0 +1,2 @@ +pub mod error; +pub mod storage; diff --git a/src/chunker/rw/error.rs b/src/chunker/rw/error.rs new file mode 100644 index 0000000..7f263a5 --- /dev/null +++ b/src/chunker/rw/error.rs @@ -0,0 +1,61 @@ +use butck_policies::error::ChunkFailed; + +use crate::chunker::context::ButckContext; + +#[derive(Debug)] +pub struct ButckRWError { + kind: ButckRWErrorKind, + ctx: ButckContext, +} + +#[derive(thiserror::Error, Debug)] +pub enum ButckRWErrorKind { + #[error("No butck storage found")] + NoButckStorageFound, + + #[error("Chunking policy not specified")] + ChunkingPolicyNotSpecified, + + #[error("Cannot enable both MemmapRead and StreamRead")] + ReadingMethodAmbiguous, + + #[error("Multiple input files specified but only one output file allowed")] + OutputCountMismatch, + + #[error("Invalid bidx file format")] + InvalidBidxFormat, + + #[error("Chunk not found in storage: {0}")] + ChunkNotFound(String), + + #[error("Failed to rebuild file: {0}")] + RebuildFailed(String), + + #[error("Chunking failed: {0}")] + ChunkFailed(#[from] ChunkFailed), + + #[error("IO error: {0}")] + IOError(#[from] std::io::Error), +} + +impl std::fmt::Display for ButckRWError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{:?}", self.kind) + } +} + +impl ButckRWError { + pub fn ctx(&self) -> &ButckContext { + &self.ctx + } + + pub fn kind(&self) -> &ButckRWErrorKind { + &self.kind + } +} + +impl ButckRWErrorKind { + pub fn pack(self, ctx: ButckContext) -> ButckRWError { + ButckRWError { kind: self, ctx } + } +} diff --git a/src/chunker/rw/storage.rs b/src/chunker/rw/storage.rs new file mode 100644 index 0000000..13452d0 --- /dev/null +++ b/src/chunker/rw/storage.rs @@ -0,0 +1,88 @@ +pub mod build; +pub mod write; + +pub use build::build; +pub use write::write; + +use std::path::{Path, PathBuf}; + +/// Information about a chunk for index file +#[derive(Debug, Clone)] +pub struct ChunkInfo { + /// Index of the chunk in the file + pub index: usize, + /// Hash of the chunk (hex string) + pub hash: String, + /// Size of the chunk in bytes + pub size: usize, + /// Start position in the original file + pub start: usize, + /// End position in the original file (exclusive) + pub end: usize, +} + +/// 根据hash值计算chunk文件的存储路径 +/// +/// # 参数 +/// - `storage_dir`: 存储目录 +/// - `hash_hex`: chunk的hash值（16进制字符串） +/// +/// # 返回 +/// 返回chunk文件的完整路径 +pub fn get_chunk_path(storage_dir: &Path, hash_hex: &str) -> PathBuf { + let first_slice = &hash_hex[0..2]; + let second_slice = &hash_hex[2..4]; + storage_dir + .join(first_slice) + .join(second_slice) + .join(hash_hex) +} + +/// 根据hash字节数组计算chunk文件的存储路径 +/// +/// # 参数 +/// - `storage_dir`: 存储目录 +/// - `hash_bytes`: chunk的hash值（字节数组） +/// +/// # 返回 +/// 返回chunk文件的完整路径 +pub fn get_chunk_path_from_bytes(storage_dir: &Path, hash_bytes: &[u8; 32]) -> PathBuf { + let hash_hex = hex::encode(hash_bytes); + get_chunk_path(storage_dir, &hash_hex) +} + +/// 生成唯一的文件路径，如果文件已存在则添加数字后缀 +/// +/// # 参数 +/// - `output_dir`: 输出目录 +/// - `desired_filename`: 期望的文件名 +/// +/// # 返回 +/// 返回唯一的文件路径 +pub fn generate_unique_path(output_dir: &Path, desired_filename: &str) -> PathBuf { + let desired_path = output_dir.join(desired_filename); + let mut candidate = desired_path.clone(); + + let mut counter = 1; + while candidate.exists() { + let path_buf = PathBuf::from(desired_filename); + if let Some(stem) = path_buf.file_stem() { + if let Some(ext) = path_buf.extension() { + let ext_str = ext.to_string_lossy(); + let new_name = if ext_str.is_empty() { + format!("{}_{}", stem.to_string_lossy(), counter) + } else { + format!("{}.{}_{}", stem.to_string_lossy(), ext_str, counter) + }; + candidate = output_dir.join(new_name); + } else { + candidate = output_dir.join(format!("{}_{}", stem.to_string_lossy(), counter)); + } + } else { + candidate = output_dir.join(format!("{}_{}", desired_filename, counter)); + } + counter += 1; + } + + candidate +} diff --git a/src/chunker/rw/storage/build.rs b/src/chunker/rw/storage/build.rs new file mode 100644 index 0000000..7608b5c --- /dev/null +++ b/src/chunker/rw/storage/build.rs @@ -0,0 +1,250 @@ +use futures::future::join_all; +use just_progress::progress; +use log::{error, info, trace}; +use memmap2::Mmap; +use std::path::PathBuf; +use tokio::{fs::File, io::AsyncWriteExt}; + +use crate::{ + chunker::{ + constants::{BUTCK_INDEX_FILE_SUFFIX, BUTCK_INDEX_MAGIC}, + context::ButckContext, + rw::error::{ButckRWError, ButckRWErrorKind}, + rw::storage, + }, + utils::size_display::size_display, +}; + +pub async fn build(ctx: ButckContext) -> Result<(), ButckRWError> { + if ctx.storage_path.is_none() { + return Err(ButckRWErrorKind::NoButckStorageFound.pack(ctx)); + } + if ctx.file_paths.is_empty() { + return Err( + ButckRWErrorKind::RebuildFailed("No bidx files specified".to_string()).pack(ctx), + ); + } + + let tasks: Vec<_> = ctx + .file_paths + .iter() + .map(|bidx_path| async { + trace!( + "Preparing to rebuild from bidx file `{}`", + bidx_path.display() + ); + rebuild_from_bidx(bidx_path, &ctx).await + }) + .collect(); + + let results = join_all(tasks).await; + + for result in results { + if let Err(e) = result { + return Err(e.pack(ctx)); + } + } + + Ok(()) +} + +async fn rebuild_from_bidx( + bidx_path: &PathBuf, + ctx: &ButckContext, +) -> Result<(), ButckRWErrorKind> { + // Validate file extension + if let Some(ext) = bidx_path.extension() + && ext != BUTCK_INDEX_FILE_SUFFIX + { + return Err(ButckRWErrorKind::InvalidBidxFormat); + } + + info!("Rebuilding from bidx file: {}", bidx_path.display()); + + // Read bidx file content + let bidx_content = if ctx.memmap_read { + let file = File::open(bidx_path).await?; + let mmap = unsafe { Mmap::map(&file)? }; + mmap.to_vec() + } else { + tokio::fs::read(bidx_path).await? + }; + + // Verify file size includes at least the header + if bidx_content.len() < 6 { + return Err(ButckRWErrorKind::InvalidBidxFormat); + } + + // Validate MAGIC bytes + if bidx_content[0..4] != BUTCK_INDEX_MAGIC { + return Err(ButckRWErrorKind::InvalidBidxFormat); + } + + // Read filename + let filename_len = u16::from_le_bytes([bidx_content[4], bidx_content[5]]) as usize; + if bidx_content.len() < 6 + filename_len { + return Err(ButckRWErrorKind::InvalidBidxFormat); + } + let filename_bytes = &bidx_content[6..6 + filename_len]; + let original_filename = String::from_utf8(filename_bytes.to_vec()) + .map_err(|_| ButckRWErrorKind::InvalidBidxFormat)?; + + trace!("Original filename from bidx: {}", original_filename); + + let hash_data_start = 6 + filename_len; + let hash_data = &bidx_content[hash_data_start..]; + + // Verify that hash data size is a multiple of 32 bytes + if hash_data.len() % 32 != 0 { + return Err(ButckRWErrorKind::InvalidBidxFormat); + } + + let chunk_count = hash_data.len() / 32; + info!("Found {} chunks in bidx file", chunk_count); + + let mut chunk_hashes = Vec::with_capacity(chunk_count); + for i in 0..chunk_count { + let start = i * 32; + let end = start + 32; + let hash_bytes: [u8; 32] = hash_data[start..end] + .try_into() + .map_err(|_| ButckRWErrorKind::InvalidBidxFormat)?; + chunk_hashes.push(hash_bytes); + } + + trace!("Parsed {} chunk hashes", chunk_hashes.len()); + + // Determine output file path + let output_path = if let Some(output_file) = &ctx.output_file { + output_file.clone() + } else { + // Use the original filename read from the bidx file + storage::generate_unique_path(&ctx.output_dir, &original_filename) + }; + + info!("Rebuilding file to: {}", output_path.display()); + + let progress_name = format!("Rebuild `{}`", output_path.display()); + progress::update_progress(progress_name.as_str(), 0.0); + let step = 1.0 / chunk_count as f64; + + let mut tasks = Vec::with_capacity(chunk_count); + + for (index, hash_bytes) in chunk_hashes.iter().enumerate() { + let hash_hex = hex::encode(hash_bytes); + tasks.push(read_chunk( + progress_name.as_str(), + step, + hash_hex, + &ctx.output_dir, + index, + )); + } + + trace!("Starting parallel read of {} chunks", tasks.len()); + let results = join_all(tasks).await; + trace!("All read tasks completed"); + + // Collect chunk data and verify order + let mut chunk_data_list = Vec::with_capacity(chunk_count); + let mut success_count = 0; + + for (index, result) in results.into_iter().enumerate() { + match result { + Ok(chunk_data) => { + let chunk_size = chunk_data.len(); + success_count += 1; + chunk_data_list.push((index, chunk_data)); + trace!( + "Chunk {} read successfully, size: {} bytes", + index, chunk_size + ); + } + Err(e) => { + error!("Failed to read chunk {}: {:?}", index, e); + return Err(e); + } + } + } + + if success_count != chunk_count { + return Err(ButckRWErrorKind::ChunkNotFound(format!( + "Only {}/{} chunks found in storage", + success_count, chunk_count + ))); + } + + info!("All {} chunks read successfully", success_count); + + // Sort by index and concatenate files + chunk_data_list.sort_by_key(|(index, _)| *index); + + // Calculate total size + let total_size: usize = chunk_data_list.iter().map(|(_, data)| data.len()).sum(); + let (total_value, total_unit) = size_display(total_size); + info!( + "Rebuilding file: {} chunks, total size: {:.2} {} ({} bytes)", + chunk_count, total_value, total_unit, total_size + ); + + // Write to output file + trace!("Writing to output file: {}", output_path.display()); + let mut output_file = File::create(&output_path).await?; + + for (index, chunk_data) in chunk_data_list { + trace!("Writing chunk {} ({} bytes)", index, chunk_data.len()); + output_file.write_all(&chunk_data).await?; + progress::increase(progress_name.as_str(), step as f32); + } + + output_file.flush().await?; + + info!("File successfully rebuilt: {}", output_path.display()); + progress::complete(progress_name.as_str()); + + Ok(()) +} + +/// Read a single chunk from storage +async fn read_chunk( + progress_name: &str, + step: f64, + hash_hex: String, + storage_dir: &PathBuf, + chunk_index: usize, +) -> Result, ButckRWErrorKind> { + trace!("read_chunk[{}]: Starting, hash: {}", chunk_index, hash_hex); + + // Build chunk file path + let file_path = storage::get_chunk_path(storage_dir, &hash_hex); + + trace!( + "read_chunk[{}]: Looking for file at: {}", + chunk_index, + file_path.display() + ); + + // Read chunk file + match tokio::fs::read(&file_path).await { + Ok(data) => { + trace!( + "read_chunk[{}]: Read {} bytes successfully", + chunk_index, + data.len() + ); + progress::increase(progress_name, step as f32); + Ok(data) + } + Err(e) if e.kind() == std::io::ErrorKind::NotFound => { + trace!("read_chunk[{}]: File not found", chunk_index); + Err(ButckRWErrorKind::ChunkNotFound(format!( + "Chunk {} (hash: {}) not found in storage", + chunk_index, hash_hex + ))) + } + Err(e) => { + trace!("read_chunk[{}]: Read failed: {:?}", chunk_index, e); + Err(ButckRWErrorKind::IOError(e)) + } + } +} diff --git a/src/chunker/rw/storage/write.rs b/src/chunker/rw/storage/write.rs new file mode 100644 index 0000000..8b3acc7 --- /dev/null +++ b/src/chunker/rw/storage/write.rs @@ -0,0 +1,118 @@ +use std::{collections::HashMap, path::PathBuf}; + +use log::trace; + +use crate::{ + chunker::{ + constants::BUTCK_INDEX_FILE_SUFFIX, + context::ButckContext, + rw::{ + error::{ButckRWError, ButckRWErrorKind}, + storage::generate_unique_path, + }, + }, + storage::{simple::write_file_simple, stream::write_file_stream}, +}; + +pub mod simple; +pub mod stream; + +pub async fn write(ctx: ButckContext) -> Result<(), ButckRWError> { + if ctx.storage_path.is_none() { + return Err(ButckRWErrorKind::NoButckStorageFound.pack(ctx)); + } + if ctx.policy_name.is_none() { + return Err(ButckRWErrorKind::ChunkingPolicyNotSpecified.pack(ctx)); + } + if ctx.file_paths.len() > 1 && ctx.output_file.is_some() { + return Err(ButckRWErrorKind::OutputCountMismatch.pack(ctx)); + } + + // Cannot enable both memory-mapped and stream reading simultaneously. + // Stream reading uses butck_policies::chunk_stream_with, + // while memory-mapped or default reading uses butck_policies::chunk_with. + if ctx.memmap_read && ctx.stream_read.is_some() { + return Err(ButckRWErrorKind::ReadingMethodAmbiguous.pack(ctx)); + } + + let param_refs: HashMap<&str, &str> = ctx + .params + .iter() + .map(|(k, v)| (k.as_str(), v.as_str())) + .collect(); + + let tasks: Vec<_> = ctx + .file_paths + .iter() + .map(|path| async { + trace!("Preparing to write file `{}`", path.display()); + write_file(path, &ctx, ¶m_refs).await + }) + .collect(); + + let results = futures::future::join_all(tasks).await; + + for result in results { + if let Err(e) = result { + return Err(e.pack(ctx)); + } + } + + Ok(()) +} + +async fn write_file( + path: &PathBuf, + ctx: &ButckContext, + params: &HashMap<&str, &str>, +) -> Result<(), ButckRWErrorKind> { + if let Some(stream_read_size) = ctx.stream_read { + write_file_stream(path, stream_read_size, ctx, params).await + } else { + write_file_simple(path, ctx, params).await + } +} + +pub fn get_index_file_name(path: &PathBuf, ctx: &ButckContext) -> PathBuf { + let output_file = if let Some(output_file) = &ctx.output_file { + return output_file.clone(); + } else { + ctx.output_dir.join(path.file_name().unwrap_or_default()) + }; + + // Append .bidx suffix directly to the original file name + let desired_filename = if let Some(ext) = output_file.extension() { + let ext_str = ext.to_string_lossy(); + if ext_str.is_empty() { + format!( + "{}.{}", + output_file + .file_stem() + .unwrap_or_default() + .to_string_lossy(), + BUTCK_INDEX_FILE_SUFFIX + ) + } else { + format!( + "{}.{}.{}", + output_file + .file_stem() + .unwrap_or_default() + .to_string_lossy(), + ext_str, + BUTCK_INDEX_FILE_SUFFIX + ) + } + } else { + format!( + "{}.{}", + output_file + .file_name() + .unwrap_or_default() + .to_string_lossy(), + BUTCK_INDEX_FILE_SUFFIX + ) + }; + + generate_unique_path(&ctx.output_dir, &desired_filename) +} diff --git a/src/chunker/rw/storage/write/simple.rs b/src/chunker/rw/storage/write/simple.rs new file mode 100644 index 0000000..75b9bd7 --- /dev/null +++ b/src/chunker/rw/storage/write/simple.rs @@ -0,0 +1,368 @@ +use futures::future::join_all; +use just_progress::progress; +use log::{error, info, trace}; +use std::{collections::HashMap, path::PathBuf}; +use tokio::{fs::File, io::AsyncReadExt}; + +use crate::{ + chunker::{ + context::ButckContext, + rw::{error::ButckRWErrorKind, storage}, + }, + core::hash::ChunkWriteHash, + storage::get_index_file_name, + utils::size_display::size_display, +}; + +pub async fn write_file_simple( + path: &PathBuf, + ctx: &ButckContext, + params: &HashMap<&str, &str>, +) -> Result<(), ButckRWErrorKind> { + read_file(path, ctx, params).await?; + Ok(()) +} + +async fn read_file( + path: &PathBuf, + ctx: &ButckContext, + params: &HashMap<&str, &str>, +) -> Result<(), ButckRWErrorKind> { + let mut file = File::open(path).await?; + + // Use butck_policies::chunk_with to locate chunk boundaries in the file + if ctx.memmap_read { + let mmap = unsafe { memmap2::Mmap::map(&file)? }; + let raw_data = &mmap[..]; + let (chunk_boundaries, total_bytes) = + (get_boundaries(raw_data, ctx, params).await?, raw_data.len()); + + // If output boundaries, do not execute actual write logic + if ctx.display_boundaries { + display_boundaries(&chunk_boundaries, total_bytes).await; + return Ok(()); + } else { + write_file_to_storage(path, ctx, chunk_boundaries, raw_data).await?; + } + } else { + let mut contents = Vec::new(); + file.read_to_end(&mut contents).await?; + let raw_data = &contents[..]; + let (chunk_boundaries, total_bytes) = + (get_boundaries(raw_data, ctx, params).await?, raw_data.len()); + + // If output boundaries, do not execute actual write logic + if ctx.display_boundaries { + display_boundaries(&chunk_boundaries, total_bytes).await; + return Ok(()); + } else { + write_file_to_storage(path, ctx, chunk_boundaries, raw_data).await?; + } + }; + progress::clear_all(); + Ok(()) +} + +async fn write_file_to_storage( + path: &PathBuf, + ctx: &ButckContext, + chunk_boundaries: Vec, + raw_data: &[u8], +) -> Result<(), ButckRWErrorKind> { + let output_index_file = get_index_file_name(path, ctx); + + let chunk_count = chunk_boundaries.len() + 1; + let progress_name = format!("Write `{}`", path.display()); + + progress::update_progress(progress_name.as_str(), 0.0); + let step = 1.0 / chunk_count as f64; + + trace!("chunks_count={}", chunk_count); + trace!("chunk_hash={:?}", ctx.chunk_hash); + trace!("file_size={}", raw_data.len()); + trace!("output_index_file={}", output_index_file.display()); + trace!("policy_name={:?}", ctx.policy_name); + trace!("storage_dir={}", ctx.output_dir.display()); + + info!( + "{} chunks will be written to {}", + chunk_count, + ctx.output_dir.display() + ); + + tokio::fs::create_dir_all(&ctx.output_dir).await?; + trace!("Output directory created or already exists"); + + let mut tasks = Vec::new(); + let mut start = 0; + let mut chunk_index = 0; + + trace!("Processing chunk boundaries:"); + + for &boundary in &chunk_boundaries { + let end = boundary as usize; + if start < end && end <= raw_data.len() { + let chunk_data = &raw_data[start..end]; + trace!( + "Chunk {}: bytes {}..{} (size: {} bytes)", + chunk_index, + start, + end - 1, + end - start + ); + tasks.push(write_chunk( + progress_name.as_str(), + step, + chunk_data, + &ctx.output_dir, + &ctx.chunk_hash, + chunk_index, + start, + end, + )); + chunk_index += 1; + } else { + trace!( + "Skipping invalid chunk boundary: start={}, end={}, data_len={}", + start, + end, + raw_data.len() + ); + } + start = end; + } + + if start < raw_data.len() { + let chunk_data = &raw_data[start..]; + trace!( + "Chunk {}: bytes {}..{} (size: {} bytes) - final chunk", + chunk_index, + start, + raw_data.len() - 1, + raw_data.len() - start + ); + tasks.push(write_chunk( + progress_name.as_str(), + step, + chunk_data, + &ctx.output_dir, + &ctx.chunk_hash, + chunk_index, + start, + raw_data.len(), + )); + } + + trace!("Total chunks prepared for writing: {}", tasks.len()); + + trace!("Starting parallel write of {} chunks", tasks.len()); + let results = join_all(tasks).await; + trace!("All write tasks completed"); + + let mut success_count = 0; + let mut chunk_infos = Vec::new(); + + for result in results { + match result { + Ok(chunk_info) => { + success_count += 1; + chunk_infos.push(chunk_info); + } + Err(e) => { + trace!("Chunk write failed: {:?}", e); + return Err(e); + } + } + } + + info!("All {} chunks written successfully", success_count); + + // Write index file + trace!("Writing index file to: {}", output_index_file.display()); + if let Err(e) = write_index_file(&output_index_file, &chunk_infos, path).await { + error!("Failed to write index file: {}", e); + return Err(ButckRWErrorKind::IOError(e)); + } + info!("Index file written to: {}", output_index_file.display()); + + trace!("write_file_to_storage completed successfully"); + + progress::complete(progress_name.as_str()); + + Ok(()) +} + +async fn write_chunk( + progress_name: &str, + step: f64, + chunk_data: &[u8], + output_dir: &PathBuf, + chunk_hash: &ChunkWriteHash, + chunk_index: usize, + start: usize, + end: usize, +) -> Result { + trace!( + "write_chunk[{}]: Starting, data size: {} bytes", + chunk_index, + chunk_data.len() + ); + + trace!( + "write_chunk[{}]: Computing hash with algorithm: {:?}", + chunk_index, chunk_hash + ); + let hash_bytes = chunk_hash.hash(chunk_data); + trace!( + "write_chunk[{}]: Hash computed: {:?}", + chunk_index, hash_bytes + ); + + let hash_hex = hex::encode(hash_bytes); + trace!("write_chunk[{}]: Hash hex: {}", chunk_index, hash_hex); + + let file_path = storage::get_chunk_path(output_dir, &hash_hex); + + if let Some(parent_dir) = file_path.parent() { + trace!( + "write_chunk[{}]: Creating directory structure: {}", + chunk_index, + parent_dir.display() + ); + tokio::fs::create_dir_all(parent_dir).await?; + trace!("write_chunk[{}]: Directory created", chunk_index); + } + + trace!( + "write_chunk[{}]: File path: {}", + chunk_index, + file_path.display() + ); + + trace!( + "write_chunk[{}]: Writing {} bytes to file", + chunk_index, + chunk_data.len() + ); + if !file_path.exists() { + tokio::fs::write(&file_path, chunk_data).await?; + } else { + trace!( + "write_chunk[{}]: File already exists, skipping", + chunk_index + ); + } + trace!("write_chunk[{}]: File written successfully", chunk_index); + progress::increase(progress_name, step as f32); + Ok(crate::chunker::rw::storage::ChunkInfo { + index: chunk_index, + hash: hash_hex, + size: chunk_data.len(), + start, + end, + }) +} + +async fn get_boundaries<'a>( + raw_data: &[u8], + ctx: &ButckContext, + params: &HashMap<&str, &str>, +) -> Result, ButckRWErrorKind> { + let policy_name = ctx.policy_name.as_ref().unwrap().as_str(); + match butck_policies::chunk_with(policy_name, raw_data, params).await { + Ok(s) => Ok(s), + Err(e) => Err(ButckRWErrorKind::ChunkFailed(e)), + } +} + +async fn write_index_file( + index_path: &PathBuf, + chunk_infos: &[crate::chunker::rw::storage::ChunkInfo], + original_file_path: &PathBuf, +) -> Result<(), std::io::Error> { + use std::io::Write; + + let file = std::fs::File::create(index_path)?; + let mut writer = std::io::BufWriter::new(file); + + // Write header: [u8; 4] magic + [u16] filename length + [u8] filename bytes + use crate::chunker::constants::BUTCK_INDEX_MAGIC; + + // Write magic bytes + writer.write_all(&BUTCK_INDEX_MAGIC)?; + + // Get original filename as bytes + let filename = original_file_path + .file_name() + .and_then(|n| n.to_str()) + .unwrap_or("unknown"); + let filename_bytes = filename.as_bytes(); + + // Write filename length as u16 (little-endian) + if filename_bytes.len() > u16::MAX as usize { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("Filename too long: {} bytes", filename_bytes.len()), + )); + } + let filename_len = filename_bytes.len() as u16; + writer.write_all(&filename_len.to_le_bytes())?; + + // Write filename bytes + writer.write_all(filename_bytes)?; + + // Write chunk hashes: [u8; 32][u8; 32][u8; 32]... + for chunk_info in chunk_infos { + // Convert hex hash to bytes + match hex::decode(&chunk_info.hash) { + Ok(hash_bytes) => { + if hash_bytes.len() == 32 { + writer.write_all(&hash_bytes)?; + } else { + // Pad or truncate to 32 bytes if needed + let mut fixed_hash = [0u8; 32]; + let len = hash_bytes.len().min(32); + fixed_hash[..len].copy_from_slice(&hash_bytes[..len]); + writer.write_all(&fixed_hash)?; + } + } + Err(e) => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Failed to decode hash hex: {}", e), + )); + } + } + } + + Ok(()) +} + +async fn display_boundaries(chunk_boundaries: &Vec, total_bytes: usize) { + let total_chunks = chunk_boundaries.len() + 1; + let (total_value, total_unit) = size_display(total_bytes); + info!( + "{} chunks, ({:.2} {}, {})", + total_chunks, total_value, total_unit, total_bytes + ); + let mut start = 0; + chunk_boundaries.iter().for_each(|p| { + let next = *p as usize; + let (size_value, size_unit) = size_display(next - start); + info!( + "{} - {} (size: {:.2} {})", + start, + next - 1, + size_value, + size_unit + ); + start = next; + }); + let last = start; + let r#final = total_bytes; + let (size_value, size_unit) = size_display(total_bytes - start); + info!( + "{} - {} (size: {:.2} {})", + last, r#final, size_value, size_unit + ); +} diff --git a/src/chunker/rw/storage/write/stream.rs b/src/chunker/rw/storage/write/stream.rs new file mode 100644 index 0000000..020cfcd --- /dev/null +++ b/src/chunker/rw/storage/write/stream.rs @@ -0,0 +1,12 @@ +use std::{collections::HashMap, path::PathBuf}; + +use crate::chunker::{context::ButckContext, rw::error::ButckRWErrorKind}; + +pub async fn write_file_stream( + path: &PathBuf, + stream_read_size: u32, + ctx: &ButckContext, + params: &HashMap<&str, &str>, +) -> Result<(), ButckRWErrorKind> { + todo!() +} diff --git a/src/core.rs b/src/core.rs new file mode 100644 index 0000000..ec5d33c --- /dev/null +++ b/src/core.rs @@ -0,0 +1 @@ +pub mod hash; diff --git a/src/core/hash.rs b/src/core/hash.rs new file mode 100644 index 0000000..36a62b3 --- /dev/null +++ b/src/core/hash.rs @@ -0,0 +1,38 @@ +use blake3::Hasher as Blake3Hasher; +use sha2::{Digest as Sha2Digest, Sha256}; + +const SALT: &[u8] = b"Dude@"; + +#[derive(Debug, Default)] +pub enum ChunkWriteHash { + #[default] + Blake3, + Sha256, +} + +impl ChunkWriteHash { + pub fn hash(&self, d: &[u8]) -> [u8; 32] { + match self { + ChunkWriteHash::Blake3 => hash_blake3(d), + ChunkWriteHash::Sha256 => hash_sha256(d), + } + } +} + +/// Compute the Blake3 hash of the data with a salt +/// Returns a 32-byte hash value +pub fn hash_blake3(d: &[u8]) -> [u8; 32] { + let mut hasher = Blake3Hasher::new(); + hasher.update(SALT); + hasher.update(d); + *hasher.finalize().as_bytes() +} + +/// Compute the SHA-256 hash of the data with a salt +/// Returns a 32-byte hash value +pub fn hash_sha256(d: &[u8]) -> [u8; 32] { + let mut hasher = Sha256::new(); + hasher.update(SALT); + hasher.update(d); + hasher.finalize().into() +} diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..e4a55c2 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,11 @@ +pub mod chunker; +pub mod core; +pub mod log; +pub mod macros; +pub mod utils; + +pub mod storage { + pub use crate::chunker::rw::error::*; + pub use crate::chunker::rw::storage::build::*; + pub use crate::chunker::rw::storage::write::*; +} diff --git a/src/log.rs b/src/log.rs new file mode 100644 index 0000000..5fc6160 --- /dev/null +++ b/src/log.rs @@ -0,0 +1,33 @@ +use env_logger::Builder; +use log::Level; +use std::io::Write; + +pub fn init_logger(level_filter: Option) { + let mut builder = match level_filter { + Some(f) => { + let mut b = Builder::new(); + b.filter_level(f); + b + } + None => return, + }; + + builder + .format(|buf, record| { + let level = record.level(); + let args = record.args(); + + let (prefix, color_code) = match level { + Level::Error => ("error: ", "\x1b[1;31m"), + Level::Warn => ("warn: ", "\x1b[1;33m"), + Level::Info => ("", "\x1b[37m"), + Level::Debug => ("debug: ", "\x1b[90m"), + Level::Trace => ("trace: ", "\x1b[36m"), + }; + + let colored_prefix = format!("{}{}\x1b[0m", color_code, prefix); + + writeln!(buf, "{}{}", colored_prefix, args) + }) + .init(); +} diff --git a/src/macros.rs b/src/macros.rs new file mode 100644 index 0000000..11b8da4 --- /dev/null +++ b/src/macros.rs @@ -0,0 +1,47 @@ +#[macro_export] +macro_rules! special_flag { + ($args:expr, $($flag:expr),+) => {{ + let mut found = false; + $( + let flag = $flag; + if $args.iter().any(|arg| arg == flag) { + found = true; + } + $args.retain(|arg| arg != flag); + )+ + found + }}; +} + +#[macro_export] +macro_rules! special_argument { + ($args:expr, $($flag:expr),+) => {{ + let mut value: Option = None; + let mut found = false; + $( + let flag = $flag; + if !found { + let mut i = 0; + while i < $args.len() { + if $args[i] == flag { + if i + 1 < $args.len() { + value = Some($args[i + 1].clone()); + $args.remove(i + 1); + $args.remove(i); + } else { + value = None; + $args.remove(i); + } + #[allow(unused_assignments)] + { + found = true; + } + break; + } + i += 1; + } + } + )+ + value + }}; +} diff --git a/src/utils.rs b/src/utils.rs new file mode 100644 index 0000000..b64c0c4 --- /dev/null +++ b/src/utils.rs @@ -0,0 +1,2 @@ +pub mod file_input_solve; +pub mod size_display; diff --git a/src/utils/file_input_solve.rs b/src/utils/file_input_solve.rs new file mode 100644 index 0000000..30d5765 --- /dev/null +++ b/src/utils/file_input_solve.rs @@ -0,0 +1,82 @@ +use std::{ + env::current_dir, + path::{Path, PathBuf}, +}; + +use just_fmt::fmt_path::fmt_path; + +pub fn parse_path_input( + files: Vec, + recursive: bool, + exclude_dir: Vec<&str>, +) -> Vec { + let current_dir = current_dir().unwrap(); + let files = if recursive { + let mut result: Vec = Vec::new(); + for arg in files.iter().skip(1) { + if exclude_dir.contains(&arg.as_str()) { + continue; + } + let path = current_dir.join(arg); + if path.is_dir() { + if let Err(e) = collect_files_recursively(&path, &mut result) { + eprintln!("Error collecting files recursively: {}", e); + continue; + } + } else { + result.push(path); + } + } + result + } else { + let mut result = Vec::new(); + for arg in files.iter().skip(1) { + if exclude_dir.contains(&arg.as_str()) { + continue; + } + let path = current_dir.join(arg); + if path.is_dir() { + if files.len() == 2 { + for entry in std::fs::read_dir(&path) + .unwrap_or_else(|e| { + eprintln!("Error reading directory: {}", e); + std::fs::read_dir(".").unwrap() + }) + .flatten() + { + let entry_path = entry.path(); + if !entry_path.is_dir() { + result.push(entry_path); + } + } + } + } else { + result.push(path); + } + } + result + }; + files + .into_iter() + .filter_map(|path| match fmt_path(path) { + Ok(formatted_path) => Some(formatted_path), + Err(e) => { + eprintln!("Error formatting path: {}", e); + None + } + }) + .collect() +} + +fn collect_files_recursively(dir: &Path, files: &mut Vec) -> std::io::Result<()> { + for entry in std::fs::read_dir(dir)? { + let entry = entry?; + let path = entry.path(); + if path.is_dir() { + collect_files_recursively(&path, files)?; + } else { + files.push(path); + } + } + Ok(()) +} diff --git a/src/utils/size_display.rs b/src/utils/size_display.rs new file mode 100644 index 0000000..3e2bc29 --- /dev/null +++ b/src/utils/size_display.rs @@ -0,0 +1,14 @@ +pub fn size_display<'a>(total_bytes: usize) -> (f64, &'a str) { + let total_bytes = total_bytes as f64; + if total_bytes >= 1024.0 * 1024.0 * 1024.0 * 1024.0 { + (total_bytes / (1024.0 * 1024.0 * 1024.0 * 1024.0), "TB") + } else if total_bytes >= 1024.0 * 1024.0 * 1024.0 { + (total_bytes / (1024.0 * 1024.0 * 1024.0), "GB") + } else if total_bytes >= 1024.0 * 1024.0 { + (total_bytes / (1024.0 * 1024.0), "MB") + } else if total_bytes >= 1024.0 { + (total_bytes / 1024.0, "KB") + } else { + (total_bytes, "B") + } +} -- cgit