Update lib.rs

src/mes/src/lib.rs

@@ -1,7 +1,9 @@
 use candid::{CandidType, Decode, Encode};
-use ic_cdk::api::time;
+use ic_cdk::{api::{time, logging, trap}, IDL};
 use std::collections::HashMap;
 
+const MAX_MESSAGE_SIZE: usize = 2048; // Example parameterized memory size
+
 #[derive(CandidType, Clone, Serialize, Deserialize)]
 struct Message {
     id: u64,
@@ -13,7 +15,6 @@ struct Message {
 }
 
 impl Message {
-    /// Create a new `Message` instance with the given parameters.
     fn new(id: u64, title: String, body: String, attachment_url: String) -> Self {
         Self {
             id,
@@ -26,32 +27,57 @@ impl Message {
     }
 }
 
-/// In-memory storage structure to manage messages.
 struct MessageStorage {
     messages: HashMap<u64, Message>,
 }
 
 impl MessageStorage {
-    /// Create a new instance of `MessageStorage`.
     fn new() -> Self {
         Self {
             messages: HashMap::new(),
         }
     }
 
-    /// Add a new message to the storage.
-    fn add_message(&mut self, message: Message) -> Option<Message> {
+    fn add_message(&mut self, message: Message) -> Result<Message, MessageError> {
+        // Input validation
+        if message.title.is_empty() || message.body.is_empty() || message.attachment_url.is_empty() {
+            return Err(MessageError::InvalidInput {
+                msg: "All fields must be non-empty".to_string(),
+            });
+        }
+
+        // Check size limitation
+        if message.title.len() + message.body.len() + message.attachment_url.len() > MAX_MESSAGE_SIZE {
+            return Err(MessageError::SizeExceeded);
+        }
+
+        // Logging
+        logging::info(&format!("Adding message with ID: {}", message.id));
+
+        // Storage
         self.messages.insert(message.id, message.clone());
-        Some(message)
+
+        Ok(message)
     }
 
-    /// Update an existing message.
-    fn update_message(
-        &mut self,
-        id: u64,
-        payload: MessagePayload,
-    ) -> Result<Message, MessageError> {
+    fn update_message(&mut self, id: u64, payload: MessagePayload) -> Result<Message, MessageError> {
         if let Some(message) = self.messages.get_mut(&id) {
+            // Input validation
+            if payload.title.is_empty() || payload.body.is_empty() || payload.attachment_url.is_empty() {
+                return Err(MessageError::InvalidInput {
+                    msg: "All fields must be non-empty".to_string(),
+                });
+            }
+
+            // Check size limitation
+            if payload.title.len() + payload.body.len() + payload.attachment_url.len() > MAX_MESSAGE_SIZE {
+                return Err(MessageError::SizeExceeded);
+            }
+
+            // Logging
+            logging::info(&format!("Updating message with ID: {}", id));
+
+            // Update
             message.attachment_url = payload.attachment_url;
             message.body = payload.body;
             message.title = payload.title;
@@ -64,9 +90,10 @@ impl MessageStorage {
         }
     }
 
-    /// Delete a message from the storage.
     fn delete_message(&mut self, id: u64) -> Result<Message, MessageError> {
         if let Some(message) = self.messages.remove(&id) {
+            // Logging
+            logging::info(&format!("Deleting message with ID: {}", id));
             Ok(message)
         } else {
             Err(MessageError::NotFound {
@@ -75,7 +102,6 @@ impl MessageStorage {
         }
     }
 
-    /// Retrieve a message by its ID.
     fn get_message(&self, id: u64) -> Result<&Message, MessageError> {
         if let Some(message) = self.messages.get(&id) {
             Ok(message)
@@ -87,52 +113,44 @@ impl MessageStorage {
     }
 }
 
-/// Payload struct for creating or updating a message.
 #[derive(CandidType, Serialize, Deserialize)]
 struct MessagePayload {
     title: String,
     body: String,
     attachment_url: String,
 }
 
-/// Custom error enum for message-related errors.
 #[derive(CandidType, Deserialize, Serialize)]
 enum MessageError {
     NotFound { msg: String },
+    InvalidInput { msg: String },
+    SizeExceeded,
 }
 
-// Instantiate the message storage globally.
 thread_local! {
     static STORAGE: MessageStorage = MessageStorage::new();
 }
 
-// Methods accessible to the canister
-
-/// Add a new message.
 #[update]
-fn add_message(message: MessagePayload) -> Option<Message> {
+fn add_message(message: MessagePayload) -> Result<Message, MessageError> {
     let id = time(); // Using time as an ID for simplicity.
     let new_message = Message::new(id, message.title, message.body, message.attachment_url);
     STORAGE.with(|storage| storage.add_message(new_message))
 }
 
-/// Update an existing message.
 #[update]
 fn update_message(id: u64, payload: MessagePayload) -> Result<Message, MessageError> {
     STORAGE.with(|storage| storage.update_message(id, payload))
 }
 
-/// Delete a message by ID.
 #[update]
 fn delete_message(id: u64) -> Result<Message, MessageError> {
     STORAGE.with(|storage| storage.delete_message(id))
 }
 
-/// Retrieve a message by ID.
 #[query]
 fn get_message(id: u64) -> Result<Message, MessageError> {
     STORAGE.with(|storage| storage.get_message(id)).map(|m| m.clone())
 }
 
-// Expose the candid interface.
 ic_cdk::export_candid!();