diff --git a/src/frontend/mod.rs b/src/frontend/mod.rs index 09febe6..e515fc1 100644 --- a/src/frontend/mod.rs +++ b/src/frontend/mod.rs @@ -128,10 +128,9 @@ impl App { .state .get() .unwrap() - .upload_requests - .lock() + .get_upload_request(id) .await - .get(&id).ok_or(LocalSendError::SessionInactive)?.clone(); + .ok_or(LocalSendError::SessionInactive)?; // TODO: replace this with ratatui widget dialog let upload_confirmed = MessageDialogBuilder::default() diff --git a/src/lib.rs b/src/lib.rs index b6ec164..4390da8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -59,7 +59,7 @@ pub struct JoecalState { pub running_state: Arc>, pub socket: Arc, pub client: reqwest::Client, - pub upload_requests: Arc>>>, + upload_requests: Arc>>>, shutdown_sender: OnceLock, // the receiving end will be held by the application so it can update the UI based on backend // events @@ -130,6 +130,9 @@ impl JoecalState { Listeners::Udp } }); + + // TODO: add a task that periodically clears out the upload requests if + // they're too old; the keys are julids so they have the time in them } pub async fn stop(&self) { @@ -147,6 +150,24 @@ impl JoecalState { let mut peers = self.peers.lock().await; peers.clear(); } + + pub async fn get_upload_request(&self, id: Julid) -> Option> { + self.upload_requests.lock().await.get(&id).cloned() + } + + pub async fn clear_upload_request(&self, id: Julid) { + let _ = self.upload_requests.lock().await.remove(&id); + } + + /// Add a transmitter for an upload request confirmation dialog that the + /// application frontend can use to tell the Axum handler whether or not to + /// accept the upload. + /// + /// IMPORTANT! Be sure to call `clear_upload_request(id)` when you're done + /// getting an answer back/before you exit! + pub async fn add_upload_request(&self, id: Julid, tx: UnboundedSender) { + self.upload_requests.lock().await.entry(id).insert_entry(tx); + } } #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/src/transfer.rs b/src/transfer.rs index c47b289..4f71ced 100644 --- a/src/transfer.rs +++ b/src/transfer.rs @@ -197,12 +197,8 @@ pub async fn register_prepare_upload( let id = Julid::new(); let (tx, mut rx) = unbounded_channel(); - state - .upload_requests - .lock() - .await - .entry(id) - .insert_entry(tx); + // be sure to clear this request before this function exits! + state.add_upload_request(id, tx).await; let dialog_send = state.transfer_event_tx.send(TransferEvent::UploadRequest { alias: req.info.alias.clone(), @@ -211,52 +207,55 @@ pub async fn register_prepare_upload( match dialog_send { Ok(_) => {} Err(_e) => { - let _ = state.upload_requests.lock().await.remove(&id); + state.clear_upload_request(id).await; return StatusCode::INTERNAL_SERVER_ERROR.into_response(); } } - // safe to unwrap because it's only `None` when there are no more transmitters, - // and we still have the `tx` we created earlier - let result = rx.recv().await.unwrap(); - let _ = state.upload_requests.lock().await.remove(&id); + let confirmation = rx.recv().await; + state.clear_upload_request(id).await; - if result == UploadDialog::UploadConfirm { - let session_id = id.as_string(); + let Some(confirmation) = confirmation else { + // the frontend must have dropped the tx before trying to send a reply back + return StatusCode::INTERNAL_SERVER_ERROR.into_response(); + }; - let file_tokens: HashMap = req - .files - .keys() - .map(|id| (id.clone(), Julid::new().to_string())) // Replace with actual token logic - .collect(); - - let session = Session { - session_id: session_id.clone(), - files: req.files.clone(), - file_tokens: file_tokens.clone(), - receiver: state.device.clone(), - sender: req.info.clone(), - status: SessionStatus::Active, - addr, - }; - - state - .sessions - .lock() - .await - .insert(session_id.clone(), session); - - ( - StatusCode::OK, - Json(PrepareUploadResponse { - session_id, - files: file_tokens, - }), - ) - .into_response() - } else { - StatusCode::FORBIDDEN.into_response() + if confirmation != UploadDialog::UploadConfirm { + return StatusCode::FORBIDDEN.into_response(); } + + let session_id = id.as_string(); + + let file_tokens: HashMap = req + .files + .keys() + .map(|id| (id.clone(), Julid::new().to_string())) // Replace with actual token logic + .collect(); + + let session = Session { + session_id: session_id.clone(), + files: req.files.clone(), + file_tokens: file_tokens.clone(), + receiver: state.device.clone(), + sender: req.info.clone(), + status: SessionStatus::Active, + addr, + }; + + state + .sessions + .lock() + .await + .insert(session_id.clone(), session); + + ( + StatusCode::OK, + Json(PrepareUploadResponse { + session_id, + files: file_tokens, + }), + ) + .into_response() } pub async fn register_upload(