From edf24cac96848db9847acd751d53a0bdf3db657c Mon Sep 17 00:00:00 2001 From: Quaternions Date: Sat, 30 Dec 2023 18:00:51 -0800 Subject: [PATCH] async download --- Cargo.lock | 233 +++++++++++++++++++++++++++++++++++++++++++++++++++- Cargo.toml | 6 +- src/main.rs | 65 ++++++++++++++- 3 files changed, 300 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fec164e..21d5f59 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -98,12 +98,16 @@ version = "0.1.0" dependencies = [ "anyhow", "clap", + "flate2", + "futures", "lazy-regex", + "pollster", "rbx_binary", "rbx_dom_weak", "rbx_reflection_database", "rbx_xml", "reqwest", + "tokio", ] [[package]] @@ -249,6 +253,34 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f7144d30dcf0fafbce74250a3963025d8d52177934239851c917d29f1df280c2" +[[package]] +name = "cookie" +version = "0.16.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e859cd57d0710d9e06c381b550c06e76992472a8c6d527aecd2fc673dcc231fb" +dependencies = [ + "percent-encoding", + "time", + "version_check", +] + +[[package]] +name = "cookie_store" +version = "0.16.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d606d0fba62e13cf04db20536c05cb7f13673c161cb47a47a82b9b9e7d3f1daa" +dependencies = [ + "cookie", + "idna 0.2.3", + "log", + "publicsuffix", + "serde", + "serde_derive", + "serde_json", + "time", + "url", +] + [[package]] name = "core-foundation" version = "0.9.4" @@ -265,6 +297,24 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06ea2b9bc92be3c2baa9334a323ebca2d6f074ff852cd1d7b11064035cd3868f" +[[package]] +name = "crc32fast" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b540bd8bc810d3885c6ea91e2018302f68baba2129ab3e88f32389ee9370880d" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "deranged" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b42b6fa04a440b495c8b04d0e71b707c585f83cb9cb28cf8cd0d976c315e31b4" +dependencies = [ + "powerfmt", +] + [[package]] name = "encoding_rs" version = "0.8.33" @@ -296,6 +346,16 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "25cbce373ec4653f1a01a31e8a5e5ec0c622dc27ff9c4e6606eefef5cbbed4a5" +[[package]] +name = "flate2" +version = "1.0.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46303f565772937ffe1d394a4fac6f411c6013172fadde9dcdb1e147a086940e" +dependencies = [ + "crc32fast", + "miniz_oxide", +] + [[package]] name = "fnv" version = "1.0.7" @@ -326,6 +386,21 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "futures" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "645c6916888f6cb6350d2550b80fb63e734897a8498abe35cfb732b6487804b0" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.30" @@ -333,6 +408,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eac8f7d7865dcb88bd4373ab671c8cf4508703796caa2b1985a9ca867b3fcb78" dependencies = [ "futures-core", + "futures-sink", ] [[package]] @@ -341,6 +417,34 @@ version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dfc6580bb841c5a68e9ef15c77ccc837b40a7504914d52e47b8b0e9bbda25a1d" +[[package]] +name = "futures-executor" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a576fc72ae164fca6b9db127eaa9a9dda0d61316034f33a0a0d4eda41f02b01d" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a44623e20b9681a318efdd71c299b6b222ed6f231972bfe2f224ebad6311f0c1" + +[[package]] +name = "futures-macro" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87750cf4b7a4c0625b1529e4c543c2182106e4dedc60a2a6455e00d212c489ac" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "futures-sink" version = "0.3.30" @@ -359,10 +463,16 @@ version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d6401deb83407ab3da39eba7e33987a73c3df0c82b4bb5813ee871c19c41d48" dependencies = [ + "futures-channel", "futures-core", + "futures-io", + "futures-macro", + "futures-sink", "futures-task", + "memchr", "pin-project-lite", "pin-utils", + "slab", ] [[package]] @@ -413,6 +523,12 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8" +[[package]] +name = "hermit-abi" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d77f7ec81a6d05a3abb01ab6eb7590f6083d08449fe5a1c8b1e620283546ccb7" + [[package]] name = "http" version = "0.2.11" @@ -484,6 +600,27 @@ dependencies = [ "tokio-native-tls", ] +[[package]] +name = "idna" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "418a0a6fab821475f634efe3ccc45c013f742efe03d853e8d3355d5cb850ecf8" +dependencies = [ + "matches", + "unicode-bidi", + "unicode-normalization", +] + +[[package]] +name = "idna" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e14ddfc70884202db2244c223200c204c2bda1bc6e0998d11b5e024d657209e6" +dependencies = [ + "unicode-bidi", + "unicode-normalization", +] + [[package]] name = "idna" version = "0.5.0" @@ -592,6 +729,12 @@ dependencies = [ "libc", ] +[[package]] +name = "matches" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2532096657941c2fea9c289d370a250971c689d4f143798ff67113ec042024a5" + [[package]] name = "memchr" version = "2.7.1" @@ -651,6 +794,16 @@ dependencies = [ "autocfg", ] +[[package]] +name = "num_cpus" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4161fcb6d602d4d2081af7c3a45852d875a03dd337a6bfdd6e06407b61342a43" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "object" version = "0.32.2" @@ -740,6 +893,18 @@ version = "0.3.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69d3587f8a9e599cc7ec2c00e331f71c4e69a5f9a4b8a6efd5b07466b9736f9a" +[[package]] +name = "pollster" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22686f4785f02a4fcc856d3b3bb19bf6c8160d103f7a99cc258bddd0251dc7f2" + +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + [[package]] name = "ppv-lite86" version = "0.2.17" @@ -774,6 +939,22 @@ dependencies = [ "syn", ] +[[package]] +name = "psl-types" +version = "2.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33cb294fe86a74cbcf50d4445b37da762029549ebeea341421c7c70370f86cac" + +[[package]] +name = "publicsuffix" +version = "2.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96a8c1bda5ae1af7f99a2962e49df150414a43d62404644d98dd5c3a93d07457" +dependencies = [ + "idna 0.3.0", + "psl-types", +] + [[package]] name = "quote" version = "1.0.33" @@ -936,6 +1117,8 @@ checksum = "37b1ae8d9ac08420c66222fb9096fc5de435c3c48542bc5336c51892cffafb41" dependencies = [ "base64 0.21.5", "bytes", + "cookie", + "cookie_store", "encoding_rs", "futures-core", "futures-util", @@ -1178,6 +1361,35 @@ dependencies = [ "syn", ] +[[package]] +name = "time" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f657ba42c3f86e7680e53c8cd3af8abbe56b5491790b46e22e19c0d57463583e" +dependencies = [ + "deranged", + "itoa", + "powerfmt", + "serde", + "time-core", + "time-macros", +] + +[[package]] +name = "time-core" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef927ca75afb808a4d64dd374f00a2adf8d0fcff8e7b184af886c3c87ec4a3f3" + +[[package]] +name = "time-macros" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26197e33420244aeb70c3e8c78376ca46571bc4e701e4791c2cd9f57dcb3a43f" +dependencies = [ + "time-core", +] + [[package]] name = "tinyvec" version = "1.6.0" @@ -1203,11 +1415,24 @@ dependencies = [ "bytes", "libc", "mio", + "num_cpus", "pin-project-lite", "socket2", + "tokio-macros", "windows-sys 0.48.0", ] +[[package]] +name = "tokio-macros" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b8a1e28f2deaa14e508979454cb3a223b10b938b45af148bc0986de36f1923b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tokio-native-tls" version = "0.3.1" @@ -1291,7 +1516,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "31e6302e3bb753d46e83516cae55ae196fc0c309407cf11ab35cc51a4c2a4633" dependencies = [ "form_urlencoded", - "idna", + "idna 0.5.0", "percent-encoding", ] @@ -1307,6 +1532,12 @@ version = "0.2.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" +[[package]] +name = "version_check" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f" + [[package]] name = "want" version = "0.3.1" diff --git a/Cargo.toml b/Cargo.toml index cd409de..870cdbd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,9 +8,13 @@ edition = "2021" [dependencies] anyhow = "1.0.75" clap = { version = "4.4.2", features = ["derive"] } +flate2 = "1.0.28" +futures = "0.3.30" lazy-regex = "3.1.0" +pollster = "0.3.0" rbx_binary = "0.7.1" rbx_dom_weak = "2.5.0" rbx_reflection_database = "0.2.7" rbx_xml = "0.13.1" -reqwest = "0.11.23" +reqwest = { version = "0.11.23", features = ["cookies"] } +tokio = { version = "1.35.1", features = ["macros", "rt-multi-thread", "fs"] } diff --git a/src/main.rs b/src/main.rs index af26209..0b24b92 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,7 @@ use std::io::{Read,Seek}; use clap::{Args,Parser,Subcommand}; use anyhow::Result as AResult; +use futures::StreamExt; type AssetID=u64; @@ -28,7 +29,8 @@ struct AssetIDList{ asset_ids:Vec, } -fn main()->AResult<()>{ +#[tokio::main] +async fn main()->AResult<()>{ let cli=Cli::parse(); match cli.command{ Commands::Download(asset_id_list)=>download_list(asset_id_list.asset_ids).await, @@ -36,10 +38,69 @@ fn main()->AResult<()>{ } } +enum ReaderType<'a,R:Read+Seek>{ + GZip(flate2::read::GzDecoder<&'a mut R>), + Raw(&'a mut R), +} + +fn maybe_gzip_decode(input:&mut R)->AResult>{ + let mut first_2=[0u8;2]; + if let (Ok(()),Ok(()))=(std::io::Read::read_exact(input,&mut first_2),std::io::Seek::rewind(input)){ + match &first_2{ + b"\x1f\x8b"=>Ok(ReaderType::GZip(flate2::read::GzDecoder::new(input))), + _=>Ok(ReaderType::Raw(input)), + } + }else{ + Err(anyhow::Error::msg("failed to peek")) + } +} + fn upload_file(_path:std::path::PathBuf,_asset_id:AssetID)->AResult<()>{ Ok(()) } -async fn download_list(_asset_ids:Vec)->AResult<()>{ +const CONCURRENT_REQUESTS:usize=8; + +fn read_readable(mut readable:impl Read)->AResult>{ + let mut contents=Vec::new(); + readable.read_to_end(&mut contents)?; + Ok(contents) +} + +async fn download_list(asset_ids:Vec)->AResult<()>{ + let cookie=format!(".ROBLOSECURITY={}",std::env::var("RBXCOOKIE_PROJECTSLIME")?); + let client=reqwest::Client::new(); + futures::stream::iter(asset_ids) + .map(|asset_id|{ + let client=&client; + let cookie=cookie.as_str(); + async move{ + let resp=client.get(format!("https://assetdelivery.roblox.com/v1/asset/?ID={}",asset_id)) + .header("Cookie",cookie) + .send().await?; + Ok((asset_id,resp.bytes().await?)) + } + }) + .buffer_unordered(CONCURRENT_REQUESTS) + .for_each(|b:AResult<_>|async{ + match b{ + Ok((asset_id,body))=>{ + let dest=std::path::PathBuf::from(asset_id.to_string()); + let contents=match maybe_gzip_decode(&mut std::io::Cursor::new(body)){ + Ok(ReaderType::GZip(readable))=>read_readable(readable), + Ok(ReaderType::Raw(readable))=>read_readable(readable), + Err(e)=>Err(e), + }; + match contents{ + Ok(data)=>match tokio::fs::write(dest,data).await{ + Err(e)=>eprintln!("fs error: {}",e), + _=>(), + }, + Err(e)=>eprintln!("gzip error: {}",e), + }; + }, + Err(e)=>eprintln!("dl error: {}",e), + } + }).await; Ok(()) } \ No newline at end of file