Add download/request speed limiter (#250)

crunchy-cli-core/Cargo.toml

@@ -13,28 +13,32 @@ openssl-tls-static = ["reqwest/native-tls", "reqwest/native-tls-alpn", "reqwest/
 
 [dependencies]
 anyhow = "1.0"
+async-speed-limit = "0.4"
 async-trait = "0.1"
 clap = { version = "4.4", features = ["derive", "string"] }
 chrono = "0.4"
-crunchyroll-rs = { version = "0.7.0", features = ["dash-stream"] }
+crunchyroll-rs = { version = "0.8.0", features = ["dash-stream", "experimental-stabilizations", "tower"] }
 ctrlc = "3.4"
 dialoguer = { version = "0.11", default-features = false }
 dirs = "5.0"
 derive_setters = "0.1"
+futures-util = { version = "0.3", features = ["io"] }
 fs2 = "0.4"
+http = "0.2"
 indicatif = "0.17"
 lazy_static = "1.4"
 log = { version = "0.4", features = ["std"] }
 num_cpus = "1.16"
 regex = "1.10"
-reqwest = { version = "0.11", default-features = false, features = ["socks"] }
+reqwest = { version = "0.11", default-features = false, features = ["socks", "stream"] }
 serde = "1.0"
 serde_json = "1.0"
 serde_plain = "1.0"
 shlex = "1.2"
 sys-locale = "0.3"
 tempfile = "3.8"
 tokio = { version = "1.34", features = ["io-util", "macros", "net", "rt-multi-thread", "time"] }
+tower-service = "0.3"
 rustls-native-certs = { version = "0.6", optional = true }
 
 [target.'cfg(not(target_os = "windows"))'.dependencies]

crunchy-cli-core/src/archive/command.rs

@@ -242,7 +242,7 @@ impl Execute for Archive {
 
                 format.visual_output(&path);
 
-                downloader.download(&ctx, &path).await?
+                downloader.download(&path).await?
             }
         }
 

crunchy-cli-core/src/download/command.rs

@@ -257,7 +257,7 @@ impl Execute for Download {
 
                 format.visual_output(&path);
 
-                downloader.download(&ctx, &path).await?
+                downloader.download(&path).await?
             }
         }
 

crunchy-cli-core/src/lib.rs

@@ -17,6 +17,7 @@ mod login;
 mod search;
 mod utils;
 
+use crate::utils::rate_limit::RateLimiterService;
 pub use archive::Archive;
 use dialoguer::console::Term;
 pub use download::Download;
@@ -66,6 +67,15 @@ pub struct Cli {
     #[arg(global = true, long)]
     user_agent: Option<String>,
 
+    #[arg(
+        help = "Maximal speed to download/request (may be a bit off here and there). Must be in format of <number>[B|KB|MB]"
+    )]
+    #[arg(
+        long_help = "Maximal speed to download/request (may be a bit off here and there). Must be in format of <number>[B|KB|MB] (e.g. 500KB or 10MB)"
+    )]
+    #[arg(global = true, long, value_parser = crate::utils::clap::clap_parse_speed_limit)]
+    speed_limit: Option<u32>,
+
     #[clap(subcommand)]
     command: Command,
 }
@@ -264,39 +274,44 @@ async fn crunchyroll_session(cli: &mut Cli) -> Result<Crunchyroll> {
         lang
     };
 
-    let mut builder = Crunchyroll::builder()
-        .locale(locale)
-        .client({
-            let mut builder = CrunchyrollBuilder::predefined_client_builder();
-            if let Some(p) = &cli.proxy {
-                builder = builder.proxy(p.clone())
-            }
-            if let Some(ua) = &cli.user_agent {
-                builder = builder.user_agent(ua)
-            }
+    let client = {
+        let mut builder = CrunchyrollBuilder::predefined_client_builder();
+        if let Some(p) = &cli.proxy {
+            builder = builder.proxy(p.clone())
+        }
+        if let Some(ua) = &cli.user_agent {
+            builder = builder.user_agent(ua)
+        }
 
-            #[cfg(any(feature = "openssl-tls", feature = "openssl-tls-static"))]
-            let client = {
-                let mut builder = builder.use_native_tls().tls_built_in_root_certs(false);
+        #[cfg(any(feature = "openssl-tls", feature = "openssl-tls-static"))]
+        let client = {
+            let mut builder = builder.use_native_tls().tls_built_in_root_certs(false);
 
-                for certificate in rustls_native_certs::load_native_certs().unwrap() {
-                    builder = builder.add_root_certificate(
-                        reqwest::Certificate::from_der(certificate.0.as_slice()).unwrap(),
-                    )
-                }
+            for certificate in rustls_native_certs::load_native_certs().unwrap() {
+                builder = builder.add_root_certificate(
+                    reqwest::Certificate::from_der(certificate.0.as_slice()).unwrap(),
+                )
+            }
 
-                builder.build().unwrap()
-            };
-            #[cfg(not(any(feature = "openssl-tls", feature = "openssl-tls-static")))]
-            let client = builder.build().unwrap();
+            builder.build().unwrap()
+        };
+        #[cfg(not(any(feature = "openssl-tls", feature = "openssl-tls-static")))]
+        let client = builder.build().unwrap();
 
-            client
-        })
+        client
+    };
+
+    let mut builder = Crunchyroll::builder()
+        .locale(locale)
+        .client(client.clone())
         .stabilization_locales(cli.experimental_fixes)
         .stabilization_season_number(cli.experimental_fixes);
     if let Command::Download(download) = &cli.command {
         builder = builder.preferred_audio_locale(download.audio.clone())
     }
+    if let Some(speed_limit) = cli.speed_limit {
+        builder = builder.middleware(RateLimiterService::new(speed_limit, client));
+    }
 
     let root_login_methods_count = cli.login_method.credentials.is_some() as u8
         + cli.login_method.etp_rt.is_some() as u8

crunchy-cli-core/src/utils/clap.rs

@@ -9,3 +9,20 @@ pub fn clap_parse_resolution(s: &str) -> Result<Resolution, String> {
 pub fn clap_parse_proxy(s: &str) -> Result<Proxy, String> {
     Proxy::all(s).map_err(|e| e.to_string())
 }
+
+pub fn clap_parse_speed_limit(s: &str) -> Result<u32, String> {
+    let quota = s.to_lowercase();
+
+    let bytes = if let Ok(b) = quota.parse() {
+        b
+    } else if let Ok(b) = quota.trim_end_matches('b').parse::<u32>() {
+        b
+    } else if let Ok(kb) = quota.trim_end_matches("kb").parse::<u32>() {
+        kb * 1024
+    } else if let Ok(mb) = quota.trim_end_matches("mb").parse::<u32>() {
+        mb * 1024 * 1024
+    } else {
+        return Err("Invalid speed limit".to_string());
+    };
+    Ok(bytes)
+}