Skip to content

Commit

Permalink
refactor info_cmd to support same http options as clone_cmd
Browse files Browse the repository at this point in the history
  • Loading branch information
@caesay
caesay committed Oct 30, 2024
1 parent 603b230 commit 495e30f
Show file tree
Hide file tree
Showing 5 changed files with 142 additions and 117 deletions.
2 changes: 1 addition & 1 deletion bitar/src/api/compress.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ pub struct CreateArchiveOptions {

/// The type of compression to use when compressing a chunk
pub compression: Option<Compression>,

/// Custom string/bytes key-value pair metadata to be stored in the archive header
pub metadata: HashMap<String, Vec<u8>>,
}
Expand Down
20 changes: 12 additions & 8 deletions bitar/src/archive.rs
View file
Original file line number Diff line number Diff line change
@@ -1,16 +1,16 @@
use crate::{
archive_reader::ArchiveReader, chunk_dictionary as dict, chunker,
compression::CompressionAlgorithm, header, ChunkIndex, ChunkOffset, CompressedArchiveChunk,
CompressedChunk, Compression, HashSum,
};
use blake2::{Blake2b512, Digest};
use futures_util::{stream::Stream, StreamExt};
use std::collections::HashMap;
use std::{
convert::TryInto,
fmt,
task::{ready, Poll},
};
use std::collections::HashMap;
use crate::{
archive_reader::ArchiveReader, chunk_dictionary as dict, chunker,
compression::CompressionAlgorithm, header, ChunkIndex, ChunkOffset, CompressedArchiveChunk,
CompressedChunk, Compression, HashSum,
};

#[derive(Debug)]
pub enum ArchiveError<R> {
Expand Down Expand Up @@ -250,9 +250,13 @@ impl<R> Archive<R> {
&self.created_by_app_version
}
/// Get the custom key-value pair metadata stored in the archive header.
pub fn metadata(&self) -> &HashMap<String, Vec<u8>> { &self.metadata }
pub fn metadata(&self) -> &HashMap<String, Vec<u8>> {
&self.metadata
}
/// Get a specific metadata value stored in the archive header, or None if it is not present.
pub fn metadata_value(&self, key: &str) -> Option<&Vec<u8>> { self.metadata.get(key) }
pub fn metadata_value(&self, key: &str) -> Option<&Vec<u8>> {
self.metadata.get(key)
}
/// Iterate chunks as ordered in source.
pub fn iter_source_chunks(&self) -> impl Iterator<Item = (u64, &ChunkDescriptor)> {
let mut chunk_offset = 0;
Expand Down
188 changes: 98 additions & 90 deletions src/cli.rs
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
use std::collections::HashMap;
use clap::error::ErrorKind;
use clap::{value_parser, Arg, ArgAction, ArgMatches, Command};
use log::LevelFilter;
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
use std::collections::HashMap;
use std::ffi::OsString;
use std::path::{Path, PathBuf};
use std::time::Duration;
Expand Down Expand Up @@ -63,7 +63,7 @@ where
.long("metadata-file")
.num_args(2) // Expect exactly 2 values (key and path) each time
.action(clap::ArgAction::Append) // Append to the list of values
.value_names(&["KEY", "PATH"])
.value_names(&["KEY", "PATH"])
.help("Custom metadata key-value pair where the value is a file contents"),
)
.arg(
Expand All @@ -76,69 +76,41 @@ where
),
);

