1  use crate::encrypt_data;
  2  use crate::receive_and_fetch_messages;
  3  use crate::send_encrypted_message;
  4  use crate::pad_message;
  5  use rocket::{get, post, routes, serde::json::Json};
  6  use rocket::fs::NamedFile;
  7  use rocket::tokio;
  8  use rocket::fs::FileServer;
  9  use serde::{Deserialize, Serialize};
 10  use std::sync::{Arc, Mutex};
 11  use std::time::Duration;
 12  use std::path::PathBuf;
 13  
 14  #[derive(Clone)]
 15  pub struct MessagingApp {
 16      username: String,
 17      messages: Arc<Mutex<Vec<String>>>,
 18      shared_hybrid_secret: Arc<String>,
 19      shared_room_id: Arc<String>,
 20      shared_url: Arc<String>,
 21  }
 22  
 23  #[derive(Serialize, Deserialize)]
 24  struct MessageInput {
 25      message: String,
 26  }
 27  
 28  impl MessagingApp {
 29      pub fn new(
 30          username: String,
 31          shared_hybrid_secret: Arc<String>,
 32          shared_room_id: Arc<Mutex<String>>,
 33          shared_url: Arc<Mutex<String>>,
 34      ) -> Self {
 35          let messages = Arc::new(Mutex::new(vec![]));
 36          let messages_clone = Arc::clone(&messages);
 37          let shared_hybrid_secret_clone = Arc::clone(&shared_hybrid_secret);
 38          let shared_room_id_clone = Arc::clone(&shared_room_id);
 39          let shared_url_clone = Arc::clone(&shared_url);
 40  
 41          let room_id = Arc::new(shared_room_id_clone.lock().unwrap_or_else(|_| panic!("Failed to lock room_id")).clone());
 42          let url = Arc::new(shared_url_clone.lock().unwrap_or_else(|_| panic!("Failed to lock url")).clone());
 43  
 44          tokio::spawn(async move {
 45              loop {
 46                  let room_id_str = shared_room_id_clone.lock().unwrap_or_else(|_| panic!("Failed to lock room_id")).clone();
 47                  let url_str = shared_url_clone.lock().unwrap_or_else(|_| panic!("Failed to lock url")).clone();
 48  
 49                  match receive_and_fetch_messages(
 50                      &room_id_str,
 51                      &shared_hybrid_secret_clone,
 52                      &url_str,
 53                      true,
 54                  ) {
 55                      Ok(new_messages) => {
 56                          let mut msgs = messages_clone.lock().unwrap_or_else(|_| panic!("Failed to lock messages"));
 57                          msgs.clear();
 58                          msgs.extend(new_messages);
 59                      }
 60                      Err(e) => {
 61                          eprintln!("Error fetching messages: {}", e);
 62                      }
 63                  }
 64                  tokio::time::sleep(Duration::from_secs(10)).await;
 65              }
 66          });
 67  
 68          MessagingApp {
 69              username,
 70              messages,
 71              shared_hybrid_secret,
 72              shared_room_id: room_id,
 73              shared_url: url,
 74          }
 75      }
 76  }
 77  
 78  #[get("/messages")]
 79  async fn get_messages(app: &rocket::State<MessagingApp>) -> Json<Vec<String>> {
 80      let result = fetch_and_update_messages(&app).await;
 81      
 82      match result {
 83          Ok(msgs) => Json(msgs),
 84          Err(e) => {
 85              eprintln!("Error fetching messages: {}", e);
 86              
 87              // Return current messages if fetching fails
 88              let msgs = app.messages.lock().unwrap_or_else(|_| panic!("Failed to lock messages"));
 89              Json(msgs.clone())
 90          }
 91      }
 92  }
 93  
 94  async fn fetch_and_update_messages(app: &rocket::State<MessagingApp>) -> Result<Vec<String>, String> {
 95      let room_id_str = app.shared_room_id.clone();
 96      let url_str = app.shared_url.clone();
 97      
 98      let new_messages = tokio::task::block_in_place(move || {
 99          receive_and_fetch_messages(
100              &room_id_str,
101              &app.shared_hybrid_secret,
102              &url_str,
103              true,
104          )
105      }).map_err(|e| format!("Error fetching messages: {}", e))?;
106  
107      // Update the in-memory message storage
108      let mut msgs = app.messages.lock().unwrap_or_else(|_| panic!("Failed to lock messages"));
109      msgs.clear();
110      msgs.extend(new_messages.clone());
111  
112      Ok(new_messages)
113  }
114  
115  #[post("/send", data = "<input>")]
116  async fn post_message(
117      input: Json<MessageInput>,
118      app: &rocket::State<MessagingApp>
119  ) -> Result<&'static str, rocket::http::Status> {
120      // Create the formatted message once
121      let formatted_message = format!("<strong>{}</strong>: {}", app.username, input.message);
122      let padded_message = pad_message(&formatted_message, 2048);
123  
124      let result = tokio::task::block_in_place(|| {
125          // Encrypt the message
126          let encrypted = encrypt_data(&padded_message, &app.shared_hybrid_secret)
127              .map_err(|e| {
128                  eprintln!("Encryption error: {}", e);
129                  rocket::http::Status::InternalServerError
130              })?;
131          
132          // Send the encrypted message
133          send_encrypted_message(&encrypted, &app.shared_room_id, &app.shared_url)
134              .map_err(|e| {
135                  eprintln!("Error sending message: {}", e);
136                  rocket::http::Status::InternalServerError
137              })
138      });
139  
140      match result {
141          Ok(_) => {
142              {
143                  let mut msgs = app.messages.lock().unwrap_or_else(|_| panic!("Failed to lock messages"));
144                  msgs.push(formatted_message);
145              }
146              Ok("Message sent")
147          }
148          Err(e) => Err(e),
149      }
150  }
151  
152  #[get("/")]
153  async fn serve_webpage() -> Option<NamedFile> {
154      NamedFile::open(PathBuf::from("static/index.html")).await.ok()
155  }
156  
157  pub fn create_rocket(app: MessagingApp) -> rocket::Rocket<rocket::Build> {
158      rocket::build()
159          .manage(app)
160          .mount("/", routes![get_messages, post_message, serve_webpage])
161          .mount("/static", FileServer::from("static"))
162  }