Add update_bot endpoint only callable by bot owner (#7073)

open-chat-labs · Dec 17, 2024 · 19b351b · 19b351b
1 parent bcada18
commit 19b351b
Show file tree

Hide file tree

Showing 14 changed files with 285 additions and 16 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/backend/canisters/user_index/CHANGELOG.md b/backend/canisters/user_index/CHANGELOG.md
@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ### Added
 
 - Add new LocalUserIndexes to the CyclesDispenser's allow list ([#7070](https://github.com/open-chat-labs/open-chat/pull/7070))
+- Add `update_bot` endpoint only callable by bot owner ([#7073](https://github.com/open-chat-labs/open-chat/pull/7073))
 
 ### Removed
 

diff --git a/backend/canisters/user_index/api/src/updates/mod.rs b/backend/canisters/user_index/api/src/updates/mod.rs
@@ -30,6 +30,7 @@ pub mod set_username;
 pub mod submit_proof_of_unique_personhood;
 pub mod suspend_user;
 pub mod unsuspend_user;
+pub mod update_bot;
 pub mod update_diamond_membership_subscription;
 pub mod upgrade_local_user_index_canister_wasm;
 pub mod upgrade_user_canister_wasm;

diff --git a/backend/canisters/user_index/api/src/updates/update_bot.rs b/backend/canisters/user_index/api/src/updates/update_bot.rs
@@ -0,0 +1,36 @@
+use candid::CandidType;
+use serde::{Deserialize, Serialize};
+use ts_export::ts_export;
+use types::{OptionUpdate, UserId};
+
+#[ts_export(user_index, update_bot)]
+#[derive(CandidType, Serialize, Deserialize, Debug)]
+pub struct Args {
+    pub bot_id: UserId,
+    pub owner: Option<UserId>,
+    pub name: Option<String>,
+    #[ts(as = "types::OptionUpdateString")]
+    pub avatar: OptionUpdate<String>, // Image as a data URL
+    pub endpoint: Option<String>,
+}
+
+#[ts_export(user_index, update_bot)]
+#[derive(CandidType, Serialize, Deserialize, Debug)]
+pub enum Response {
+    Success,
+    NameInvalid,
+    NameAlreadyExists,
+    AvatarInvalid,
+    EndpointInvalid,
+    BotNotFound,
+    BotSuspended,
+    NotAuthorised,
+    OwnerNotFound,
+    OwnerSuspended,
+    NewOwnerNotFound,
+    NewOwnerSuspended,
+    DefinitionNotFound,
+    DefinitionInvalid,
+    DescriptionTooLong,
+    TooManyCommands,
+}
diff --git a/backend/canisters/user_index/impl/Cargo.toml b/backend/canisters/user_index/impl/Cargo.toml
@@ -31,7 +31,7 @@ group_canister_c2c_client = { path = "../../group/c2c_client" }
 hex = { workspace = true }
 http_request = { path = "../../../libraries/http_request" }
 human_readable = { path = "../../../libraries/human_readable" }
-ic-cdk = { workspace = true }
+ic-cdk = { workspace = true, features = ["transform-closure"] }
 ic-cdk-timers = { workspace = true }
 ic-ledger-types = { workspace = true }
 ic-stable-structures = { workspace = true }

diff --git a/backend/canisters/user_index/impl/src/lifecycle/inspect_message.rs b/backend/canisters/user_index/impl/src/lifecycle/inspect_message.rs
@@ -20,6 +20,7 @@ fn accept_if_valid(state: &RuntimeState) {
         | "set_moderation_flags"
         | "set_username"
         | "submit_proof_of_unique_personhood"
+        | "update_bot"
         | "update_diamond_membership_subscription" => {
             let caller = state.env.caller();
             let is_user = state.data.users.get_by_principal(&caller).is_some();

diff --git a/backend/canisters/user_index/impl/src/model/mod.rs b/backend/canisters/user_index/impl/src/model/mod.rs
@@ -10,3 +10,7 @@ pub mod storage_index_user_config_batch;
 pub mod streak_insurance_logs;
 pub mod user;
 pub mod user_map;
+
+pub const MAX_AVATAR_SIZE: usize = 250_000;
+pub const MAX_DESCRIPTION_LEN: usize = 10_000;
+pub const MAX_COMMANDS: usize = 100;
diff --git a/backend/canisters/user_index/impl/src/model/user_map.rs b/backend/canisters/user_index/impl/src/model/user_map.rs
@@ -40,7 +40,7 @@ pub struct UserMap {
     pub users_with_duplicate_principals: Vec<(UserId, UserId)>,
 }
 
-#[derive(Serialize, Deserialize)]
+#[derive(Serialize, Deserialize, Clone)]
 pub struct Bot {
     pub name: String,
     pub avatar: Option<Document>,
@@ -131,7 +131,13 @@ impl UserMap {
         }
     }
 
-    pub fn update(&mut self, mut user: User, now: TimestampMillis, ignore_principal_clash: bool) -> UpdateUserResult {
+    pub fn update(
+        &mut self,
+        mut user: User,
+        now: TimestampMillis,
+        ignore_principal_clash: bool,
+        bot: Option<Bot>,
+    ) -> UpdateUserResult {
         let user_id = user.user_id;
 
         if let Some(previous) = self.users.get(&user_id) {
@@ -152,7 +158,7 @@ impl UserMap {
                 }
             }
 
-            if username_case_insensitive_changed && self.does_username_exist(username, false) {
+            if username_case_insensitive_changed && self.does_username_exist(username, bot.is_some()) {
                 return UpdateUserResult::UsernameTaken;
             }
 
@@ -166,15 +172,25 @@ impl UserMap {
             }
 
             if username_case_insensitive_changed {
-                self.username_to_user_id.remove(previous_username);
-                self.username_to_user_id.insert(username, user_id);
+                if bot.is_some() {
+                    self.botname_to_user_id.remove(previous_username);
+                    self.botname_to_user_id.insert(username, user_id);
+                } else {
+                    self.username_to_user_id.remove(previous_username);
+                    self.username_to_user_id.insert(username, user_id);
+                }
             }
 
             if previous.display_name != user.display_name {
                 user.display_name_upper = user.display_name.as_ref().map(|s| s.to_uppercase());
             }
 
             self.users.insert(user_id, user);
+
+            if let Some(bot) = bot {
+                self.bots.insert(user_id, bot);
+            }
+
             UpdateUserResult::Success
         } else {
             UpdateUserResult::UserNotFound
@@ -504,7 +520,7 @@ impl UserMap {
             UserType::User,
             None,
         );
-        self.update(user, date_created, false);
+        self.update(user, date_created, false, None);
     }
 }
 
@@ -644,7 +660,7 @@ mod tests {
             let mut updated = original.clone();
             updated.username.clone_from(&username2);
 
-            assert!(matches!(user_map.update(updated, 3, false), UpdateUserResult::Success));
+            assert!(matches!(user_map.update(updated, 3, false, None), UpdateUserResult::Success));
 
             assert_eq!(user_map.users.keys().collect_vec(), vec!(&user_id));
             assert_eq!(user_map.username_to_user_id.len(), 1);
@@ -688,7 +704,10 @@ mod tests {
 
         user_map.add_test_user(original);
         user_map.add_test_user(other);
-        assert!(matches!(user_map.update(updated, 3, false), UpdateUserResult::UsernameTaken));
+        assert!(matches!(
+            user_map.update(updated, 3, false, None),
+            UpdateUserResult::UsernameTaken
+        ));
     }
 
     #[test]
@@ -712,6 +731,6 @@ mod tests {
         user_map.add_test_user(original);
         updated.username = "ABC".to_string();
 
-        assert!(matches!(user_map.update(updated, 2, false), UpdateUserResult::Success));
+        assert!(matches!(user_map.update(updated, 2, false, None), UpdateUserResult::Success));
     }
 }
diff --git a/backend/canisters/user_index/impl/src/updates/mod.rs b/backend/canisters/user_index/impl/src/updates/mod.rs
@@ -29,6 +29,7 @@ pub mod set_username;
 pub mod submit_proof_of_unique_personhood;
 pub mod suspend_user;
 pub mod unsuspend_user;
+pub mod update_bot;
 pub mod update_diamond_membership_subscription;
 pub mod upgrade_local_user_index_canister_wasm;
 pub mod upgrade_user_canister_wasm;

diff --git a/backend/canisters/user_index/impl/src/updates/register_bot.rs b/backend/canisters/user_index/impl/src/updates/register_bot.rs
@@ -1,5 +1,6 @@
 use crate::guards::caller_is_governance_principal;
 use crate::model::user_map::Bot;
+use crate::model::{MAX_AVATAR_SIZE, MAX_COMMANDS, MAX_DESCRIPTION_LEN};
 use crate::{mutate_state, read_state, RuntimeState, USER_LIMIT};
 use candid::Principal;
 use canister_api_macros::{proposal, update};
@@ -14,10 +15,6 @@ use user_index_canister::register_bot::{Response::*, *};
 use utils::document::try_parse_data_url;
 use utils::text_validation::{validate_username, UsernameValidationError};
 
-const MAX_AVATAR_SIZE: usize = 250_000;
-const MAX_DESCRIPTION_LEN: usize = 10_000;
-const MAX_COMMANDS: usize = 100;
-
 #[update(msgpack = true)]
 #[trace]
 fn register_bot(args: Args) -> Response {

diff --git a/backend/canisters/user_index/impl/src/updates/set_display_name.rs b/backend/canisters/user_index/impl/src/updates/set_display_name.rs
@@ -29,7 +29,7 @@ fn set_display_name_impl(args: Args, state: &mut RuntimeState) -> Response {
         let mut user_to_update = user.clone();
         user_to_update.display_name.clone_from(&args.display_name);
         let user_id = user.user_id;
-        match state.data.users.update(user_to_update, now, false) {
+        match state.data.users.update(user_to_update, now, false, None) {
             UpdateUserResult::Success => {
                 state.push_event_to_local_user_index(
                     user_id,

diff --git a/backend/canisters/user_index/impl/src/updates/set_username.rs b/backend/canisters/user_index/impl/src/updates/set_username.rs
@@ -31,7 +31,7 @@ fn set_username_impl(args: Args, state: &mut RuntimeState) -> Response {
         user_to_update.username.clone_from(&username);
         let user_id = user.user_id;
         let now = state.env.now();
-        match state.data.users.update(user_to_update, now, false) {
+        match state.data.users.update(user_to_update, now, false, None) {
             UpdateUserResult::Success => {
                 state.push_event_to_local_user_index(
                     user_id,