From b21964eb4180637174ca077a8ecf716cdc3e65c0 Mon Sep 17 00:00:00 2001 From: Hatter Jiang Date: Sun, 25 Aug 2024 10:32:19 +0800 Subject: [PATCH] feat: update osssendfile-rs, check sha256 --- osssendfile-rs/Cargo.lock | 7 ++++ osssendfile-rs/Cargo.toml | 1 + osssendfile-rs/src/main.rs | 73 ++++++++++++++++++++++++++++---------- 3 files changed, 62 insertions(+), 19 deletions(-) diff --git a/osssendfile-rs/Cargo.lock b/osssendfile-rs/Cargo.lock index bde57fe..d33f903 100644 --- a/osssendfile-rs/Cargo.lock +++ b/osssendfile-rs/Cargo.lock @@ -485,6 +485,12 @@ version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + [[package]] name = "http" version = "1.1.0" @@ -840,6 +846,7 @@ dependencies = [ "aes-gcm-stream", "base64 0.22.1", "clap", + "hex", "oss", "rand 0.8.5", "reqwest", diff --git a/osssendfile-rs/Cargo.toml b/osssendfile-rs/Cargo.toml index af91d76..e934c0e 100644 --- a/osssendfile-rs/Cargo.toml +++ b/osssendfile-rs/Cargo.toml @@ -7,6 +7,7 @@ edition = "2021" aes-gcm-stream = "0.2" base64 = "0.22.1" clap = { version = "4.5", features = ["derive"] } +hex = "0.4.3" oss = "0.3.2" rand = "0.8.5" reqwest = "0.12" diff --git a/osssendfile-rs/src/main.rs b/osssendfile-rs/src/main.rs index 89b4dd0..6fce897 100755 --- a/osssendfile-rs/src/main.rs +++ b/osssendfile-rs/src/main.rs @@ -5,6 +5,7 @@ //! aes-gcm-stream = "0.2" //! base64 = "0.22.1" //! clap = { version = "4.5", features = ["derive"] } +//! hex = "0.4.3" //! oss = "0.3.2" //! rand = "0.8.5" //! reqwest = "0.12" @@ -41,7 +42,7 @@ const CREATE_STS_URL: &str = "https://hatter.ink/oidc/create_sts.json"; const ADD_DOC_URL: &str = "https://playsecurity.org/doc/addDoc.jsonp"; #[derive(Debug, Parser)] -#[command(name = "osssendfile-rs")] +#[command(name = "osssendfile-rs", bin_name = "osssendfile.rs")] #[command(about = "OSS send file Rust edition", long_about = None)] struct OssSendFileArgs { /// Config file, default location: ~/.jssp/config/osssendfile.json @@ -51,8 +52,8 @@ struct OssSendFileArgs { // #[arg(long)] // no_enc: bool, // /// Do remove source file - // #[arg(long)] - // remove_source_file: bool, + #[arg(long)] + remove_source_file: bool, // /// JWK // #[arg(long, short = 'j')] // jwk: Option, @@ -93,6 +94,16 @@ struct Sts { security_token: String, } +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +struct SendFileResponse { + status: i32, + id: i64, + length: i64, + etag: String, + sha256: String, +} + #[tokio::main] async fn main() -> XResult<()> { let args = OssSendFileArgs::parse(); @@ -147,13 +158,13 @@ async fn main() -> XResult<()> { let object_url = oss_client.generate_signed_get_url(&oss_send_file_config.bucket, &temp_oss_filename, 600); information!("Send URL to play security: {}", object_url); - send_file(&client, - &oss_send_file_config.token, - &object_url, - &args.filename.clone().unwrap_or_else(|| get_filename(filename)), - args.keywords.as_deref().unwrap_or(""), - temp_key).await?; - information!("Send URL success"); + let send_file_response = send_file(&client, + &oss_send_file_config.token, + &object_url, + &args.filename.clone().unwrap_or_else(|| get_filename(filename)), + args.keywords.as_deref().unwrap_or(""), + temp_key).await?; + information!("Send URL success, file id: {}", &send_file_response.id); information!("Delete object: {}", &temp_oss_filename); let delete_object_response = oss_client.delete_file( @@ -164,10 +175,24 @@ async fn main() -> XResult<()> { information!("Delete object success"); match fs::remove_file(&temp_file) { - Ok(_) => information!("Delete temp file: {:?}", &temp_file), + Ok(_) => information!("Delete temp file success: {:?}", &temp_file), Err(e) => warning!("Delete temp file failed: {}", e), } + // check digest + let source_file_sha256 = hex::encode(digest_sha256(&source_file)?); + if source_file_sha256 != send_file_response.sha256 { + failure!("Check send file digest failed: {} vs {}", source_file_sha256, send_file_response.sha256); + return simple_error!("Check send file digest failed."); + } + + if args.remove_source_file { + match fs::remove_file(&source_file) { + Ok(_) => information!("Delete source file success: {:?}", &source_file), + Err(e) => warning!("Delete source file failed: {}", e), + } + } + success!("File {} upload success", filename); Ok(()) } @@ -221,16 +246,16 @@ fn encrypt(source_file: &PathBuf, dest_file: &PathBuf) -> XResult<[u8; 16]> { let mut source_file_read = File::open(source_file)?; let mut dest_file_write = File::create(dest_file)?; - let mut written = 0u64; + // let mut written = 0u64; let mut buf: [u8; DEFAULT_BUF_SIZE] = [0u8; DEFAULT_BUF_SIZE]; loop { let len = match source_file_read.read(&mut buf) { Ok(0) => { let (final_block, tag) = encryptor.finalize(); dest_file_write.write(&final_block)?; - written += final_block.len() as u64; + // written += final_block.len() as u64; dest_file_write.write(&tag)?; - written += tag.len() as u64; + // written += tag.len() as u64; return Ok(temp_key); } Ok(len) => len, @@ -239,10 +264,18 @@ fn encrypt(source_file: &PathBuf, dest_file: &PathBuf) -> XResult<[u8; 16]> { }; let encrypted = encryptor.update(&buf[0..len]); dest_file_write.write(&encrypted)?; - written += encrypted.len() as u64; + // written += encrypted.len() as u64; } } +fn digest_sha256(source_file: &PathBuf) -> XResult> { + let mut sha256 = sha2::Sha256::new(); + let content = opt_result!(fs::read(&source_file), "Read file: {:?} failed: {}", source_file); + sha256.update(&content); + let digest = sha256.finalize().as_slice().to_vec(); + Ok(digest) +} + async fn request_sts(client: &Client, oss_send_file_config: &OssSendFileConfig) -> XResult { let mut params = HashMap::new(); params.insert("client_id", &oss_send_file_config.oidc.client_id); @@ -256,7 +289,7 @@ async fn request_sts(client: &Client, oss_send_file_config: &OssSendFileConfig) parse_sts_response(response).await } -async fn send_file(client: &Client, token: &str, url: &str, title: &str, keywords: &str, key: [u8; 16]) -> XResult<()> { +async fn send_file(client: &Client, token: &str, url: &str, title: &str, keywords: &str, key: [u8; 16]) -> XResult { let mut params = HashMap::new(); params.insert("jsonp", "1".to_string()); params.insert("token", token.to_string()); @@ -275,9 +308,11 @@ async fn send_file(client: &Client, token: &str, url: &str, title: &str, keyword } let response_bytes = response.bytes().await?.as_ref().to_vec(); - let add_doc_result = String::from_utf8_lossy(&response_bytes); - success!("Add doc success: {}", add_doc_result); - Ok(()) + let send_file_response: SendFileResponse = opt_result!( + serde_json::from_slice(&response_bytes), "Parse send file response failed: {}"); + + success!("Add doc id: {} success, sha256: {}", send_file_response.id, send_file_response.sha256); + Ok(send_file_response) } async fn parse_sts_response(response: Response) -> XResult {