1  #[macro_use]
  2  extern crate rocket;
  3  
  4  use rocket::response::{Redirect, content::RawHtml};
  5  use rocket::serde::{Serialize, Deserialize};
  6  use rocket::State;
  7  use rocket::http::Status;
  8  use std::collections::{HashMap, HashSet};
  9  use std::sync::Arc;
 10  use tokio::sync::Mutex;
 11  use tokio::time::{interval, Duration};
 12  use std::time::{SystemTime, UNIX_EPOCH};
 13  use html_escape::encode_text;
 14  use tokio::time::sleep;
 15  use zeroize::Zeroize;
 16  use rocket::serde::json::Json;
 17  
 18  mod encryption;
 19  use crate::encryption::{encrypt_message, decrypt_message, is_message_encrypted};
 20  
 21  // Constants
 22  const TIME_WINDOW: u64 = 60;
 23  const MESSAGE_LIMIT: usize = 20; // 20 messages in 60 seconds
 24  const MAX_MESSAGE_LENGTH: usize = 2 * 1024 * 1024; // 2 megabytes
 25  const RECENT_MESSAGE_LIMIT: usize = 1000; // Maximum number of messages
 26  const MESSAGE_EXPIRY_DURATION: u64 = 86400; // 1 day
 27  
 28  #[derive(Debug, Clone, Serialize, Deserialize)]
 29  struct Message {
 30      content: String,
 31      timestamp: u64,
 32  }
 33  
 34  #[derive(Debug, Deserialize)]
 35  struct MessageData {
 36      message: String,
 37      room_id: String,
 38  }
 39  
 40  #[derive(Debug)]
 41  struct ChatState {
 42      messages: Arc<Mutex<Vec<Message>>>,
 43      user_request_timestamps: Arc<Mutex<HashMap<String, (u64, u64)>>>,
 44      recent_messages: Arc<Mutex<HashSet<String>>>,
 45      global_message_timestamps: Arc<Mutex<Vec<u64>>>,
 46  }
 47  
 48  // Manually implement Clone for ChatState with Arc
 49  impl Clone for ChatState {
 50      fn clone(&self) -> Self {
 51          ChatState {
 52              messages: Arc::clone(&self.messages),
 53              user_request_timestamps: Arc::clone(&self.user_request_timestamps),
 54              recent_messages: Arc::clone(&self.recent_messages),
 55              global_message_timestamps: Arc::clone(&self.global_message_timestamps),
 56          }
 57      }
 58  }
 59  
 60  // Helper function to format the timestamp into HH:MM:SS
 61  fn format_timestamp(timestamp: u64) -> String {
 62      let seconds = timestamp % 60;
 63      let minutes = (timestamp / 60) % 60;
 64      let hours = (timestamp / 3600) % 24;
 65      format!("{:02}:{:02}:{:02}", hours, minutes, seconds)
 66  }
 67  
 68  // Function to check if the message count exceeds the limit globally
 69  async fn check_message_limit(state: &ChatState) -> bool {
 70      let mut global_timestamps = state.global_message_timestamps.lock().await;
 71      let current_time = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs();
 72  
 73      // Remove messages older than the time window (60 seconds)
 74      global_timestamps.retain(|&timestamp| current_time - timestamp <= TIME_WINDOW);
 75  
 76      // Check if we have exceeded the message limit (20 messages in 60 seconds)
 77      if global_timestamps.len() >= MESSAGE_LIMIT {
 78          return false; // Exceeded the limit
 79      }
 80  
 81      // Record the current message timestamp
 82      global_timestamps.push(current_time);
 83      true
 84  }
 85  
 86  // Function to check if the message is valid (length and total message count)
 87  async fn is_message_valid(message: &str, state: &ChatState) -> bool {
 88      // Check if the message length exceeds the maximum limit
 89      if message.len() > MAX_MESSAGE_LENGTH {
 90          return false;
 91      }
 92  
 93      // Lock the messages state to access the total message count
 94      let mut messages = state.messages.lock().await;
 95  
 96      // Check if the total number of messages exceeds the limit
 97      if messages.len() >= RECENT_MESSAGE_LIMIT {
 98          // Wipe the content of the oldest message before removing it
 99          wipe_message_content(&mut messages[0]);
100          messages.remove(0); // Remove the first message in the vector (oldest)
101      }
102  
103      true
104  }
105  
106  #[get("/messages?<room_id>")]
107  async fn messages(room_id: Option<String>, state: &State<Arc<ChatState>>) -> String {
108      let chat_state = state.inner();
109      let messages = chat_state.messages.lock().await;
110  
111      let mut html = String::new();
112      for message in messages.iter() {
113          // Format timestamp for display
114          let timestamp = format_timestamp(message.timestamp);
115  
116          // Decrypt message content based on the provided room_id, if available
117          let decrypted_content = match &room_id {
118              Some(ref pw) => decrypt_message(&message.content, pw).unwrap_or_else(|_| {
119                  // Return nothing when decryption fails
120                  return String::new();  // Empty string will effectively remove the message from HTML
121              }),
122              None => String::new(),  // No room_id provided, return empty content
123          };
124  
125          // If decryption fails (decrypted_content is empty), skip appending the message
126          if decrypted_content.is_empty() {
127              continue;  // Skip this message entirely
128          }
129  
130          // Display the decrypted content, or nothing if decryption failed
131          html.push_str(&format!(
132              r#"<p>[{}]: {}</p>"#,
133              timestamp,
134              encode_text(&decrypted_content)
135          ));
136      }
137  
138      html
139  }
140  
141  #[get("/?<room_id>")]
142  async fn index(room_id: Option<String>, state: &State<Arc<ChatState>>) -> Result<RawHtml<String>, Status> {
143      // Read the static HTML template
144      let mut html = tokio::fs::read_to_string("static/index.html")
145          .await
146          .map_err(|_error| Status::InternalServerError)?;
147  
148      // Get room_id, defaulting to empty string if not provided
149      // Safely handle temporary value by assigning them to variable
150      let room_id_value = room_id.clone().unwrap_or_else(|| "".to_string());
151  
152      // Safely encode them using encode_text.
153      let encoded_room_id = encode_text(&room_id_value);
154  
155      // Replace placeholder with actual values
156      html = html.replace("room_id_PLACEHOLDER", &encoded_room_id);
157  
158      // Get current chat messages and generate HTML for them
159      let messages = state.messages.lock().await;
160      let mut messages_html = String::new();
161  
162      for msg in messages.iter() {
163          let timestamp = format_timestamp(msg.timestamp);
164  
165          // Decrypt message content based on provided room_id
166          let decrypted_content = if let Some(ref pw) = room_id {
167              decrypt_message(&msg.content, pw).unwrap_or_else(|_| "Decryption failed".to_string())
168          } else {
169              "room_id not provided".to_string()
170          };
171  
172          // Add the message to the HTML string
173          messages_html.push_str(&format!(
174              "<p>[{}]: {}</p>",
175              timestamp,
176              encode_text(&decrypted_content)
177          ));
178      }
179  
180      // Insert messages into the HTML template
181      html = html.replace("<!-- Messages will be dynamically inserted here -->", &messages_html);
182  
183      // Return the final HTML to the client
184      Ok(RawHtml(html))
185  }
186  
187  // Route for sending a message with encryption
188  #[post("/send", data = "<message_data>")]
189  async fn send(message_data: Json<MessageData>, state: &State<Arc<ChatState>>) -> Result<Redirect, RawHtml<String>> {
190      let message = message_data.message.trim();
191      let room_id = message_data.room_id.trim();
192  
193      // Delay message processing by 2 seconds
194      sleep(Duration::from_secs(2)).await;
195  
196      // Check if the message limit has been exceeded globally
197      if !check_message_limit(&state.inner()).await {
198          return Err(RawHtml("Too many messages sent in a short period. Please wait for 2 minutes.".to_string()));
199      }
200  
201      // Reject the request if the room room_id field is empty
202      if room_id.is_empty() {
203          return Err(RawHtml("Room room_id cannot be empty. Please provide a room_id.".to_string()));
204      }
205  
206      // Check if the room_id is at least 8 characters long
207      if room_id.len() < 8 {
208          return Err(RawHtml("Room room_id must be at least 8 characters long.".to_string()));
209      }
210  
211      // Check if the message is valid (length and total message count)
212      if !is_message_valid(message, state).await {
213          return Err(RawHtml("Invalid message. Make sure it's less than 10MB.".to_string()));
214      }
215  
216      // Check if the message is encrypted
217      if !is_message_encrypted(message) {
218          return Err(RawHtml("Message is not encrypted. Please provide an encrypted message.".to_string()));
219      }
220  
221      let mut messages = state.messages.lock().await;
222      let timestamp = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs();
223  
224      let encrypted_content = encrypt_message(message, room_id).map_err(|_| RawHtml("Encryption failed.".to_string()))?;
225  
226      messages.push(Message {
227          content: encrypted_content,
228          timestamp,
229      });
230  
231      Ok(Redirect::to(format!("/")))
232  }
233  
234  // Function to wipe message content securely
235  fn wipe_message_content(message: &mut Message) {
236      // Securely zero out the message content
237      message.content.zeroize();
238  }
239  
240  // Cleanup task to remove expired messages and securely wipe their contents
241  async fn message_cleanup_task(state: Arc<ChatState>) {
242      let mut interval = interval(Duration::from_secs(1)); // Check every second
243  
244      loop {
245          interval.tick().await; // Wait for the next tick of the interval
246  
247          let current_time = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs();
248  
249          // Acquire the lock on messages
250          let mut messages = state.messages.lock().await;
251  
252          // Check if there are messages that should be wiped
253          if let Some(oldest_message_index) = messages.iter().position(|message| {
254              current_time - message.timestamp >= MESSAGE_EXPIRY_DURATION
255          }) {
256              // Securely wipe the content of the oldest message
257              wipe_message_content(&mut messages[oldest_message_index]);
258  
259              // Remove the oldest message
260              messages.remove(oldest_message_index);
261          }
262      }
263  }
264  // Main function to launch the Rocket server
265  #[tokio::main]
266  async fn main() {
267      let chat_state = Arc::new(ChatState {
268          messages: Arc::new(Mutex::new(Vec::new())),
269          user_request_timestamps: Arc::new(Mutex::new(HashMap::new())),
270          recent_messages: Arc::new(Mutex::new(HashSet::new())),
271          global_message_timestamps: Arc::new(Mutex::new(Vec::new())),
272      });
273  
274      // Launch the message cleanup task
275      tokio::spawn(message_cleanup_task(Arc::clone(&chat_state)));
276  
277      // Launch the Rocket application
278      rocket::build()
279          .manage(chat_state)
280          .mount("/", routes![index, send, messages])
281          .mount("/static", rocket::fs::FileServer::from("static"))
282          .launch()
283          .await
284          .unwrap(); // Ensure the Rocket server is awaited and handle any errors
285  }