feat(mirrors)!: added ratelimiter (hopefully less bans)

shockpast · shockpast · commit f4a94aaf12b2 · 2025-03-30T22:56:39.000+03:00
chore!: updated User-Agent to keep track of old users
chore: limited "concurrent_downloads" up to 6
diff --git a/src/collector/mod.rs b/src/collector/mod.rs
@@ -33,12 +33,6 @@ pub struct CollectionUploader {
     pub username: String,
 }
 
-#[derive(Clone, Debug, Deserialize)]
-struct CollectionDate {}
-
-#[derive(Clone, Debug, Deserialize)]
-struct CollectionModes {}
-
 #[derive(Clone, Debug, Deserialize)]
 #[serde(rename_all = "lowercase")]
 enum Gamemode {
diff --git a/src/config.rs b/src/config.rs
@@ -90,6 +90,10 @@ pub fn init() -> Config {
     let contents = fs::read_to_string("config.toml").expect("config.toml doesn't exist!");
     let mut config = toml::from_str::<Config>(&contents).unwrap();
 
+    if config.user.concurrent_downloads > 6 {
+        panic!("It's highly recommended, that you won't use more than 6 \"threads\" to download maps, otherwise you will get banned from mirrors.");
+    }
+
     let osu_path = osu::find_game().unwrap();
     config.osu.songs_path = format!("{}\\Songs", osu_path);
     config.osu.collection_path = format!("{}\\collection.db", osu_path);
diff --git a/src/main.rs b/src/main.rs
@@ -7,6 +7,7 @@ use std::{
 };
 
 use clap::Parser;
+use mirrors::Ratelimiter;
 use osu_db::CollectionList;
 use sanitise_file_name::sanitise;
 use tokio::sync::{RwLock, Semaphore};
@@ -69,6 +70,7 @@ async fn main() {
     );
 
     let downloaded = Arc::new(AtomicI32::new(1));
+    let rate_limiter = Arc::new(Ratelimiter::default());
     let semaphore = Arc::new(Semaphore::new(CONFIG.user.concurrent_downloads));
 
     for beatmapset in remote_collection_info.beatmapsets {
@@ -105,11 +107,13 @@ async fn main() {
         let downloaded = Arc::clone(&downloaded);
         let remote_collection_beatmaps = Arc::clone(&remote_collection_beatmaps);
         let semaphore = Arc::clone(&semaphore);
+        let rate_limiter = Arc::clone(&rate_limiter);
 
         tokio::task::spawn(async move {
             let _permit = semaphore.acquire().await.unwrap();
+            let _rate_limiter: &Ratelimiter = &rate_limiter;
 
-            match mirror.get_file(beatmapset.id).await {
+            match mirror.get_file(beatmapset.id, _rate_limiter).await {
                 Ok(bytes) => {
                     let beatmapset_entity = &remote_collection_beatmaps
                         .beatmapsets
@@ -143,6 +147,8 @@ async fn main() {
                             .find(|b| b.checksum == beatmap.checksum)
                             .unwrap();
 
+                        info!("requests left: {}", _rate_limiter.info.read().await.remaining);
+
                         info!(
                             "({}/{}) {} - {} [{}]",
                             downloaded.load(Ordering::SeqCst),
diff --git a/src/mirrors/beatconnect.rs b/src/mirrors/beatconnect.rs
@@ -1,6 +1,6 @@
 use serde::Deserialize;
 
-use super::Mirror;
+use super::{Mirror, Ratelimiter};
 
 #[derive(Deserialize)]
 pub struct Beatconnect;
@@ -15,15 +15,19 @@ impl Mirror for Beatconnect {
         "https://beatconnect.io/b"
     }
 
-    async fn get_file(&self, id: i32) -> Result<Vec<u8>, String> {
+    async fn get_file(&self, id: i32, rate_limit: &Ratelimiter) -> Result<Vec<u8>, String> {
+        rate_limit.wait_if_needed().await;
+
         let client = reqwest::Client::new();
         let response = client
             .get(format!("{}/{}", self.get_base_url(), id))
-            .header("User-Agent", "shockpast/osu-collector-cli: 1.0.0")
+            .header("User-Agent", "shockpast/ecstasy: 1.1.2")
             .send()
             .await
             .unwrap();
 
+        rate_limit.update_rate_limit(response.headers()).await;
+
         let content_type = response
             .headers()
             .get(reqwest::header::CONTENT_TYPE)
diff --git a/src/mirrors/catboy.rs b/src/mirrors/catboy.rs
@@ -1,6 +1,6 @@
 use serde::Deserialize;
 
-use super::Mirror;
+use super::{Mirror, Ratelimiter};
 
 #[derive(Debug, Clone, Deserialize)]
 struct ErrorResponse {
@@ -20,15 +20,19 @@ impl Mirror for Catboy {
         "https://catboy.best/d"
     }
 
-    async fn get_file(&self, id: i32) -> Result<Vec<u8>, String> {
+    async fn get_file(&self, id: i32, rate_limit: &Ratelimiter) -> Result<Vec<u8>, String> {
+        rate_limit.wait_if_needed().await;
+
         let client = reqwest::Client::new();
         let response = client
             .get(format!("{}/{}", self.get_base_url(), id))
-            .header("User-Agent", "shockpast/osu-collector-cli: 1.0.0")
+            .header("User-Agent", "shockpast/ecstasy: 1.1.2")
             .send()
             .await
             .unwrap();
 
+        rate_limit.update_rate_limit(response.headers()).await;
+
         let content_type = response
             .headers()
             .get(reqwest::header::CONTENT_TYPE)
diff --git a/src/mirrors/mod.rs b/src/mirrors/mod.rs
@@ -1,12 +1,52 @@
+use std::{sync::Arc, time::{Duration, Instant}};
+
+use reqwest::header::HeaderMap;
+use tokio::sync::RwLock;
+use tracing::warn;
+
 pub mod beatconnect;
 pub mod catboy;
 pub mod nerinyan;
 pub mod osudirect;
 pub mod sayobot;
 
+#[derive(Default)]
+pub struct RatelimitInfo {
+    pub remaining: u32,
+    pub reset_at: Option<Instant>,
+}
+
+#[derive(Default)]
+pub struct Ratelimiter {
+    pub info: Arc<RwLock<RatelimitInfo>>,
+}
+
+impl Ratelimiter {
+    async fn wait_if_needed(&self) {
+        if let Some(reset_at) = self.info.read().await.reset_at {
+            if reset_at > Instant::now() {
+                let wait_duration = reset_at.duration_since(Instant::now());
+                warn!("You've hit an rate-limit, chill out, and wait for 60 seconds.");
+
+                tokio::time::sleep(wait_duration).await;
+            }
+        }
+    }
+
+    async fn update_rate_limit(&self, headers: &HeaderMap) {
+        if let Some(remaining) = headers.get("x-ratelimit-remaining").and_then(|v| v.to_str().ok()) {
+            self.info.write().await.remaining = remaining.parse().unwrap_or(0);
+        }
+        
+        if self.info.read().await.remaining <= 1 {
+            self.info.write().await.reset_at = Some(Instant::now() + Duration::from_secs(70));
+        }
+    }
+}
+
 #[async_trait::async_trait]
 pub trait Mirror {
     fn get_name(&self) -> &'static str;
     fn get_base_url(&self) -> &'static str;
-    async fn get_file(&self, id: i32) -> Result<Vec<u8>, String>;
-}
+    async fn get_file(&self, id: i32, rate_limiter: &Ratelimiter) -> Result<Vec<u8>, String>;
+}
diff --git a/src/mirrors/nerinyan.rs b/src/mirrors/nerinyan.rs
@@ -1,7 +1,6 @@
-use reqwest::StatusCode;
 use serde::Deserialize;
 
-use super::Mirror;
+use super::{Mirror, Ratelimiter};
 
 #[derive(Deserialize)]
 pub struct Nerinyan;
@@ -16,43 +15,28 @@ impl Mirror for Nerinyan {
         "https://api.nerinyan.moe/d"
     }
 
-    async fn get_file(&self, id: i32) -> Result<Vec<u8>, String> {
+    async fn get_file(&self, id: i32, rate_limit: &Ratelimiter) -> Result<Vec<u8>, String> {
+        rate_limit.wait_if_needed().await;
+
         let client = reqwest::Client::new();
         let response = client
             .get(format!("{}/{}", self.get_base_url(), id))
-            .header("User-Agent", "shockpast/osu-collector-cli: 1.0.0")
+            .header("User-Agent", "shockpast/ecstasy: 1.1.2")
             .send()
             .await
             .unwrap();
 
+        rate_limit.update_rate_limit(response.headers()).await;
+
         let content_type = response
             .headers()
             .get(reqwest::header::CONTENT_TYPE)
             .and_then(|v| v.to_str().ok())
             .map(|s| s.to_string())
             .unwrap_or_default();
-        let status_code = response.status();
 
         let bytes = response.bytes().await.map_err(|e| e.to_string())?;
 
-        if status_code.is_client_error() {
-            match status_code {
-                StatusCode::FORBIDDEN => {
-                    return Err(format!("{} possibly banned us.", self.get_name()));
-                }
-                _ => todo!(),
-            };
-        }
-
-        if status_code.is_server_error() {
-            match status_code {
-                StatusCode::BAD_GATEWAY => {
-                    panic!("{} is down, consider using other mirror.", self.get_name());
-                }
-                _ => todo!(),
-            };
-        }
-
         if content_type.contains("application/json") {
             if let Ok(json) = serde_json::from_slice::<serde_json::Value>(&bytes) {
                 return Err(json.to_string());
diff --git a/src/mirrors/osudirect.rs b/src/mirrors/osudirect.rs
@@ -1,6 +1,6 @@
 use serde::Deserialize;
 
-use super::Mirror;
+use super::{Mirror, Ratelimiter};
 
 #[derive(Debug, Clone, Deserialize)]
 struct ErrorResponse {
@@ -20,15 +20,19 @@ impl Mirror for OsuDirect {
         "https://osu.direct/api/d"
     }
 
-    async fn get_file(&self, id: i32) -> Result<Vec<u8>, String> {
+    async fn get_file(&self, id: i32, rate_limit: &Ratelimiter) -> Result<Vec<u8>, String> {
+        rate_limit.wait_if_needed().await;
+
         let client = reqwest::Client::new();
         let response = client
             .get(format!("{}/{}", self.get_base_url(), id))
-            .header("User-Agent", "shockpast/osu-collector-cli: 1.0.0")
+            .header("User-Agent", "shockpast/ecstasy: 1.1.2")
             .send()
             .await
             .unwrap();
 
+        rate_limit.update_rate_limit(response.headers()).await;
+
         let content_type = response
             .headers()
             .get(reqwest::header::CONTENT_TYPE)
diff --git a/src/mirrors/sayobot.rs b/src/mirrors/sayobot.rs
@@ -1,6 +1,6 @@
 use serde::Deserialize;
 
-use super::Mirror;
+use super::{Mirror, Ratelimiter};
 
 #[derive(Debug, Clone, Deserialize)]
 struct ErrorResponse {
@@ -20,15 +20,19 @@ impl Mirror for Sayobot {
         "https://txy1.sayobot.cn/beatmaps/download/full"
     }
 
-    async fn get_file(&self, id: i32) -> Result<Vec<u8>, String> {
+    async fn get_file(&self, id: i32, rate_limit: &Ratelimiter) -> Result<Vec<u8>, String> {
+        rate_limit.wait_if_needed().await;
+
         let client = reqwest::Client::new();
         let response = client
-            .get(format!("{}/{}?server=auto", self.get_base_url(), id))
-            .header("User-Agent", "shockpast/osu-collector-cli: 1.0.0")
+            .get(format!("{}/{}", self.get_base_url(), id))
+            .header("User-Agent", "shockpast/ecstasy: 1.1.2")
             .send()
             .await
             .unwrap();
 
+        rate_limit.update_rate_limit(response.headers()).await;
+
         let content_type = response
             .headers()
             .get(reqwest::header::CONTENT_TYPE)
diff --git a/src/utilities/osu.rs b/src/utilities/osu.rs
@@ -43,6 +43,7 @@ pub fn find_game() -> Result<String, std::io::Error> {
 }
 
 #[cfg(target_os = "linux")]
-pub async fn find_game() -> Option<&'static str> {
-    todo!("Linux is currently not supported, since we don't have any experience with it.")
+pub async fn find_game() -> Result<String, std::io::Error> {
+    let path = std::env::var("OSU_FOLDER").expect("'OSU_FOLDER' export is not defined (e.g.: '$HOME\\osu')");
+    Ok(path)
 }
diff --git a/src/utilities/speedtest.rs b/src/utilities/speedtest.rs
@@ -2,7 +2,7 @@ use std::{collections::VecDeque, time::Instant};
 
 use tracing::{debug, info};
 
-use crate::config::MirrorType;
+use crate::{config::MirrorType, mirrors::Ratelimiter};
 
 async fn test_download(client: &reqwest::Client, payload_size_bytes: usize) -> f64 {
     let req = client.get(format!(
@@ -39,10 +39,11 @@ pub async fn benchmark() {
 
     for (index, mirror_type) in MirrorType::ALL.iter().enumerate() {
         let mirror = mirror_type.get_mirror();
+        let rate_limiter = Ratelimiter::default();
 
         let start = Instant::now();
 
-        let file = mirror.get_file(1030499).await.unwrap();
+        let file = mirror.get_file(1030499, &rate_limiter).await.unwrap();
         if file.len() as f64 <= 0.0 {
             continue;
         }

Original file line number	Diff line number	Diff line change
`@@ -43,6 +43,7 @@ pub fn find_game() -> Result<String, std::io::Error> {`
`43`	`43`	`}`
`44`	`44`
`45`	`45`	`#[cfg(target_os = "linux")]`
`46`		`-pub async fn find_game() -> Option<&'static str> {`
`47`		`- todo!("Linux is currently not supported, since we don't have any experience with it.")`
	`46`	`+pub async fn find_game() -> Result<String, std::io::Error> {`
	`47`	`+ let path = std::env::var("OSU_FOLDER").expect("'OSU_FOLDER' export is not defined (e.g.: '$HOME\\osu')");`
	`48`	`+ Ok(path)`
`48`	`49`	`}`