Retry guest login when throttled

Author: Hamuko

Cargo.lock

@@ -10,6 +10,7 @@ crossfire = { version = "2.1.6", optional = true }
 futures-util = "0.3"
 html_parser = "0.7"
 log = "0.4"
+regex = "1.12"
 reqwest = { version = "0.12", default-features = false, features = ["json"] }
 rust_socketio = { version = "0.6", features = ["async"] }
 serde = { version = "1.0", features = ["derive"] }

Cargo.toml

@@ -9,11 +9,12 @@ use rust_socketio::asynchronous::{Client, ClientBuilder};
 use rust_socketio::{Payload, TransportType};
 use serde_json::{Value, json};
 use simple_logger::SimpleLogger;
+use std::sync::Arc;
 use tokio::fs::File;
 use tokio::io::{AsyncWriteExt, BufWriter};
 use tokio::select;
 use tokio::signal;
-use tokio::time::Duration;
+use tokio::time::{Duration, sleep};
 use tokio_util::sync::CancellationToken;
 
 const WRITE_BUFFER_SIZE: usize = 8 * 1024; // 8 KiB
@@ -54,6 +55,11 @@ enum Event {
     Terminate,
 }
 
+enum LoginError {
+    Throttled(u64),
+    Unknown(Option<String>),
+}
+
 enum SocketAddressError {
     NotFound,
     Parse(serde_json::Error),
@@ -73,7 +79,7 @@ async fn create_chat_log_file(channel: &str) -> File {
     file
 }
 
-fn handle_login_event(values: Vec<Value>) {
+fn handle_login_event(values: Vec<Value>) -> Result<(), LoginError> {
     for value in values {
         let login: data::Login = match serde_json::from_value(value) {
             Ok(v) => v,
@@ -86,15 +92,21 @@ fn handle_login_event(values: Vec<Value>) {
         if login.success {
             log::info!(
                 "Logged in as guest {}",
-                login.name.unwrap_or("Unknown".into())
+                login.name.unwrap_or("<UNKNOWN>".into())
             );
+            return Ok(());
         } else {
-            log::warn!(
-                "Login failed: {}",
-                login.error.unwrap_or("Unknown error".into())
-            );
+            if let Some(error_message) = login.error {
+                if let Some(duration) = utils::check_throttling_error(&error_message) {
+                    return Err(LoginError::Throttled(duration));
+                } else {
+                    return Err(LoginError::Unknown(Some(error_message)));
+                }
+            }
+            return Err(LoginError::Unknown(None));
         }
     }
+    Err(LoginError::Unknown(None))
 }
 
 /// Join a channel on the Cytube server.
@@ -211,16 +223,18 @@ async fn main() {
     };
 
     let channel_name = args.channel.clone();
+    let guest_login = Arc::new(args.guest_login);
+    let connect_guest_login = guest_login.clone();
     let socket = ClientBuilder::new(socket_address)
         .transport_type(TransportType::Any)
         .on(rust_socketio::Event::Connect, move |_, client| {
             let channel_name = channel_name.clone();
-            let guest_login = args.guest_login.clone();
+            let username = connect_guest_login.clone();
             async move {
                 log::info!("Connected to server");
                 join_channel(&client, &channel_name).await;
-                if let Some(username) = guest_login {
-                    login_as_guest(&client, &username).await;
+                if let Some(username) = &*username {
+                    login_as_guest(&client, username).await;
                 }
             }
             .boxed()
@@ -286,6 +300,7 @@ async fn main() {
         .expect("Connection failed");
 
     let channel_name = args.channel.clone();
+    let manager_socket = socket.clone();
     let manager = tokio::spawn(async move {
         let mut last_timestamp: u64 = 0;
         while let Some(event) = channel::read_event(&mut rx).await {
@@ -327,7 +342,24 @@ async fn main() {
                 Event::Disconnect => {
                     log::warn!("Client disconnected from server");
                 }
-                Event::Login(values) => handle_login_event(values),
+                Event::Login(values) => match handle_login_event(values) {
+                    Ok(_) => {}
+                    Err(LoginError::Throttled(seconds)) => {
+                        log::warn!("Guest login throttled; retrying in {} seconds", seconds);
+                        let username = guest_login.clone();
+                        let socket = manager_socket.clone();
+                        let _ = tokio::spawn(async move {
+                            sleep(Duration::from_secs(seconds)).await;
+                            if let Some(username) = &*username {
+                                login_as_guest(&socket, username).await;
+                            }
+                        });
+                    }
+                    Err(LoginError::Unknown(Some(message))) => {
+                        log::warn!("Login failed: {}", message)
+                    }
+                    Err(LoginError::Unknown(None)) => log::warn!("Login failed"),
+                },
                 Event::RotateLog => {
                     log::info!("Rotating log file...");
                     match file_buffer.flush().await {

src/main.rs

@@ -1,3 +1,18 @@
+use regex::Regex;
+use std::sync::LazyLock;
+
+/// Determine if the error message indicates a throttling error.
+pub fn check_throttling_error(string: &str) -> Option<u64> {
+    static RE: LazyLock<Regex> = LazyLock::new(|| {
+        Regex::new(r"Guest logins are restricted to one per IP address per (\d+) seconds").unwrap()
+    });
+    if let Some(captures) = RE.captures(string) {
+        captures.get(1).unwrap().as_str().parse::<u64>().ok()
+    } else {
+        None
+    }
+}
+
 /// Parse host from plain domain name or URL.
 pub fn parse_domain(s: &str) -> Result<url::Host, String> {
     if let Ok(host) = url::Host::parse(s) {
@@ -15,6 +30,14 @@ pub fn parse_domain(s: &str) -> Result<url::Host, String> {
 mod tests {
     use test_case::test_case;
 
+    #[test_case("Guest logins are restricted to one per IP address per 60 seconds.", Some(60); "normal error")]
+    #[test_case("Guest logins are restricted to one per IP address per 9999 seconds.", Some(9999); "longer error")]
+    #[test_case("Guest logins are restricted to one per IP address per N seconds.", None; "non-numerical error")]
+    #[test_case("Something bad happened", None; "unknown error")]
+    fn check_throttling_error(input: &str, expected: Option<u64>) {
+        assert_eq!(super::check_throttling_error(input), expected);
+    }
+
     #[test_case("cytu.be", Some("cytu.be"); "plain domain")]
     #[test_case("https://cytu.be", Some("cytu.be"); "URL")]
     #[test_case("@t!", None; "invalid characters")]