From e3038ca6c5981808d6b1d5aedd0f76215086be18 Mon Sep 17 00:00:00 2001
From: Bulat Kurbanov <kurbanovbul@gmail.com>
Date: Sun, 7 Jan 2024 22:18:58 +0100
Subject: [PATCH] Add bot registration limit

---
 src/bots/bots_manager/mod.rs      |  3 ++-
 src/bots/bots_manager/register.rs | 40 ++++++++++++++++++++++---------
 src/bots/bots_manager/strings.rs  |  2 ++
 3 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/src/bots/bots_manager/mod.rs b/src/bots/bots_manager/mod.rs
index 53f9cb1..0068ffe 100644
--- a/src/bots/bots_manager/mod.rs
+++ b/src/bots/bots_manager/mod.rs
@@ -22,8 +22,9 @@ pub async fn message_handler(
 
     let message_text = match result {
         register::RegisterStatus::Success { ref username } => format_registered_message(username),
-        register::RegisterStatus::WrongToken => strings::ERROR_MESSAGE.to_string(),
         register::RegisterStatus::RegisterFail => strings::ALREADY_REGISTERED.to_string(),
+        register::RegisterStatus::LimitExtended => strings::LIMIT_EXTENDED_MESSAGE.to_string(),
+        register::RegisterStatus::WrongToken => strings::ERROR_MESSAGE.to_string(),
     };
 
     bot.send_message(message.chat.id, message_text)
diff --git a/src/bots/bots_manager/register.rs b/src/bots/bots_manager/register.rs
index cc87abc..e2d9ca6 100644
--- a/src/bots/bots_manager/register.rs
+++ b/src/bots/bots_manager/register.rs
@@ -11,6 +11,14 @@ pub enum RegisterStatus {
     Success { username: String },
     WrongToken,
     RegisterFail,
+    LimitExtended,
+}
+
+#[derive(Debug)]
+pub enum RegisterRequestStatus {
+    Success,
+    LimitExtended,
+    UnknownError { status_code: u16 },
 }
 
 async fn get_bot_username(token: &str) -> Option<String> {
@@ -27,7 +35,7 @@ async fn make_register_request(
     user_id: UserId,
     username: &str,
     token: &str,
-) -> Result<(), Box<dyn Error + Send + Sync>> {
+) -> Result<RegisterRequestStatus, Box<dyn Error + Send + Sync>> {
     let body = json!({
         "token": token,
         "user": user_id,
@@ -36,16 +44,19 @@ async fn make_register_request(
         "username": username,
     });
 
-    reqwest::Client::new()
+    let result = reqwest::Client::new()
         .post(config::CONFIG.manager_url.clone())
         .body(body.to_string())
         .header("Authorization", config::CONFIG.manager_api_key.clone())
         .header("Content-Type", "application/json")
         .send()
-        .await?
-        .error_for_status()?;
+        .await?;
 
-    Ok(())
+    Ok(match result.status().as_u16() {
+        200 => RegisterRequestStatus::Success,
+        402 => RegisterRequestStatus::LimitExtended,
+        status_code => RegisterRequestStatus::UnknownError { status_code },
+    })
 }
 
 pub async fn register(user_id: UserId, message_text: &str) -> RegisterStatus {
@@ -58,13 +69,20 @@ pub async fn register(user_id: UserId, message_text: &str) -> RegisterStatus {
 
     let register_request_status = make_register_request(user_id, &bot_username, token).await;
 
-    if let Err(err) = register_request_status {
-        log::error!("Bot reg error: {:?}", err);
+    let result = match register_request_status {
+        Ok(v) => v,
+        Err(err) => {
+            log::error!("Bot reg error: {:?}", err);
 
-        return RegisterStatus::RegisterFail;
-    }
+            return RegisterStatus::RegisterFail;
+        }
+    };
 
-    RegisterStatus::Success {
-        username: bot_username,
+    match result {
+        RegisterRequestStatus::Success => RegisterStatus::Success {
+            username: bot_username,
+        },
+        RegisterRequestStatus::LimitExtended => RegisterStatus::LimitExtended,
+        RegisterRequestStatus::UnknownError { .. } => RegisterStatus::RegisterFail,
     }
 }
diff --git a/src/bots/bots_manager/strings.rs b/src/bots/bots_manager/strings.rs
index 92c62db..f8b542e 100644
--- a/src/bots/bots_manager/strings.rs
+++ b/src/bots/bots_manager/strings.rs
@@ -8,3 +8,5 @@ pub fn format_registered_message(username: &str) -> String {
 pub const ALREADY_REGISTERED: &str = "Ошибка! Возможно бот уже зарегистрирован!";
 
 pub const ERROR_MESSAGE: &str = "Ошибка! Что-то не так с ботом!";
+
+pub const LIMIT_EXTENDED_MESSAGE: &str = "Вы достигли максимального количества ботов!";