let clone_subcmd = Command::new("clone")
.about("Clone a remote (or local archive). The archive is unpacked while being cloned.")
.arg(input_archive_arg())
.arg(
Arg::new("http-retry-count")
.long("http-retry-count")
.value_name("COUNT")
.value_parser(value_parser!(u32))
.default_value("0")
.help("Retry transfer on failure"),
)
.arg(
Arg::new("http-retry-delay")
.long("http-retry-delay")
.value_name("SECONDS")
.value_parser(value_parser!(u64))
.default_value("0")
.help("Delay retry for some time on transfer failure"),
)
.arg(
Arg::new("http-timeout")
.long("http-timeout")
.value_name("SECONDS")
.value_parser(value_parser!(u64))
.help("Fail transfer if unresponsive for some time"),
)
.arg(
Arg::new("http-header")
.long("http-header")
.value_name("HEADER")
.action(ArgAction::Append)
.help("Provide custom http header(s)"),
)
.arg(
Arg::new("verify-header")
.long("verify-header")
.value_name("CHECKSUM")
.value_parser(parse_hash_sum)
.help("Verify that the archive header checksum is the one given"),
)
.arg(output_file_arg())
.arg(
Arg::new("seed")
.value_name("FILE")
.value_parser(value_parser!(OsString))
.action(ArgAction::Append)
.long("seed")
.help("File to use as seed while cloning or '-' to read from stdin"),
)
.arg(
Arg::new("seed-output")
.long("seed-output")
.action(ArgAction::SetTrue)
.help("Use the output file as seed and update in-place"),
)
.arg(force_create_arg())
.arg(
Arg::new("verify-output")
.long("verify-output")
.action(ArgAction::SetTrue)
.help("Verify that the checksum of the output matches with the archive"),
)
.arg(buffered_chunks_arg());
let clone_subcmd = add_archive_input_http_args(
Command::new("clone")
.about("Clone a remote (or local archive). The archive is unpacked while being cloned.")
.arg(input_archive_arg())
.arg(
Arg::new("verify-header")
.long("verify-header")
.value_name("CHECKSUM")
.value_parser(parse_hash_sum)
.help("Verify that the archive header checksum is the one given"),
)
.arg(output_file_arg())
.arg(
Arg::new("seed")
.value_name("FILE")
.value_parser(value_parser!(OsString))
.action(ArgAction::Append)
.long("seed")
.help("File to use as seed while cloning or '-' to read from stdin"),
)
.arg(
Arg::new("seed-output")
.long("seed-output")
.action(ArgAction::SetTrue)
.help("Use the output file as seed and update in-place"),
)
.arg(force_create_arg())
.arg(
Arg::new("verify-output")
.long("verify-output")
.action(ArgAction::SetTrue)
.help("Verify that the checksum of the output matches with the archive"),
)
.arg(buffered_chunks_arg()),
);

let diff_subcmd = add_chunker_args(
Command::new("diff")
Expand All @@ -159,16 +131,18 @@ where
)
.arg(buffered_chunks_arg()),
);

let info_subcmd = Command::new("info")
.about("Print archive details")
.arg(input_archive_arg())
.arg(
Arg::new("metadata-key")
.long("metadata-key")
.value_name("KEY")
.help("Print only the metadata value for the given key"),
);

let info_subcmd = add_archive_input_http_args(
Command::new("info")
.about("Print archive details")
.arg(
Arg::new("metadata-key")
.long("metadata-key")
.value_name("KEY")
.help("Print only the metadata value for the given key"),
)
.arg(input_archive_arg()),
);

let mut cmd = Command::new(PKG_NAME)
.version(PKG_VERSION)
Expand Down Expand Up @@ -218,7 +192,7 @@ where
}
}
}

let mut metadata_strings: HashMap<String, String> = HashMap::new();
if let Some(values) = matches.get_many::<String>("metadata-value") {
let values: Vec<_> = values.collect();
Expand All @@ -228,7 +202,7 @@ where
}
}
}

