use send_state, fix sending state without tx ready, switch to u32 for download token

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]

src/models.rs

@@ -8,6 +8,7 @@ use crate::utility::search;
 
 #[derive(Debug, Serialize, Clone)]
 pub struct Session {
+    pub id: usize,
     pub downloads: Vec<SessionDownload>,
 
     pub first_seen: chrono::DateTime<chrono::Utc>,
@@ -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<chrono::Utc>,
     pub download_time: chrono::DateTime<chrono::Utc>,
@@ -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<ExecutableJson> },
 }