use salvo::{http::cookie::Cookie, websocket::Message, Response}; use serde::{Deserialize, Serialize}; use std::{collections::HashMap, path}; use tokio::sync::{mpsc::UnboundedSender, Mutex}; use crate::utility::search; #[derive(Debug, Serialize, Clone)] pub struct Session { pub id: u32, pub downloads: Vec, pub first_seen: chrono::DateTime, // The last time a request OR websocket message from/to this session was made pub last_seen: chrono::DateTime, // The last time a request was made with this session pub last_request: chrono::DateTime, // The sender for the websocket connection #[serde(skip_serializing)] pub tx: Option>>, } impl Session { // Update the last seen time(s) for the session pub fn seen(&mut self, socket: bool) { self.last_seen = chrono::Utc::now(); if !socket { self.last_request = chrono::Utc::now(); } } // Add a download to the session pub fn add_download(&mut self, exe: &Executable) -> &SessionDownload { let token: u32 = rand::random(); let download = SessionDownload { token, filename: format!( "{}-{:08x}{}{}", exe.name, token, if exe.extension.len() > 0 { "." } else { "" }, exe.extension ), last_used: chrono::Utc::now(), download_time: chrono::Utc::now(), }; self.downloads.push(download); return self.downloads.last().unwrap(); } // Delete a download from the session // Returns true if the download was deleted, false if it was not found pub fn delete_download(&mut self, token: u32) -> bool { if let Some(index) = self.downloads.iter().position(|d| d.token == token) { self.downloads.remove(index); true } else { tracing::warn!("Attempted to delete non-existent download token: {}", token); false } } // This function's failure is not a failure to transmit the message, but a failure to buffer it into the channel (or any preceding steps). pub fn send_message(&mut self, message: OutgoingMessage) -> Result<(), anyhow::Error> { if self.tx.is_none() { return Err(anyhow::anyhow!("Session {} has no sender", self.id)); } // TODO: Error handling let tx = self.tx.as_ref().unwrap(); let result = tx.send(Ok(Message::text(serde_json::to_string(&message).unwrap()))); match result { Ok(_) => return Ok(()), Err(e) => return Err(anyhow::anyhow!("Error sending message: {}", e)), } } pub fn send_state(&mut self) -> Result<(), anyhow::Error> { let message = OutgoingMessage::State { session: self.clone(), }; self.send_message(message) } } #[derive(Serialize, Debug, Clone)] pub struct SessionDownload { pub token: u32, pub filename: String, pub last_used: chrono::DateTime, pub download_time: chrono::DateTime, } impl SessionDownload {} #[derive(Clone, Debug)] pub struct State<'a> { // A map of executables, keyed by their type/platform pub executables: HashMap<&'a str, Executable>, // A map of sessions, keyed by their identifier (a random number) pub sessions: HashMap, // Provided on startup, the URL to the build log of the current deployment pub build_log: Option, } impl<'a> State<'a> { pub fn new() -> Mutex { Mutex::new(Self { build_log: None, executables: HashMap::new(), sessions: HashMap::new(), }) } pub fn add_executable(&mut self, exe_type: &'a str, exe_path: &str) { let data = std::fs::read(&exe_path).expect("Unable to read file"); let pattern = "a".repeat(1024); let key_start = search(&data, pattern.as_bytes(), 0).unwrap(); let key_end = key_start + pattern.len(); let path = path::Path::new(&exe_path); let name = path.file_stem().unwrap().to_str().unwrap(); let extension = match path.extension() { Some(s) => s.to_str().unwrap(), None => "", }; let exe = Executable { data, filename: path.file_name().unwrap().to_str().unwrap().to_string(), name: name.to_string(), extension: extension.to_string(), key_start: key_start, key_end: key_end, }; self.executables.insert(exe_type, exe); } pub async fn new_session(&mut self, res: &mut Response) -> u32 { let id: u32 = rand::random(); let now = chrono::Utc::now(); self.sessions.insert( id, Session { id, downloads: Vec::new(), last_seen: now, last_request: now, first_seen: now, tx: None, }, ); tracing::info!("New session created: {}", id); res.add_cookie( Cookie::build(("Session", id.to_string())) .http_only(true) .partitioned(true) .secure(cfg!(debug_assertions) == false) .path("/") // Use SameSite=None only in development .same_site(if cfg!(debug_assertions) { salvo::http::cookie::SameSite::None } else { salvo::http::cookie::SameSite::Strict }) .permanent() .build(), ); return id; } pub fn executable_json(&self) -> Vec { let mut executables = Vec::new(); for (key, exe) in &self.executables { executables.push(ExecutableJson { id: key.to_string(), size: exe.data.len(), filename: exe.filename.clone(), }); } return executables; } } #[derive(Default, Clone, Debug)] pub struct Executable { pub data: Vec, // the raw data of the executable pub filename: String, pub name: String, // the name before the extension pub extension: String, // may be empty string pub key_start: usize, // the index of the byte where the key starts pub key_end: usize, // the index of the byte where the key ends } impl Executable { pub fn with_key(&self, new_key: &[u8]) -> Vec { let mut data = self.data.clone(); // Copy the key into the data for i in 0..new_key.len() { data[self.key_start + i] = new_key[i]; } // If the new key is shorter than the old key, we just write over the remaining data if new_key.len() < self.key_end - self.key_start { for i in self.key_start + new_key.len()..self.key_end { data[i] = b' '; } } return data; } } #[derive(Debug, Deserialize)] #[serde(tag = "type", rename_all = "kebab-case")] pub enum IncomingMessage { // A request from the client to delete a download token DeleteDownloadToken { id: u32 }, } #[derive(Debug, Serialize)] #[serde(tag = "type", rename_all = "kebab-case")] pub enum OutgoingMessage { // An alert to the client that a session download has been used. #[serde(rename = "notify")] TokenAlert { token: u32, }, // A message describing the current session state State { session: Session, }, Executables { build_log: Option, executables: Vec, }, } #[derive(Debug, Serialize)] pub struct ExecutableJson { pub id: String, pub size: usize, pub filename: String, }