Ok((
CommandOpts::Compress(compress_cmd::Options {
input: input.cloned(),
Expand Down Expand Up @@ -278,11 +252,11 @@ where
log_opts,
))
} else if let Some(matches) = matches.subcommand_matches("info") {
let input = matches.get_one::<OsString>("ARCHIVE").unwrap();
let metadata_key = matches.get_one::<String>("metadata-key");
let input_archive = parse_input_archive_config(&mut cmd, matches)?;
Ok((
CommandOpts::Info(info_cmd::Options {
input: input.clone(),
input_archive,
metadata_key: metadata_key.cloned(),
}),
log_opts,
Expand Down Expand Up @@ -443,6 +417,39 @@ fn parse_hash_sum(hex_str: &str) -> Result<HashSum, std::num::ParseIntError> {
hex_str_to_vec(hex_str).map(HashSum::from)
}

fn add_archive_input_http_args(cmd: Command) -> Command {
cmd.arg(
Arg::new("http-retry-count")
.long("http-retry-count")
.value_name("COUNT")
.value_parser(value_parser!(u32))
.default_value("0")
.help("Retry transfer on failure"),
)
.arg(
Arg::new("http-retry-delay")
.long("http-retry-delay")
.value_name("SECONDS")
.value_parser(value_parser!(u64))
.default_value("0")
.help("Delay retry for some time on transfer failure"),
)
.arg(
Arg::new("http-timeout")
.long("http-timeout")
.value_name("SECONDS")
.value_parser(value_parser!(u64))
.help("Fail transfer if unresponsive for some time"),
)
.arg(
Arg::new("http-header")
.long("http-header")
.value_name("HEADER")
.action(ArgAction::Append)
.help("Provide custom http header(s)"),
)
}

fn add_chunker_args(cmd: Command) -> Command {
cmd.arg(
Arg::new("avg-chunk-size")
Expand Down Expand Up @@ -563,14 +570,10 @@ fn force_create_arg() -> Arg {

#[cfg(test)]
mod tests {

use super::*;
use reqwest::header::{AUTHORIZATION, CONTENT_TYPE};
use tempfile::NamedTempFile;

use crate::clone_cmd::RemoteInput;

use super::*;

fn get_num_chunk_buffers() -> usize {
match num_cpus::get() {
0 | 1 => 1,
Expand Down Expand Up @@ -780,14 +783,14 @@ mod tests {
&input.path().to_string_lossy(),
"./output.img",
])
.unwrap_or_else(|e| panic!("{}", e));
.unwrap_or_else(|e| panic!("{}", e));
assert_eq!(log, LogOpts::new(LevelFilter::Info));
assert_eq!(
opts,
CommandOpts::Clone(clone_cmd::Options {
force_create: false,
input_archive: clone_cmd::InputArchive::Local(input_path.into()),
header_checksum: Some(parse_hash_sum("5520529d1175327f9a39df0a75fe6bd314f9e6bedd89734c508a043c66066c7ada2a7b493659794f916840d976e9f0b10ec94a09caec0296ced9666998ec7977").unwrap()),
header_checksum: Some(parse_hash_sum("5520529d1175327f9a39df0a75fe6bd314f9e6bedd89734c508a043c66066c7ada2a7b493659794f916840d976e9f0b10ec94a09caec0296ced9666998ec7977").unwrap()),
output: "./output.img".into(),
seed_stdin: false,
seed_files: vec![],
Expand Down Expand Up @@ -816,7 +819,7 @@ mod tests {
opts,
CommandOpts::Clone(clone_cmd::Options {
force_create: false,
input_archive: clone_cmd::InputArchive::Remote(Box::new(RemoteInput {
input_archive: clone_cmd::InputArchive::Remote(Box::new(clone_cmd::RemoteInput {
url: "https://some-url.com/archive.cba".try_into().unwrap(),
headers: HeaderMap::new(),
receive_timeout: None,
Expand Down Expand Up @@ -861,7 +864,7 @@ mod tests {
opts,
CommandOpts::Clone(clone_cmd::Options {
force_create: true,
input_archive: clone_cmd::InputArchive::Remote(Box::new(RemoteInput {
input_archive: clone_cmd::InputArchive::Remote(Box::new(clone_cmd::RemoteInput {
url: "https://some-url.com/archive.cba".try_into().unwrap(),
headers,
receive_timeout: None,
Expand All @@ -881,13 +884,18 @@ mod tests {

#[test]
fn info_command() {
let (info, log) = parse_opts(["bita", "info", "-v", "an-input-file.cba"])
let input = NamedTempFile::new().unwrap();
let input_path = input.path();
let mut temp_file_path = input_path.to_path_buf();
temp_file_path.set_extension(".tmp");

let (info, log) = parse_opts(["bita", "info", "-v", &input.path().to_string_lossy()])
.unwrap_or_else(|e| panic!("{}", e));
assert_eq!(log, LogOpts::new(LevelFilter::Debug));
assert_eq!(
info,
CommandOpts::Info(info_cmd::Options {
input: "an-input-file.cba".into(),
input_archive: clone_cmd::InputArchive::Local(input_path.into()),
metadata_key: None,
}),
);
Expand Down
10 changes: 4 additions & 6 deletions src/compress_cmd.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -226,20 +226,18 @@ pub async fn compress_cmd(opts: Options) -> Result<()> {
chunking_algorithm: dict::chunker_parameters::ChunkingAlgorithm::FixedSize as i32,
},
};

// Construct custom metadata hashmap
let mut metadata = HashMap::new();
for (key, path) in opts.metadata_files {
let content = std::fs::read(&path).context(format!(
"Failed to read metadata file {}",
path.display()
))?;
let content = std::fs::read(&path)
.context(format!("Failed to read metadata file {}", path.display()))?;
metadata.insert(key, content);
}
for (key, value) in opts.metadata_strings {
metadata.insert(key, value.into());
}

// Build the final archive
let file_header = dict::ChunkDictionary {
rebuild_order: chunk_order.iter().map(|&index| index as u32).collect(),
Expand Down
Loading

0 comments on commit 495e30f

Please sign in to comment.