diff --git a/src/mes/src/lib.rs b/src/mes/src/lib.rs index a66ed33..2cebc22 100644 --- a/src/mes/src/lib.rs +++ b/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, } 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 { + fn add_message(&mut self, message: Message) -> Result { + // 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 { + fn update_message(&mut self, id: u64, payload: MessagePayload) -> Result { 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 { 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,7 +113,6 @@ impl MessageStorage { } } -/// Payload struct for creating or updating a message. #[derive(CandidType, Serialize, Deserialize)] struct MessagePayload { title: String, @@ -95,44 +120,37 @@ struct MessagePayload { 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 { +fn add_message(message: MessagePayload) -> Result { 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 { STORAGE.with(|storage| storage.update_message(id, payload)) } -/// Delete a message by ID. #[update] fn delete_message(id: u64) -> Result { STORAGE.with(|storage| storage.delete_message(id)) } -/// Retrieve a message by ID. #[query] fn get_message(id: u64) -> Result { STORAGE.with(|storage| storage.get_message(id)).map(|m| m.clone()) } -// Expose the candid interface. ic_cdk::export_candid!();