mirror of
https://github.com/nimbold/Firelink.git
synced 2026-09-04 06:55:23 +00:00
fix(downloads): harden queue and media retry controls
Add aria2 control epochs so delayed resume workers cannot unpause a transfer after a newer pause, remove, or reconfigure command wins. Thread captured cookies into yt-dlp media runs, sanitize yt-dlp config values, reject malformed media headers, and preserve resumable media artifacts across transient retry backoff. Cover the new queue epoch and media helper behavior with focused regression tests.
This commit is contained in:
+182
-30
@@ -951,6 +951,67 @@ async fn cleanup_media_artifacts(out_path: &std::path::Path, remove_primary: boo
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn sanitize_ytdlp_config_value(value: &str) -> String {
|
||||||
|
value.replace(['\n', '\r'], "")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn append_ytdlp_config_option(config: &mut String, option: &str, value: &str) {
|
||||||
|
let safe_value = sanitize_ytdlp_config_value(value);
|
||||||
|
if !safe_value.is_empty() {
|
||||||
|
config.push_str(option);
|
||||||
|
config.push('\n');
|
||||||
|
config.push_str(&safe_value);
|
||||||
|
config.push('\n');
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn append_ytdlp_add_header(config: &mut String, header: &str) -> Result<bool, String> {
|
||||||
|
let safe_header = sanitize_ytdlp_config_value(header).trim().to_string();
|
||||||
|
if safe_header.is_empty() {
|
||||||
|
return Ok(false);
|
||||||
|
}
|
||||||
|
let Some((name, _)) = safe_header.split_once(':') else {
|
||||||
|
return Err(format!("invalid HTTP header: {safe_header}"));
|
||||||
|
};
|
||||||
|
if name.trim().is_empty() {
|
||||||
|
return Err(format!("invalid HTTP header: {safe_header}"));
|
||||||
|
}
|
||||||
|
append_ytdlp_config_option(config, "--add-header", &safe_header);
|
||||||
|
Ok(name.trim().eq_ignore_ascii_case("cookie"))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn append_ytdlp_http_headers(
|
||||||
|
config: &mut String,
|
||||||
|
headers: Option<&str>,
|
||||||
|
cookies: Option<&str>,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
let mut has_cookie_header = false;
|
||||||
|
if let Some(headers) = headers {
|
||||||
|
for header in headers.lines() {
|
||||||
|
has_cookie_header |= append_ytdlp_add_header(config, header)?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if !has_cookie_header {
|
||||||
|
if let Some(cookies) = cookies {
|
||||||
|
let safe_cookies = sanitize_ytdlp_config_value(cookies).trim().to_string();
|
||||||
|
if !safe_cookies.is_empty() {
|
||||||
|
append_ytdlp_add_header(config, &format!("Cookie: {safe_cookies}"))?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn should_cleanup_media_artifacts_after_failure(
|
||||||
|
failure_reason: &str,
|
||||||
|
strike: usize,
|
||||||
|
max_retries: usize,
|
||||||
|
) -> bool {
|
||||||
|
!(crate::retry::is_transient_network_error(failure_reason) && strike < max_retries)
|
||||||
|
}
|
||||||
|
|
||||||
async fn validate_url_ssrf(url: &str) -> Result<Option<(String, std::net::SocketAddr)>, String> {
|
async fn validate_url_ssrf(url: &str) -> Result<Option<(String, std::net::SocketAddr)>, String> {
|
||||||
let parsed = reqwest::Url::parse(url).map_err(|_| "SSRF blocked: Invalid URL")?;
|
let parsed = reqwest::Url::parse(url).map_err(|_| "SSRF blocked: Invalid URL")?;
|
||||||
if parsed.scheme() != "http" && parsed.scheme() != "https" {
|
if parsed.scheme() != "http" && parsed.scheme() != "https" {
|
||||||
@@ -1343,14 +1404,12 @@ async fn fetch_media_metadata_uncached(
|
|||||||
let mut config_content = String::new();
|
let mut config_content = String::new();
|
||||||
if let Some(user) = username {
|
if let Some(user) = username {
|
||||||
if !user.is_empty() {
|
if !user.is_empty() {
|
||||||
let safe_user = user.replace(['\n', '\r'], "");
|
append_ytdlp_config_option(&mut config_content, "--username", &user);
|
||||||
config_content.push_str(&format!("--username\n{}\n", safe_user));
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if let Some(pass) = password {
|
if let Some(pass) = password {
|
||||||
if !pass.is_empty() {
|
if !pass.is_empty() {
|
||||||
let safe_pass = pass.replace(['\n', '\r'], "");
|
append_ytdlp_config_option(&mut config_content, "--password", &pass);
|
||||||
config_content.push_str(&format!("--password\n{}\n", safe_pass));
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
use std::io::Write;
|
use std::io::Write;
|
||||||
@@ -2485,6 +2544,7 @@ pub(crate) async fn start_media_download_internal(
|
|||||||
username: Option<String>,
|
username: Option<String>,
|
||||||
password: Option<String>,
|
password: Option<String>,
|
||||||
headers: Option<String>,
|
headers: Option<String>,
|
||||||
|
cookies: Option<String>,
|
||||||
proxy: Option<String>,
|
proxy: Option<String>,
|
||||||
user_agent: Option<String>,
|
user_agent: Option<String>,
|
||||||
max_tries: Option<i32>,
|
max_tries: Option<i32>,
|
||||||
@@ -2526,23 +2586,15 @@ pub(crate) async fn start_media_download_internal(
|
|||||||
let mut config_content = String::new();
|
let mut config_content = String::new();
|
||||||
if let Some(user) = username {
|
if let Some(user) = username {
|
||||||
if !user.is_empty() {
|
if !user.is_empty() {
|
||||||
config_content.push_str(&format!("--username\n{}\n", user));
|
append_ytdlp_config_option(&mut config_content, "--username", &user);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if let Some(pass) = password {
|
if let Some(pass) = password {
|
||||||
if !pass.is_empty() {
|
if !pass.is_empty() {
|
||||||
config_content.push_str(&format!("--password\n{}\n", pass));
|
append_ytdlp_config_option(&mut config_content, "--password", &pass);
|
||||||
}
|
|
||||||
}
|
|
||||||
if let Some(headers) = headers {
|
|
||||||
for header in headers
|
|
||||||
.lines()
|
|
||||||
.map(str::trim)
|
|
||||||
.filter(|header| !header.is_empty())
|
|
||||||
{
|
|
||||||
config_content.push_str(&format!("--add-header\n{}\n", header));
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
append_ytdlp_http_headers(&mut config_content, headers.as_deref(), cookies.as_deref())?;
|
||||||
use std::io::Write;
|
use std::io::Write;
|
||||||
config_file
|
config_file
|
||||||
.write_all(config_content.as_bytes())
|
.write_all(config_content.as_bytes())
|
||||||
@@ -2806,7 +2858,6 @@ pub(crate) async fn start_media_download_internal(
|
|||||||
}
|
}
|
||||||
Some(tauri_plugin_shell::process::CommandEvent::Error(err)) => {
|
Some(tauri_plugin_shell::process::CommandEvent::Error(err)) => {
|
||||||
log::error!("yt-dlp shell error [{}]: {}", id, err);
|
log::error!("yt-dlp shell error [{}]: {}", id, err);
|
||||||
cleanup_media_artifacts(&out_path, false).await;
|
|
||||||
break err;
|
break err;
|
||||||
}
|
}
|
||||||
Some(tauri_plugin_shell::process::CommandEvent::Terminated(payload)) => {
|
Some(tauri_plugin_shell::process::CommandEvent::Terminated(payload)) => {
|
||||||
@@ -2832,7 +2883,6 @@ pub(crate) async fn start_media_download_internal(
|
|||||||
return Ok(completed_path);
|
return Ok(completed_path);
|
||||||
}
|
}
|
||||||
log::error!("yt-dlp exited with non-zero code {:?} for id: {}", payload.code, id);
|
log::error!("yt-dlp exited with non-zero code {:?} for id: {}", payload.code, id);
|
||||||
cleanup_media_artifacts(&out_path, false).await;
|
|
||||||
break if stderr_tail.is_empty() {
|
break if stderr_tail.is_empty() {
|
||||||
format!("yt-dlp exited with code {:?}", payload.code)
|
format!("yt-dlp exited with code {:?}", payload.code)
|
||||||
} else {
|
} else {
|
||||||
@@ -2841,7 +2891,6 @@ pub(crate) async fn start_media_download_internal(
|
|||||||
}
|
}
|
||||||
Some(_) => {}
|
Some(_) => {}
|
||||||
None => {
|
None => {
|
||||||
cleanup_media_artifacts(&out_path, false).await;
|
|
||||||
break if stderr_tail.is_empty() {
|
break if stderr_tail.is_empty() {
|
||||||
"yt-dlp process ended unexpectedly".to_string()
|
"yt-dlp process ended unexpectedly".to_string()
|
||||||
} else {
|
} else {
|
||||||
@@ -2855,6 +2904,9 @@ pub(crate) async fn start_media_download_internal(
|
|||||||
|
|
||||||
let transient = is_transient_network_error(&failure_reason);
|
let transient = is_transient_network_error(&failure_reason);
|
||||||
let strikes_left = strike < max_retries;
|
let strikes_left = strike < max_retries;
|
||||||
|
if should_cleanup_media_artifacts_after_failure(&failure_reason, strike, max_retries) {
|
||||||
|
cleanup_media_artifacts(&out_path, false).await;
|
||||||
|
}
|
||||||
if !(transient && strikes_left) {
|
if !(transient && strikes_left) {
|
||||||
return Err(failure_reason);
|
return Err(failure_reason);
|
||||||
}
|
}
|
||||||
@@ -2899,9 +2951,13 @@ async fn pause_download(
|
|||||||
.await?;
|
.await?;
|
||||||
match status.as_str() {
|
match status.as_str() {
|
||||||
"paused" => {
|
"paused" => {
|
||||||
|
state.queue_manager.next_aria2_control_epoch(&id).await;
|
||||||
|
state.queue_manager.cancel_aria2_retries(&id).await;
|
||||||
log::info!("aria2 pause [{}]: gid {} was already paused", id, gid);
|
log::info!("aria2 pause [{}]: gid {} was already paused", id, gid);
|
||||||
}
|
}
|
||||||
"active" | "waiting" => {
|
"active" | "waiting" => {
|
||||||
|
state.queue_manager.next_aria2_control_epoch(&id).await;
|
||||||
|
state.queue_manager.cancel_aria2_retries(&id).await;
|
||||||
let result = rpc_call(
|
let result = rpc_call(
|
||||||
state.aria2_port.load(std::sync::atomic::Ordering::Relaxed),
|
state.aria2_port.load(std::sync::atomic::Ordering::Relaxed),
|
||||||
&state.aria2_secret,
|
&state.aria2_secret,
|
||||||
@@ -2914,9 +2970,23 @@ async fn pause_download(
|
|||||||
log::info!("aria2 pause [{}]: gid {} paused", id, gid);
|
log::info!("aria2 pause [{}]: gid {} paused", id, gid);
|
||||||
}
|
}
|
||||||
terminal => {
|
terminal => {
|
||||||
|
let retrying = state.queue_manager.has_aria2_retry_state(&id).await;
|
||||||
state.queue_manager.clear_aria2_retry_state(&id).await;
|
state.queue_manager.clear_aria2_retry_state(&id).await;
|
||||||
state.queue_manager.forget_aria2_gid(&id).await;
|
state.queue_manager.forget_aria2_gid(&id).await;
|
||||||
state.queue_manager.release_permit(&id).await;
|
state.queue_manager.release_permit(&id).await;
|
||||||
|
state.queue_manager.next_aria2_control_epoch(&id).await;
|
||||||
|
state.queue_manager.cancel_aria2_retries(&id).await;
|
||||||
|
if retrying && matches!(terminal, "error" | "removed") {
|
||||||
|
use tauri::Emitter;
|
||||||
|
let _ = app_handle.emit(
|
||||||
|
"download-state",
|
||||||
|
crate::ipc::DownloadStateEvent::new(
|
||||||
|
id,
|
||||||
|
crate::ipc::DownloadStatus::Paused,
|
||||||
|
),
|
||||||
|
);
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
state.queue_manager.release_registered_id(&id).await;
|
state.queue_manager.release_registered_id(&id).await;
|
||||||
return Err(format!(
|
return Err(format!(
|
||||||
"cannot pause aria2 gid {gid} in terminal state {terminal}"
|
"cannot pause aria2 gid {gid} in terminal state {terminal}"
|
||||||
@@ -2997,6 +3067,8 @@ async fn resume_download(
|
|||||||
.await?;
|
.await?;
|
||||||
match status.as_str() {
|
match status.as_str() {
|
||||||
"paused" => {
|
"paused" => {
|
||||||
|
let control_epoch = state.queue_manager.next_aria2_control_epoch(&id).await;
|
||||||
|
state.queue_manager.allow_aria2_retries(&id).await;
|
||||||
use tauri::Emitter;
|
use tauri::Emitter;
|
||||||
let _ = app_handle.emit(
|
let _ = app_handle.emit(
|
||||||
"download-state",
|
"download-state",
|
||||||
@@ -3012,6 +3084,20 @@ async fn resume_download(
|
|||||||
let app_handle_clone = app_handle.clone();
|
let app_handle_clone = app_handle.clone();
|
||||||
tauri::async_runtime::spawn(async move {
|
tauri::async_runtime::spawn(async move {
|
||||||
let acquired = queue_manager.ensure_aria2_permit(&id_clone).await;
|
let acquired = queue_manager.ensure_aria2_permit(&id_clone).await;
|
||||||
|
if !acquired {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
if queue_manager.is_aria2_retry_cancelled(&id_clone).await
|
||||||
|
|| !queue_manager
|
||||||
|
.is_aria2_control_epoch_current(&id_clone, control_epoch)
|
||||||
|
.await
|
||||||
|
|| queue_manager.aria2_gid_for_download(&id_clone).as_deref()
|
||||||
|
!= Some(gid_clone.as_str())
|
||||||
|
|| !queue_manager.is_registered(&id_clone).await
|
||||||
|
{
|
||||||
|
queue_manager.release_permit(&id_clone).await;
|
||||||
|
return;
|
||||||
|
}
|
||||||
let result = match rpc_call(
|
let result = match rpc_call(
|
||||||
aria2_port,
|
aria2_port,
|
||||||
&aria2_secret,
|
&aria2_secret,
|
||||||
@@ -3022,9 +3108,7 @@ async fn resume_download(
|
|||||||
{
|
{
|
||||||
Ok(result) => result,
|
Ok(result) => result,
|
||||||
Err(error) => {
|
Err(error) => {
|
||||||
if acquired {
|
queue_manager.release_permit(&id_clone).await;
|
||||||
queue_manager.release_permit(&id_clone).await;
|
|
||||||
}
|
|
||||||
log::error!("failed to resume aria2 gid {}: {}", gid_clone, error);
|
log::error!("failed to resume aria2 gid {}: {}", gid_clone, error);
|
||||||
let _ = app_handle_clone.emit(
|
let _ = app_handle_clone.emit(
|
||||||
"download-state",
|
"download-state",
|
||||||
@@ -3037,9 +3121,7 @@ async fn resume_download(
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
if let Err(error) = ensure_aria2_gid_result("unpause", &gid_clone, &result) {
|
if let Err(error) = ensure_aria2_gid_result("unpause", &gid_clone, &result) {
|
||||||
if acquired {
|
queue_manager.release_permit(&id_clone).await;
|
||||||
queue_manager.release_permit(&id_clone).await;
|
|
||||||
}
|
|
||||||
log::error!("failed to resume aria2 gid {}: {}", gid_clone, error);
|
log::error!("failed to resume aria2 gid {}: {}", gid_clone, error);
|
||||||
let _ = app_handle_clone.emit(
|
let _ = app_handle_clone.emit(
|
||||||
"download-state",
|
"download-state",
|
||||||
@@ -3050,6 +3132,23 @@ async fn resume_download(
|
|||||||
);
|
);
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
if queue_manager.is_aria2_retry_cancelled(&id_clone).await
|
||||||
|
|| !queue_manager
|
||||||
|
.is_aria2_control_epoch_current(&id_clone, control_epoch)
|
||||||
|
.await
|
||||||
|
|| queue_manager.aria2_gid_for_download(&id_clone).as_deref()
|
||||||
|
!= Some(gid_clone.as_str())
|
||||||
|
{
|
||||||
|
let _ = rpc_call(
|
||||||
|
aria2_port,
|
||||||
|
&aria2_secret,
|
||||||
|
"aria2.forcePause",
|
||||||
|
serde_json::json!([gid_clone]),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
queue_manager.release_permit(&id_clone).await;
|
||||||
|
return;
|
||||||
|
}
|
||||||
log::info!("aria2 resume [{}]: unpaused gid {}", id_clone, gid_clone);
|
log::info!("aria2 resume [{}]: unpaused gid {}", id_clone, gid_clone);
|
||||||
let _ = app_handle_clone.emit(
|
let _ = app_handle_clone.emit(
|
||||||
"download-state",
|
"download-state",
|
||||||
@@ -3109,6 +3208,7 @@ async fn remove_download(
|
|||||||
let active_kind = state.queue_manager.active_kind(&id).await;
|
let active_kind = state.queue_manager.active_kind(&id).await;
|
||||||
state.queue_manager.remove_from_pending(&id).await;
|
state.queue_manager.remove_from_pending(&id).await;
|
||||||
|
|
||||||
|
state.queue_manager.next_aria2_control_epoch(&id).await;
|
||||||
state.queue_manager.cancel_aria2_retries(&id).await;
|
state.queue_manager.cancel_aria2_retries(&id).await;
|
||||||
|
|
||||||
let gid = state.queue_manager.aria2_gid_for_download(&id);
|
let gid = state.queue_manager.aria2_gid_for_download(&id);
|
||||||
@@ -3251,6 +3351,7 @@ async fn detach_download_for_reconfigure(
|
|||||||
log::info!("detach_download_for_reconfigure called for id: {}", id);
|
log::info!("detach_download_for_reconfigure called for id: {}", id);
|
||||||
let active_kind = state.queue_manager.active_kind(&id).await;
|
let active_kind = state.queue_manager.active_kind(&id).await;
|
||||||
state.queue_manager.remove_from_pending(&id).await;
|
state.queue_manager.remove_from_pending(&id).await;
|
||||||
|
state.queue_manager.next_aria2_control_epoch(&id).await;
|
||||||
state.queue_manager.cancel_aria2_retries(&id).await;
|
state.queue_manager.cancel_aria2_retries(&id).await;
|
||||||
|
|
||||||
let gid = state.queue_manager.aria2_gid_for_download(&id);
|
let gid = state.queue_manager.aria2_gid_for_download(&id);
|
||||||
@@ -4308,12 +4409,13 @@ fn set_extension_frontend_ready(state: tauri::State<'_, AppState>, ready: bool)
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::{
|
use super::{
|
||||||
aggregate_media_fraction, build_media_format_options, collect_download_uris,
|
aggregate_media_fraction, append_ytdlp_http_headers, build_media_format_options,
|
||||||
filename_from_content_disposition, filename_from_url_disposition_query,
|
collect_download_uris, filename_from_content_disposition,
|
||||||
filename_from_url_path, is_excluded_yt_dlp_format, json_lower, media_output_template,
|
filename_from_url_disposition_query, filename_from_url_path, is_excluded_yt_dlp_format,
|
||||||
media_progress_speed, normalize_speed_limit_for_aria2, parse_firelink_deep_link,
|
json_lower, media_output_template, media_progress_speed, normalize_speed_limit_for_aria2,
|
||||||
parse_ffmpeg_version, parse_media_progress_line, redact_log_line, FirelinkDeepLink, MediaProgress,
|
parse_firelink_deep_link, parse_ffmpeg_version, parse_media_progress_line, redact_log_line,
|
||||||
MEDIA_PROGRESS_PREFIX,
|
sanitize_ytdlp_config_value, should_cleanup_media_artifacts_after_failure,
|
||||||
|
FirelinkDeepLink, MediaProgress, MEDIA_PROGRESS_PREFIX,
|
||||||
};
|
};
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
use std::time::{Duration, Instant};
|
use std::time::{Duration, Instant};
|
||||||
@@ -4335,6 +4437,56 @@ mod tests {
|
|||||||
assert_eq!(template, destination.join("clip.mp4"));
|
assert_eq!(template, destination.join("clip.mp4"));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn ytdlp_config_values_cannot_inject_extra_lines() {
|
||||||
|
assert_eq!(
|
||||||
|
sanitize_ytdlp_config_value("user\n--exec\rmalicious"),
|
||||||
|
"user--execmalicious"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn ytdlp_media_headers_include_captured_cookies_once() {
|
||||||
|
let mut config = String::new();
|
||||||
|
append_ytdlp_http_headers(
|
||||||
|
&mut config,
|
||||||
|
Some("Referer: https://example.com/video"),
|
||||||
|
Some("session=abc\r\n--proxy=http://bad.invalid"),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert!(config.contains("--add-header\nReferer: https://example.com/video\n"));
|
||||||
|
assert!(config.contains("--add-header\nCookie: session=abc--proxy=http://bad.invalid\n"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn ytdlp_media_headers_reject_invalid_lines() {
|
||||||
|
let mut config = String::new();
|
||||||
|
let error = append_ytdlp_http_headers(&mut config, Some("not a header"), None)
|
||||||
|
.expect_err("invalid header line should be rejected");
|
||||||
|
|
||||||
|
assert!(error.contains("invalid HTTP header"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn retryable_media_failures_preserve_resumable_artifacts() {
|
||||||
|
assert!(!should_cleanup_media_artifacts_after_failure(
|
||||||
|
"The response status is not successful. status=503",
|
||||||
|
0,
|
||||||
|
1
|
||||||
|
));
|
||||||
|
assert!(should_cleanup_media_artifacts_after_failure(
|
||||||
|
"The response status is not successful. status=503",
|
||||||
|
1,
|
||||||
|
1
|
||||||
|
));
|
||||||
|
assert!(should_cleanup_media_artifacts_after_failure(
|
||||||
|
"HTTP 404 Not Found",
|
||||||
|
0,
|
||||||
|
3
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn metadata_filename_prefers_content_disposition_filename() {
|
fn metadata_filename_prefers_content_disposition_filename() {
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
|
|||||||
+32
-1
@@ -113,6 +113,10 @@ pub struct QueueManager<R: tauri::Runtime = tauri::Wry> {
|
|||||||
/// Download ids whose aria2 retry loop must not create another job.
|
/// Download ids whose aria2 retry loop must not create another job.
|
||||||
aria2_retry_cancelled: Mutex<HashSet<String>>,
|
aria2_retry_cancelled: Mutex<HashSet<String>>,
|
||||||
|
|
||||||
|
/// Monotonic per-download aria2 control generation. Long-running queued
|
||||||
|
/// resume tasks capture this and abort when a later pause/remove wins.
|
||||||
|
aria2_control_epochs: Mutex<HashMap<String, u64>>,
|
||||||
|
|
||||||
spawner: Arc<dyn SidecarSpawner>,
|
spawner: Arc<dyn SidecarSpawner>,
|
||||||
app_handle: AppHandle<R>,
|
app_handle: AppHandle<R>,
|
||||||
}
|
}
|
||||||
@@ -147,6 +151,7 @@ impl<R: tauri::Runtime> QueueManager<R> {
|
|||||||
aria2_payloads: Mutex::new(HashMap::new()),
|
aria2_payloads: Mutex::new(HashMap::new()),
|
||||||
aria2_retry_strikes: Mutex::new(HashMap::new()),
|
aria2_retry_strikes: Mutex::new(HashMap::new()),
|
||||||
aria2_retry_cancelled: Mutex::new(HashSet::new()),
|
aria2_retry_cancelled: Mutex::new(HashSet::new()),
|
||||||
|
aria2_control_epochs: Mutex::new(HashMap::new()),
|
||||||
spawner,
|
spawner,
|
||||||
app_handle,
|
app_handle,
|
||||||
}
|
}
|
||||||
@@ -168,10 +173,35 @@ impl<R: tauri::Runtime> QueueManager<R> {
|
|||||||
self.registered_ids.lock().await.remove(id);
|
self.registered_ids.lock().await.remove(id);
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn is_registered(&self, id: &str) -> bool {
|
pub async fn is_registered(&self, id: &str) -> bool {
|
||||||
self.registered_ids.lock().await.contains(id)
|
self.registered_ids.lock().await.contains(id)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub async fn next_aria2_control_epoch(&self, id: &str) -> u64 {
|
||||||
|
let mut epochs = self.aria2_control_epochs.lock().await;
|
||||||
|
let epoch = epochs.get(id).copied().unwrap_or_default().wrapping_add(1);
|
||||||
|
epochs.insert(id.to_string(), epoch);
|
||||||
|
epoch
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn is_aria2_control_epoch_current(&self, id: &str, epoch: u64) -> bool {
|
||||||
|
self.aria2_control_epochs
|
||||||
|
.lock()
|
||||||
|
.await
|
||||||
|
.get(id)
|
||||||
|
.copied()
|
||||||
|
.unwrap_or_default()
|
||||||
|
== epoch
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn is_aria2_retry_cancelled(&self, id: &str) -> bool {
|
||||||
|
self.aria2_retry_cancelled.lock().await.contains(id)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn has_aria2_retry_state(&self, id: &str) -> bool {
|
||||||
|
self.aria2_retry_strikes.lock().await.contains_key(id)
|
||||||
|
}
|
||||||
|
|
||||||
/// Enqueue a task. Checks the centralized `registered_ids` for deduplication.
|
/// Enqueue a task. Checks the centralized `registered_ids` for deduplication.
|
||||||
pub async fn push(&self, task: QueuedTask) -> Result<(), String> {
|
pub async fn push(&self, task: QueuedTask) -> Result<(), String> {
|
||||||
let id = task.id.clone();
|
let id = task.id.clone();
|
||||||
@@ -1267,6 +1297,7 @@ impl SidecarSpawner for ProductionSpawner {
|
|||||||
payload.username.clone(),
|
payload.username.clone(),
|
||||||
payload.password.clone(),
|
payload.password.clone(),
|
||||||
payload.headers.clone(),
|
payload.headers.clone(),
|
||||||
|
payload.cookies.clone(),
|
||||||
payload.proxy.clone(),
|
payload.proxy.clone(),
|
||||||
payload.user_agent.clone(),
|
payload.user_agent.clone(),
|
||||||
payload.max_tries,
|
payload.max_tries,
|
||||||
|
|||||||
@@ -139,6 +139,18 @@ async fn ensure_aria2_permit_does_not_double_acquire() {
|
|||||||
assert_eq!(mgr.available_permits(), 2);
|
assert_eq!(mgr.available_permits(), 2);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn aria2_control_epoch_invalidates_stale_resume_workers() {
|
||||||
|
let (mgr, _spawner) = make_manager(1);
|
||||||
|
let first_resume = mgr.next_aria2_control_epoch("a").await;
|
||||||
|
assert!(mgr.is_aria2_control_epoch_current("a", first_resume).await);
|
||||||
|
|
||||||
|
let pause = mgr.next_aria2_control_epoch("a").await;
|
||||||
|
assert_ne!(pause, first_resume);
|
||||||
|
assert!(!mgr.is_aria2_control_epoch_current("a", first_resume).await);
|
||||||
|
assert!(mgr.is_aria2_control_epoch_current("a", pause).await);
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn forgetting_aria2_gid_clears_mapping_without_releasing_twice() {
|
async fn forgetting_aria2_gid_clears_mapping_without_releasing_twice() {
|
||||||
let (mgr, _spawner) = make_manager(1);
|
let (mgr, _spawner) = make_manager(1);
|
||||||
|
|||||||
Reference in New Issue
Block a user