diff --git a/src/main.rs b/src/main.rs index f1e1109..c735464 100644 --- a/src/main.rs +++ b/src/main.rs @@ -102,10 +102,7 @@ async fn handle_socket(session_id: usize, websocket: WebSocket) { .expect("Unable to get session"); session.tx = Some(tx_channel); - session.send_message(OutgoingMessage::State { - id: session_id, - session: session.clone(), - }); + session.send_state(); session.send_message(executable_message); // Handle incoming messages @@ -193,6 +190,13 @@ pub async fn download(req: &mut Request, res: &mut Response, depot: &mut Depot) "Content-Type", HeaderValue::from_static("application/octet-stream"), ); + + // Don't try to send state if somehow the session has not connected + if session.tx.is_some() { + session.send_state(); + } else { + tracing::warn!("Download being made without any connection websocket"); + } } #[handler] diff --git a/src/models.rs b/src/models.rs index ef78308..519c7d2 100644 --- a/src/models.rs +++ b/src/models.rs @@ -8,6 +8,7 @@ use crate::utility::search; #[derive(Debug, Serialize, Clone)] pub struct Session { + pub id: usize, pub downloads: Vec, pub first_seen: chrono::DateTime, @@ -33,7 +34,7 @@ impl Session { // Add a download to the session pub fn add_download(&mut self, exe: &Executable) -> &SessionDownload { let mut rng = rand::thread_rng(); - let token: u64 = rng.gen(); + let token: u32 = rand::random(); let download = SessionDownload { token, @@ -47,6 +48,10 @@ impl Session { } pub fn send_message(&mut self, message: OutgoingMessage) { + if self.tx.is_none() { + return; + } + // TODO: Error handling, check tx exists let result = self @@ -59,11 +64,19 @@ impl Session { tracing::error!("Failed to initial session state: {}", e); } } + + pub fn send_state(&mut self) { + let message = OutgoingMessage::State { + session: self.clone(), + }; + + self.send_message(message); + } } #[derive(Serialize, Debug, Clone)] pub struct SessionDownload { - pub token: u64, + pub token: u32, pub filename: String, pub last_used: chrono::DateTime, pub download_time: chrono::DateTime, @@ -121,6 +134,7 @@ impl<'a> State<'a> { self.sessions.insert( id, Session { + id, downloads: Vec::new(), last_seen: now, last_request: now, @@ -201,7 +215,7 @@ pub enum OutgoingMessage { // An alert to the client that a session download has been used. TokenAlert { token: u64 }, // A message describing the current session state - State { session: Session, id: usize }, + State { session: Session }, Executables { executables: Vec }, }