feat(bsk): harden browser file transfer transactions

This commit is contained in:
Ljy-0827
2026-08-24 15:02:52 +08:00
parent a86cbf724f
commit bf59753aa3
24 changed files with 1397 additions and 407 deletions
+75
View File
@@ -0,0 +1,75 @@
//! Atomic visibility boundary for files received into a same-directory temp.
use std::path::Path;
pub fn commit(temp: &Path, out: &Path, overwrite: bool) -> std::io::Result<()> {
if !overwrite {
std::fs::hard_link(temp, out)?;
std::fs::remove_file(temp)?;
return Ok(());
}
replace(temp, out)
}
#[cfg(unix)]
fn replace(temp: &Path, out: &Path) -> std::io::Result<()> {
// POSIX rename replaces an existing non-directory destination atomically.
std::fs::rename(temp, out)
}
#[cfg(windows)]
fn replace(temp: &Path, out: &Path) -> std::io::Result<()> {
use std::os::windows::ffi::OsStrExt;
use windows_sys::Win32::Storage::FileSystem::{
MOVEFILE_REPLACE_EXISTING, MOVEFILE_WRITE_THROUGH, MoveFileExW,
};
let from: Vec<u16> = temp.as_os_str().encode_wide().chain(Some(0)).collect();
let to: Vec<u16> = out.as_os_str().encode_wide().chain(Some(0)).collect();
// SAFETY: both buffers are NUL-terminated and remain alive for the call.
let ok = unsafe {
MoveFileExW(
from.as_ptr(),
to.as_ptr(),
MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH,
)
};
if ok == 0 {
Err(std::io::Error::last_os_error())
} else {
Ok(())
}
}
#[cfg(not(any(unix, windows)))]
fn replace(temp: &Path, out: &Path) -> std::io::Result<()> {
std::fs::rename(temp, out)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn no_replace_never_overwrites_existing_output() {
let dir = tempfile::tempdir().unwrap();
let temp = dir.path().join("new.part");
let out = dir.path().join("out.bin");
std::fs::write(&temp, b"new").unwrap();
std::fs::write(&out, b"old").unwrap();
assert!(commit(&temp, &out, false).is_err());
assert_eq!(std::fs::read(&out).unwrap(), b"old");
}
#[test]
fn overwrite_replaces_existing_output_without_predelete() {
let dir = tempfile::tempdir().unwrap();
let temp = dir.path().join("new.part");
let out = dir.path().join("out.bin");
std::fs::write(&temp, b"new").unwrap();
std::fs::write(&out, b"old").unwrap();
commit(&temp, &out, true).unwrap();
assert_eq!(std::fs::read(&out).unwrap(), b"new");
assert!(!temp.exists());
}
}
+4 -16
View File
@@ -15,6 +15,7 @@ use bsk_protocol::tools::{
use clap::Args;
use uuid::Uuid;
use crate::cli::atomic_output;
use crate::cli::ensure_daemon::ensure_daemon;
use crate::cli::error::{CliError, Format};
use crate::cli::interaction::split_target;
@@ -56,6 +57,7 @@ pub fn dispatch(args: DownloadArgs, format: Format) -> Result<(), CliError> {
tab_id: args.tab_id,
timeout_ms: Some(args.timeout),
browser_relative_dir: None,
max_byte_size: None,
};
let reply: DownloadResult = crate::cli::business_rpc::call(
info.sock_path.clone(),
@@ -130,22 +132,8 @@ fn write_transfer(sock: &Path, id: &str, out: &Path, overwrite: bool) -> Result<
}
file.sync_all().map_err(|e| CliError::Local(e.into()))?;
drop(file);
if !overwrite {
std::fs::hard_link(&temp, out)
.with_context(|| format!("commit download without overwriting {}", out.display()))
.map_err(CliError::Local)?;
std::fs::remove_file(&temp)
.with_context(|| format!("remove temporary download {}", temp.display()))
.map_err(CliError::Local)?;
return Ok(());
}
if out.exists() {
std::fs::remove_file(out)
.with_context(|| format!("replace existing output {}", out.display()))
.map_err(CliError::Local)?;
}
std::fs::rename(&temp, out)
.with_context(|| format!("commit download to {}", out.display()))
atomic_output::commit(&temp, out, overwrite)
.with_context(|| format!("atomically commit download to {}", out.display()))
.map_err(CliError::Local)
})();
if result.is_err() {
+1
View File
@@ -2,6 +2,7 @@
use std::time::Duration;
mod atomic_output;
pub mod browser_wait;
pub mod browsers;
pub mod business_rpc;
+25
View File
@@ -53,6 +53,8 @@ pub mod reason {
pub const FILE_INPUT_NOT_ACTIVATED: &str = "file_input_not_activated";
pub const SET_FILE_INPUT_FAILED: &str = "set_file_input_failed";
pub const DOWNLOAD_CAPTURE_FAILED: &str = "download_capture_failed";
pub const TRANSFER_OUTCOME_UNKNOWN: &str = "transfer_outcome_unknown";
pub const TRANSFER_TIMEOUT: &str = "transfer_timeout";
pub const SESSION_BUSY: &str = crate::rpc_reason::SESSION_BUSY;
pub const RECORD_START_PAGE_UNREACHABLE: &str = "record_start_page_unreachable";
}
@@ -201,6 +203,20 @@ pub fn info_for_error(code: ErrorCode, data: Option<&serde_json::Value>) -> Rend
return base;
};
match (code, reason) {
(_, reason::TRANSFER_OUTCOME_UNKNOWN) => RenderInfo {
summary: "the file transfer outcome could not be confirmed",
hint: Some(
"do not retry the transfer: the browser may already have applied it; inspect the page or stop safely",
),
exit_code: base.exit_code,
},
(ErrorCode::Timeout, reason::TRANSFER_TIMEOUT) => RenderInfo {
summary: "the file transfer timed out after browser dispatch",
hint: Some(
"do not retry when effect_state is unknown or committed; inspect the page before taking another action",
),
exit_code: base.exit_code,
},
(ErrorCode::PermissionDenied, reason::ELEMENT_NOT_VISIBLE) => RenderInfo {
summary: "target element has no visible geometry",
hint: Some(
@@ -474,6 +490,15 @@ mod tests {
let download = serde_json::json!({ "reason": reason::DOWNLOAD_CAPTURE_FAILED });
let info = info_for_error(ErrorCode::CdpFailed, Some(&download));
assert!(info.summary.contains("download could not be attributed"));
let unknown = serde_json::json!({ "reason": reason::TRANSFER_OUTCOME_UNKNOWN });
let info = info_for_error(ErrorCode::ProtocolError, Some(&unknown));
assert!(info.summary.contains("outcome could not be confirmed"));
assert!(info.hint.unwrap().contains("do not retry"));
let timeout = serde_json::json!({ "reason": reason::TRANSFER_TIMEOUT });
let info = info_for_error(ErrorCode::Timeout, Some(&timeout));
assert!(info.summary.contains("timed out after browser dispatch"));
}
#[test]
+35 -30
View File
@@ -49,38 +49,43 @@ pub fn dispatch(args: UploadArgs, format: Format) -> Result<(), CliError> {
let info = ensure_daemon().context("ensure daemon is running")?;
let (ref_, selector) = split_target(args.target, args.ref_, args.selector)?;
let mut staged = Vec::new();
let result = (|| {
for path in &args.files {
staged.push(stage_file(&info.sock_path, &args.session, path)?);
}
let params = UploadParams {
session_id: args.session,
ref_,
selector,
tab_id: args.tab_id,
files: staged
.iter()
.map(|(id, name)| UploadFile {
transfer_id: id.clone(),
name: name.clone(),
staged_path: None,
})
.collect(),
timeout_ms: Some(args.timeout),
};
crate::cli::business_rpc::call::<_, UploadResult>(
info.sock_path.clone(),
"upload",
Method::ToolUpload,
Some(params),
ipc_timeout(args.timeout),
)
})();
if result.is_err() {
for (id, _) in &staged {
let _ = release(&info.sock_path, id);
for path in &args.files {
match stage_file(&info.sock_path, &args.session, path) {
Ok(file) => staged.push(file),
Err(err) => {
for (id, _) in &staged {
let _ = release(&info.sock_path, id);
}
return Err(err);
}
}
}
let params = UploadParams {
session_id: args.session,
ref_,
selector,
tab_id: args.tab_id,
files: staged
.iter()
.map(|(id, name)| UploadFile {
transfer_id: id.clone(),
name: name.clone(),
staged_path: None,
})
.collect(),
timeout_ms: Some(args.timeout),
};
let result = crate::cli::business_rpc::call::<_, UploadResult>(
info.sock_path.clone(),
"upload",
Method::ToolUpload,
Some(params),
ipc_timeout(args.timeout),
);
// Once tool.upload is dispatched, staging ownership belongs to the
// session. A transport timeout cannot prove that Chrome did not attach
// the file, so releasing here could invalidate a late successful attach.
// Session teardown remains the single cleanup boundary after dispatch.
let reply = result?;
match format {
Format::Json => println!("{}", serde_json::to_string_pretty(&reply).unwrap()),
+108 -25
View File
@@ -40,6 +40,67 @@ struct Entry {
ready: bool,
}
/// Browser download path after the daemon has verified that it belongs to
/// the capability minted for this transfer. From this point onward the daemon
/// owns cleanup on every error path; unvalidated arbitrary paths are never
/// removed.
#[derive(Debug)]
struct ValidatedBrowserFile {
path: PathBuf,
parent: PathBuf,
cleanup_armed: bool,
}
impl ValidatedBrowserFile {
fn validate(reported_path: &Path, transfer_id: &str) -> Result<Self, RpcError> {
if !reported_path.is_absolute() {
return Err(permission(
"download path is outside its browser capability",
));
}
let expected_parent = Path::new("BrowserSkill").join(transfer_id);
let parent = reported_path
.parent()
.ok_or_else(|| permission("download path has no parent directory"))?;
if !parent.ends_with(&expected_parent) {
return Err(permission(
"download escaped its browser capability directory",
));
}
for path in [reported_path, parent] {
if fs::symlink_metadata(path)
.map_err(io_error)?
.file_type()
.is_symlink()
{
return Err(permission("download capability path contains a symlink"));
}
}
Ok(Self {
path: reported_path.to_path_buf(),
parent: parent.to_path_buf(),
cleanup_armed: true,
})
}
fn remove_source(mut self) -> Result<(), RpcError> {
remove_if_present(&self.path).map_err(io_error)?;
remove_dir_if_present(&self.parent).map_err(io_error)?;
self.cleanup_armed = false;
Ok(())
}
}
impl Drop for ValidatedBrowserFile {
fn drop(&mut self) {
if !self.cleanup_armed {
return;
}
let _ = remove_if_present(&self.path);
let _ = remove_dir_if_present(&self.parent);
}
}
#[derive(Debug)]
pub struct DownloadStaging {
pub transfer_id: String,
@@ -259,30 +320,11 @@ impl TransferRegistry {
let entry = entries
.get_mut(id)
.ok_or_else(|| not_found("download transfer not found"))?;
if entry.direction != Direction::Download || !reported_path.is_absolute() {
return Err(permission(
"download path is outside its browser capability",
));
if entry.direction != Direction::Download {
return Err(permission("transfer is not a download capability"));
}
let expected_parent = Path::new("BrowserSkill").join(id);
let reported_parent = reported_path
.parent()
.ok_or_else(|| permission("download path has no parent directory"))?;
if !reported_parent.ends_with(&expected_parent) {
return Err(permission(
"download escaped its browser capability directory",
));
}
for path in [reported_path, reported_parent] {
if fs::symlink_metadata(path)
.map_err(io_error)?
.file_type()
.is_symlink()
{
return Err(permission("download capability path contains a symlink"));
}
}
let canonical = reported_path.canonicalize().map_err(io_error)?;
let browser_file = ValidatedBrowserFile::validate(reported_path, id)?;
let canonical = browser_file.path.canonicalize().map_err(io_error)?;
let meta = fs::metadata(&canonical).map_err(io_error)?;
if !meta.is_file() || meta.len() > MAX_TRANSFER_BYTES {
return Err(invalid(
@@ -310,8 +352,7 @@ impl TransferRegistry {
entry.written = copied;
entry.expected_size = Some(copied);
entry.ready = true;
let _ = fs::remove_file(&canonical);
let _ = fs::remove_dir(reported_parent);
browser_file.remove_source()?;
Ok(copied)
}
@@ -391,6 +432,22 @@ fn set_private_file(file: &File) -> std::io::Result<()> {
Ok(())
}
fn remove_if_present(path: &Path) -> std::io::Result<()> {
match fs::remove_file(path) {
Ok(()) => Ok(()),
Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(err) => Err(err),
}
}
fn remove_dir_if_present(path: &Path) -> std::io::Result<()> {
match fs::remove_dir(path) {
Ok(()) => Ok(()),
Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(err) => Err(err),
}
}
fn safe_upload_basename(name: &str) -> Result<&str, RpcError> {
if name.is_empty() || name.contains(['/', '\\', '\0']) {
return Err(invalid("upload name must be a safe basename"));
@@ -570,6 +627,7 @@ mod tests {
.import_download(&staging.transfer_id, &outside)
.is_err()
);
assert_eq!(fs::read(&outside).unwrap(), b"secret");
let browser_dir = temp
.path()
@@ -601,6 +659,31 @@ mod tests {
);
}
#[test]
fn validated_oversized_browser_download_is_removed_on_rejection() {
let (temp, registry) = registry();
let staging = registry.begin_download("s1").unwrap();
let browser_dir = temp
.path()
.join("Downloads")
.join("BrowserSkill")
.join(&staging.transfer_id);
fs::create_dir_all(&browser_dir).unwrap();
let inside = browser_dir.join("oversized.bin");
File::create(&inside)
.unwrap()
.set_len(MAX_TRANSFER_BYTES + 1)
.unwrap();
assert!(
registry
.import_download(&staging.transfer_id, &inside)
.is_err()
);
assert!(!inside.exists());
assert!(!browser_dir.exists());
}
#[test]
fn releasing_a_session_removes_all_staging() {
let (_temp, registry) = registry();
+2 -1
View File
@@ -382,6 +382,7 @@ async fn handle_tool_dispatch(
Err(err) => return ResponseBody::Err(err),
};
download.browser_relative_dir = Some(staging.browser_relative_dir);
download.max_byte_size = Some(super::file_transfer::MAX_TRANSFER_BYTES);
download_transfer_id = Some(staging.transfer_id);
params = serde_json::to_value(download).unwrap_or(Value::Null);
}
@@ -702,7 +703,7 @@ fn tool_dispatch_timeout(params: &Value) -> Result<Duration, RpcError> {
fn tool_dispatch_transport_timeout(method: &Method, params: &Value) -> Result<Duration, RpcError> {
tool_dispatch_timeout(params).map(|timeout| {
if method == &Method::ToolUpload {
if matches!(method, Method::ToolUpload | Method::ToolDownload) {
timeout.saturating_add(EXTENSION_RESPONSE_GRACE)
} else {
timeout
+226 -27
View File
@@ -438,6 +438,7 @@ async fn dispatch_with_sender(
lifecycle_cancel: Option<AbortToken>,
cancel_cleanup_timeout: Duration,
) -> Result<Value, DispatchError> {
let effect_aware_transfer = is_effect_aware_transfer(&method);
let (respond_tx, respond_rx) = oneshot::channel();
let job = ToolJob {
method,
@@ -474,15 +475,21 @@ async fn dispatch_with_sender(
mpsc::error::TrySendError::Closed(_) => DispatchError::QueueClosed,
});
}
let mut response_timeout = timeout.saturating_add(Duration::from_secs(1));
if lifecycle_cancellable {
// A lifecycle cancel can arrive at the original RPC deadline and
// intentionally keeps the worker alive while extension cleanup
// settles. Do not let this outer waiter abandon that reconciliation
// early, which could reopen a daemon session whose window is gone.
response_timeout = response_timeout.saturating_add(cancel_cleanup_timeout);
}
let waited = tokio::time::timeout(response_timeout, respond_rx).await;
// Lifecycle teardown and effect-aware transfers both keep the worker busy
// for bounded compensation after their original deadline. Keep the outer
// waiter alive for the same grace period so it cannot abandon reconciliation.
let response_grace = if lifecycle_cancellable || effect_aware_transfer {
cancel_cleanup_timeout
} else {
Duration::ZERO
};
let waited = tokio::time::timeout(
timeout
.saturating_add(response_grace)
.saturating_add(Duration::from_secs(1)),
respond_rx,
)
.await;
match waited {
Ok(Ok(Ok(v))) => Ok(v),
Ok(Ok(Err(rpc))) => Err(DispatchError::Rpc(rpc)),
@@ -636,44 +643,85 @@ async fn forward_one(
.lifecycle_cancel
.as_ref()
.map(|_| &forward_lifecycle_cancel as &(dyn Fn() + Sync));
let send_deadline_cancel = || {
client
.sink
.send(Frame::Request(RequestFrame {
id: format!("deadline-cancel-{rpc_id}"),
method: Method::Cancel,
params: Some(serde_json::json!({ "rpc_id": rpc_id })),
}))
.is_ok()
};
let deadline_cancel: Option<&(dyn Fn() -> bool + Sync)> =
is_effect_aware_transfer(&job.method).then_some(&send_deadline_cancel);
let waited = await_with_optional_cancel(
job.timeout,
job.cancel_cleanup_timeout,
waiter,
cancel_token.as_ref(),
on_abort,
deadline_cancel,
)
.await;
let response = match waited {
WaitOutcome::Response(resp) => resp,
WaitOutcome::CancelledAfterResponse(resp) => {
if job.method == Method::ToolSessionStop
&& let ResponseBody::Ok(value) = resp.body
&& let ResponseBody::Ok(value) = &resp.body
{
// Teardown crossed its final irreversible boundary before
// the cancel landed. Commit the real extension result so the
// daemon does not retain a session whose window is gone.
return Ok(value);
return Ok(value.clone());
}
// Cancellation wins the external verdict, but only after the
// extension's original RPC has settled. Preserve any non-cancel
// error so compensation failures remain explicit instead of
// being hidden behind a generic cancelled result.
if let ResponseBody::Err(err) = resp.body
&& !matches!(err.code, ErrorCode::Cancelled | ErrorCode::UserAborted)
{
return Err(err);
// File transfer commits are irreversible. A late cancel cannot
// overwrite a confirmed success, and an unknown transfer effect
// must remain explicit so callers do not retry or release upload
// staging as though nothing happened.
if is_effect_aware_transfer(&job.method) {
match &resp.body {
ResponseBody::Ok(_) => resp,
ResponseBody::Err(err)
if transfer_effect(err) == Some("unknown")
|| transfer_effect(err) == Some("committed") =>
{
return Err(err.clone());
}
ResponseBody::Err(err)
if !matches!(err.code, ErrorCode::Cancelled | ErrorCode::UserAborted) =>
{
return Err(err.clone());
}
ResponseBody::Err(_) => {
return Err(cancelled_error(
job.inflight.as_deref(),
"tool dispatch cancelled after extension cleanup",
));
}
}
} else {
// For ordinary tools cancellation keeps the existing verdict,
// while non-cancel errors still expose compensation failures.
if let ResponseBody::Err(err) = resp.body
&& !matches!(err.code, ErrorCode::Cancelled | ErrorCode::UserAborted)
{
return Err(err);
}
return Err(cancelled_error(
job.inflight.as_deref(),
"tool dispatch cancelled after extension cleanup",
));
}
return Err(cancelled_error(
job.inflight.as_deref(),
"tool dispatch cancelled after extension cleanup",
));
}
WaitOutcome::TimedOutAfterResponse(resp) => match resp.body {
ResponseBody::Ok(value) if job.method == Method::ToolSessionStop => {
ResponseBody::Ok(value)
if job.method == Method::ToolSessionStop
|| is_effect_aware_transfer(&job.method) =>
{
// The close crossed its irreversible boundary during the
// timeout cleanup grace. Reconcile daemon state instead of
// reopening a session whose window is already gone.
// timeout cleanup grace, or the transfer committed before its
// deadline cancel settled. Preserve the irreversible result.
return Ok(value);
}
ResponseBody::Err(err)
@@ -681,6 +729,9 @@ async fn forward_one(
{
return Err(err);
}
ResponseBody::Err(err) if is_effect_aware_transfer(&job.method) => {
return Err(timed_out_transfer_error(&err));
}
_ => {
return Err(RpcError {
code: ErrorCode::Timeout,
@@ -691,6 +742,14 @@ async fn forward_one(
},
WaitOutcome::CleanupTimeout => {
client.pending.lock().unwrap().cancel(&rpc_id);
if is_effect_aware_transfer(&job.method) {
return Err(unknown_transfer_error(
ErrorCode::Timeout,
"cancelled file transfer did not confirm its outcome before cleanup timed out",
"cleanup",
true,
));
}
return Err(RpcError {
code: ErrorCode::Timeout,
message: format!(
@@ -700,8 +759,25 @@ async fn forward_one(
data: Some(serde_json::json!({ "reason": "cancel_cleanup_timeout" })),
});
}
WaitOutcome::TimeoutCleanupFailed => {
client.pending.lock().unwrap().cancel(&rpc_id);
return Err(unknown_transfer_error(
ErrorCode::Timeout,
"file transfer timed out and cleanup could not be confirmed",
"cleanup",
true,
));
}
WaitOutcome::WaiterClosed => {
client.pending.lock().unwrap().cancel(&rpc_id);
if is_effect_aware_transfer(&job.method) {
return Err(unknown_transfer_error(
ErrorCode::ProtocolError,
"file transfer transport closed after dispatch; outcome is unknown",
"transport",
false,
));
}
return Err(RpcError {
code: ErrorCode::ProtocolError,
message: "transport closed mid-call".into(),
@@ -710,6 +786,14 @@ async fn forward_one(
}
WaitOutcome::Timeout => {
client.pending.lock().unwrap().cancel(&rpc_id);
if is_effect_aware_transfer(&job.method) {
return Err(unknown_transfer_error(
ErrorCode::Timeout,
"file transfer timed out after dispatch; outcome is unknown",
"transport",
false,
));
}
return Err(RpcError {
code: ErrorCode::Timeout,
message: format!("tool RPC timed out after {:?}", job.timeout),
@@ -723,12 +807,65 @@ async fn forward_one(
}
}
fn is_effect_aware_transfer(method: &Method) -> bool {
matches!(method, Method::ToolUpload | Method::ToolDownload)
}
fn transfer_effect(err: &RpcError) -> Option<&str> {
err.data.as_ref()?.get("effect_state")?.as_str()
}
fn unknown_transfer_error(
code: ErrorCode,
message: impl Into<String>,
phase: &str,
cleanup_failed: bool,
) -> RpcError {
let mut data = serde_json::json!({
"reason": "transfer_outcome_unknown",
"effect_state": "unknown",
"phase": phase,
});
if cleanup_failed {
data["cleanup_state"] = serde_json::json!("failed");
}
RpcError {
code,
message: message.into(),
data: Some(data),
}
}
fn timed_out_transfer_error(err: &RpcError) -> RpcError {
let data = err.data.as_ref();
RpcError {
code: ErrorCode::Timeout,
message: "file transfer timed out after dispatch".into(),
data: Some(serde_json::json!({
"reason": "transfer_timeout",
"effect_state": data
.and_then(|value| value.get("effect_state"))
.and_then(Value::as_str)
.unwrap_or("unknown"),
"phase": data
.and_then(|value| value.get("phase"))
.and_then(Value::as_str)
.unwrap_or("cleanup"),
"cleanup_state": data
.and_then(|value| value.get("cleanup_state"))
.and_then(Value::as_str)
.unwrap_or("complete"),
})),
}
}
#[derive(Debug)]
enum WaitOutcome {
Response(bsk_protocol::ResponseFrame),
CancelledAfterResponse(bsk_protocol::ResponseFrame),
TimedOutAfterResponse(bsk_protocol::ResponseFrame),
CleanupTimeout,
TimeoutCleanupFailed,
WaiterClosed,
Timeout,
}
@@ -739,6 +876,7 @@ async fn await_with_optional_cancel(
mut waiter: oneshot::Receiver<bsk_protocol::ResponseFrame>,
cancel: Option<&super::abort::AbortToken>,
on_abort: Option<&(dyn Fn() + Sync)>,
on_deadline: Option<&(dyn Fn() -> bool + Sync)>,
) -> WaitOutcome {
match cancel {
Some(token) => {
@@ -764,6 +902,16 @@ async fn await_with_optional_cancel(
Err(_) => WaitOutcome::WaiterClosed,
},
_ = &mut deadline => {
if let Some(send_cancel) = on_deadline {
return if send_cancel() {
match tokio::time::timeout(cleanup_timeout, &mut waiter).await {
Ok(Ok(resp)) => WaitOutcome::TimedOutAfterResponse(resp),
Ok(Err(_)) | Err(_) => WaitOutcome::TimeoutCleanupFailed,
}
} else {
WaitOutcome::TimeoutCleanupFailed
};
}
let Some(on_abort) = on_abort else {
return WaitOutcome::Timeout;
};
@@ -779,7 +927,16 @@ async fn await_with_optional_cancel(
None => match tokio::time::timeout(timeout, &mut waiter).await {
Ok(Ok(resp)) => WaitOutcome::Response(resp),
Ok(Err(_)) => WaitOutcome::WaiterClosed,
Err(_) => WaitOutcome::Timeout,
Err(_) => match on_deadline {
Some(send_cancel) if send_cancel() => {
match tokio::time::timeout(cleanup_timeout, &mut waiter).await {
Ok(Ok(resp)) => WaitOutcome::TimedOutAfterResponse(resp),
Ok(Err(_)) | Err(_) => WaitOutcome::TimeoutCleanupFailed,
}
}
Some(_) => WaitOutcome::TimeoutCleanupFailed,
None => WaitOutcome::Timeout,
},
},
}
}
@@ -864,6 +1021,7 @@ mod await_with_optional_cancel_tests {
rx,
Some(&token),
None,
None,
)
.await;
assert!(
@@ -887,6 +1045,7 @@ mod await_with_optional_cancel_tests {
rx,
Some(&token),
None,
None,
)
.await;
match outcome {
@@ -907,6 +1066,7 @@ mod await_with_optional_cancel_tests {
rx,
None,
None,
None,
)
.await;
assert!(matches!(outcome, WaitOutcome::Response(_)));
@@ -924,6 +1084,7 @@ mod await_with_optional_cancel_tests {
rx,
Some(&token),
None,
None,
)
.await
});
@@ -948,6 +1109,7 @@ mod await_with_optional_cancel_tests {
rx,
Some(&token),
None,
None,
)
.await;
assert!(matches!(outcome, WaitOutcome::CleanupTimeout));
@@ -972,12 +1134,49 @@ mod await_with_optional_cancel_tests {
rx,
Some(&token),
Some(&forward_abort),
None,
)
.await;
assert!(abort_forwarded.load(Ordering::SeqCst));
assert!(matches!(outcome, WaitOutcome::TimedOutAfterResponse(_)));
}
#[tokio::test]
async fn transfer_deadline_sends_cancel_and_waits_for_cleanup_response() {
let (tx, rx) = oneshot::channel();
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(5)).await;
tx.send(dummy_response()).unwrap();
});
let send_cancel = || true;
let outcome = await_with_optional_cancel(
Duration::from_millis(1),
Duration::from_secs(1),
rx,
None,
None,
Some(&send_cancel),
)
.await;
assert!(matches!(outcome, WaitOutcome::TimedOutAfterResponse(_)));
}
#[tokio::test]
async fn transfer_deadline_is_unknown_when_cancel_cannot_be_sent() {
let (_tx, rx) = oneshot::channel();
let send_cancel = || false;
let outcome = await_with_optional_cancel(
Duration::from_millis(1),
Duration::from_secs(1),
rx,
None,
None,
Some(&send_cancel),
)
.await;
assert!(matches!(outcome, WaitOutcome::TimeoutCleanupFailed));
}
}
#[cfg(test)]