From 9133e3b05b57563939d146f15c13b1ddc8253104 Mon Sep 17 00:00:00 2001 From: NimBold Date: Sun, 12 Jul 2026 07:55:47 +0330 Subject: [PATCH] fix(downloads): harden release-critical transfer paths --- src-tauri/src/extension_server.rs | 26 ++++++++++++-- src-tauri/src/lib.rs | 24 +++++++++++++ src-tauri/src/queue.rs | 51 ++++++++++++++++------------ src-tauri/tests/queue_manager.rs | 49 ++++++++++++++++++++++++++ src/components/AddDownloadsModal.tsx | 25 ++++++++++++-- src/store/useDownloadStore.test.ts | 3 +- src/store/useDownloadStore.ts | 15 +++++++- src/utils/clipboard.test.ts | 3 +- src/utils/url.ts | 2 +- 9 files changed, 168 insertions(+), 30 deletions(-) diff --git a/src-tauri/src/extension_server.rs b/src-tauri/src/extension_server.rs index 089cc7c..3c3ea5a 100644 --- a/src-tauri/src/extension_server.rs +++ b/src-tauri/src/extension_server.rs @@ -308,13 +308,14 @@ fn normalize_download(payload: ExtensionRequest) -> Option { matches!(url.scheme(), "http" | "https").then(|| url.to_string()) }); let filename = payload.filename.and_then(|value| sanitize_filename(&value)); + let headers = normalize_headers(payload.headers, payload.media); Some(ExtensionDownload { urls, referer, silent: payload.silent, filename, - headers: payload.headers.filter(|value| !value.trim().is_empty()), + headers, // Explicit media is resolved by yt-dlp, which must use Firelink's // configured browser-cookie source. Forwarding a browser's complete // Cookie header can exceed upstream limits and makes old extension @@ -328,6 +329,24 @@ fn normalize_download(payload: ExtensionRequest) -> Option { }) } +fn normalize_headers(headers: Option, media: bool) -> Option { + let headers = headers?; + if !media { + return (!headers.trim().is_empty()).then_some(headers); + } + + let filtered = headers + .lines() + .filter(|line| { + line.split_once(':') + .map(|(name, _)| !name.trim().eq_ignore_ascii_case("cookie")) + .unwrap_or(true) + }) + .collect::>() + .join("\n"); + (!filtered.trim().is_empty()).then_some(filtered) +} + fn normalize_url(raw_url: &str) -> Option { let url = Url::parse(raw_url.trim()).ok()?; matches!(url.scheme(), "http" | "https" | "ftp" | "sftp").then(|| url.to_string()) @@ -507,7 +526,10 @@ mod tests { referer: None, silent: false, filename: None, - headers: Some("User-Agent: Firefox".to_string()), + headers: Some(format!( + "Cookie: stale={};\nUser-Agent: Firefox", + "x".repeat(64 * 1024) + )), cookies: Some(format!("large={}", "x".repeat(64 * 1024))), media: true, }) diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 99ede8e..57f4ba0 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -126,6 +126,16 @@ fn metadata_filename_from_response( .unwrap_or_else(|| "download".to_string()) } +fn metadata_response_error(status: reqwest::StatusCode) -> Option { + (!status.is_success()).then(|| { + format!( + "Metadata request failed with HTTP {} ({})", + status.as_u16(), + status.canonical_reason().unwrap_or("unknown status") + ) + }) +} + #[derive(Serialize, TS)] #[ts(export, export_to = "../../src/bindings/")] pub struct MetadataResponse { @@ -1367,6 +1377,10 @@ async fn fetch_metadata( } } + if let Some(error) = metadata_response_error(current_res.status()) { + return Err(error); + } + res = current_res; break; } @@ -4814,6 +4828,7 @@ mod tests { filename_from_url_disposition_query, filename_from_url_path, is_excluded_yt_dlp_format, is_browser_cookie_extraction_error, json_lower, media_metadata_cache_key, media_output_template, media_progress_args, media_progress_speed, + metadata_response_error, normalize_speed_limit_for_aria2, parse_firelink_deep_link, parse_ffmpeg_version, parse_media_progress_line, redact_log_line, redact_log_line_for_output, sanitize_ytdlp_config_value, @@ -4841,6 +4856,15 @@ mod tests { assert_eq!(template, destination.join("clip.mp4")); } + #[test] + fn metadata_rejects_final_http_errors_but_accepts_partial_content() { + assert!(metadata_response_error(reqwest::StatusCode::PARTIAL_CONTENT).is_none()); + assert_eq!( + metadata_response_error(reqwest::StatusCode::NOT_FOUND).as_deref(), + Some("Metadata request failed with HTTP 404 (Not Found)") + ); + } + #[test] fn recognizes_resume_sidecars_without_treating_the_primary_file_as_partial() { let directory = tempfile::tempdir().unwrap(); diff --git a/src-tauri/src/queue.rs b/src-tauri/src/queue.rs index a586980..0a0c1ec 100644 --- a/src-tauri/src/queue.rs +++ b/src-tauri/src/queue.rs @@ -93,7 +93,6 @@ pub struct QueueManager { target_capacity: AtomicUsize, slots_to_retire: AtomicUsize, notify: Notify, - notify_permit_released: Notify, /// aria2 gid -> download id map (shared with the WS poller). pub aria2_gids: Arc>>, @@ -110,6 +109,8 @@ pub struct QueueManager { /// Download ids whose aria2 retry loop must not create another job. aria2_retry_cancelled: Mutex>, + /// Wakes retry backoff workers when a pause/remove action cancels them. + aria2_retry_cancel_notify: Notify, /// Monotonic per-download aria2 control generation. Long-running queued /// resume tasks capture this and abort when a later pause/remove wins. @@ -145,12 +146,12 @@ impl QueueManager { target_capacity: AtomicUsize::new(capacity), slots_to_retire: AtomicUsize::new(0), notify: Notify::new(), - notify_permit_released: Notify::new(), aria2_gids: Arc::new(std::sync::RwLock::new(HashMap::new())), pending_completion: Arc::new(Mutex::new(HashMap::new())), aria2_payloads: Mutex::new(HashMap::new()), aria2_retry_strikes: Mutex::new(HashMap::new()), aria2_retry_cancelled: Mutex::new(HashSet::new()), + aria2_retry_cancel_notify: Notify::new(), aria2_control_epochs: Mutex::new(HashMap::new()), spawner, app_handle, @@ -384,7 +385,6 @@ impl QueueManager { let removed = self.active_permits.lock().await.remove(id).is_some(); self.active_kinds.lock().await.remove(id); if removed { - self.notify_permit_released.notify_waiters(); self.notify.notify_one(); } } @@ -644,6 +644,7 @@ impl QueueManager { .lock() .await .insert(id.to_string()); + self.aria2_retry_cancel_notify.notify_waiters(); } pub async fn allow_aria2_retries(&self, id: &str) { @@ -686,22 +687,6 @@ impl QueueManager { removed.last().cloned() } - async fn wait_permit_released(self: &Arc, id: &str) { - loop { - if !self.active_permits.lock().await.contains_key(id) { - return; - } - let notified = self.notify_permit_released.notified(); - if !self.active_permits.lock().await.contains_key(id) { - return; - } - tokio::select! { - _ = notified => {} - _ = tokio::time::sleep(std::time::Duration::from_secs(5)) => {} - } - } - } - /// Intercept transient `onDownloadError` events: backoff, re-issue /// `addUri`, and rotate the gid mapping. Permanent errors and exhausted /// strikes fall through to a hard `Failed` state. @@ -774,10 +759,24 @@ impl QueueManager { let id_for_task = id.clone(); let error_for_emit = error.clone(); tauri::async_runtime::spawn(async move { + let retry_cancel = async { + loop { + if this.is_aria2_retry_cancelled(&id_for_task).await { + break; + } + let notified = this.aria2_retry_cancel_notify.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + if this.is_aria2_retry_cancelled(&id_for_task).await { + break; + } + notified.await; + } + }; let outcome = backoff_and_emit( strike, error_for_emit, - this.wait_permit_released(&id_for_task), + retry_cancel, |reason| { use tauri::Emitter; let _ = this.app_handle.emit( @@ -938,7 +937,11 @@ fn automatic_retry_limit(max_tries: Option) -> usize { } fn aria2_attempt_limit(max_tries: Option) -> u32 { - (automatic_retry_limit(max_tries) + 1) as u32 + // Firelink owns the retry budget and performs the backoff/GID rotation. + // Keep each aria2 GID to one attempt so `max_tries` is not multiplied by + // aria2's own internal retry loop. + let _ = max_tries; + 1 } fn is_retryable_aria2_error(error: &str) -> bool { @@ -1514,4 +1517,10 @@ mod tests { "aria2 error code 3: Resource not found" )); } + + #[test] + fn aria2_internal_attempts_do_not_multiply_firelink_retry_budget() { + assert_eq!(aria2_attempt_limit(Some(0)), 1); + assert_eq!(aria2_attempt_limit(Some(10)), 1); + } } diff --git a/src-tauri/tests/queue_manager.rs b/src-tauri/tests/queue_manager.rs index 2902f9b..2ce0049 100644 --- a/src-tauri/tests/queue_manager.rs +++ b/src-tauri/tests/queue_manager.rs @@ -514,6 +514,55 @@ async fn aria2_permit_survives_rpc_return() { handle.abort(); } +#[tokio::test] +async fn transient_aria2_error_reissues_after_backoff() { + use firelink_lib::queue::PendingOutcome; + + let (mgr, spawner) = make_manager(1); + let manager = Arc::new(mgr); + let mut task = aria2_task("retry"); + task.payload.max_tries = Some(1); + manager.push(task).await.unwrap(); + + let dispatcher = { + let manager = Arc::clone(&manager); + tokio::spawn(async move { manager.run_dispatcher().await }) + }; + tokio::time::sleep(Duration::from_millis(100)).await; + assert_eq!(spawner.add_uri_calls.load(Ordering::SeqCst), 1); + + manager + .handle_aria2_event( + "gid-1", + PendingOutcome::Error("Timeout.".to_string()), + ) + .await; + + timeout(Duration::from_secs(4), async { + loop { + if spawner.add_uri_calls.load(Ordering::SeqCst) >= 2 { + break; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + }) + .await + .expect("transient aria2 errors must retry after backoff"); + + manager + .handle_aria2_event( + "gid-2", + PendingOutcome::Error("Timeout.".to_string()), + ) + .await; + assert_eq!(spawner.add_uri_calls.load(Ordering::SeqCst), 2); + assert!(manager.aria2_gid_for_download("retry").is_none()); + assert_eq!(manager.available_permits(), 1); + + manager.release_permit("retry").await; + dispatcher.abort(); +} + #[tokio::test] async fn gid_completion_before_store_buffers_and_reconciles() { use firelink_lib::queue::PendingOutcome; diff --git a/src/components/AddDownloadsModal.tsx b/src/components/AddDownloadsModal.tsx index 395d174..d030f13 100644 --- a/src/components/AddDownloadsModal.tsx +++ b/src/components/AddDownloadsModal.tsx @@ -49,9 +49,25 @@ const normalizeComparableUrl = (rawUrl: string) => { } }; +const urlsHaveDifferentHosts = (sourceUrl: string, targetUrl: string) => { + try { + return new URL(sourceUrl).hostname.toLowerCase() !== new URL(targetUrl).hostname.toLowerCase(); + } catch { + return false; + } +}; + const extensionHeaders = (context: PendingAddRequestContext | undefined) => [ context?.referer ? `Referer: ${context.referer.replace(/[\r\n]/g, '')}` : '', - context?.headers + context?.media + ? (context.headers || '') + .split(/\r?\n/) + .filter(line => { + const separator = line.indexOf(':'); + return separator < 0 || line.slice(0, separator).trim().toLowerCase() !== 'cookie'; + }) + .join('\n') + : context?.headers ].filter(Boolean).join('\n'); export const AddDownloadsModal = () => { @@ -121,9 +137,12 @@ export const AddDownloadsModal = () => { if (context) return extensionHeaders(context).trim(); return hasExtensionRequestContext ? '' : headers.trim(); }; - const cookiesForRow = (sourceUrl: string) => { + const cookiesForRow = (sourceUrl: string, targetUrl = sourceUrl) => { if (cookiesManuallyEditedRef.current) return cookies.trim(); const context = requestContextForUrl(sourceUrl); + if (context && context.cookies && urlsHaveDifferentHosts(sourceUrl, targetUrl)) { + return ''; + } if (context) return context.cookies.trim(); return hasExtensionRequestContext ? '' : cookies.trim(); }; @@ -723,7 +742,7 @@ export const AddDownloadsModal = () => { checksum: checksumEnabled && checksumValue.trim() ? `${checksumAlgo}=${checksumValue.trim()}` : undefined, - cookies: cookiesForRow(item.sourceUrl) || undefined, + cookies: cookiesForRow(item.sourceUrl, item.downloadUrl) || undefined, mirrors: mirrors.trim() || undefined, destination: useSharedDestination ? finalLocation diff --git a/src/store/useDownloadStore.test.ts b/src/store/useDownloadStore.test.ts index 833eeea..39f3231 100644 --- a/src/store/useDownloadStore.test.ts +++ b/src/store/useDownloadStore.test.ts @@ -845,7 +845,7 @@ describe('useDownloadStore', () => { referer: 'https://adult.example/watch/123', silent: false, filename: null, - headers: 'User-Agent: Firefox Test', + headers: `Cookie: stale=${'x'.repeat(64 * 1024)}\nUser-Agent: Firefox Test`, cookies: `oversized=${'x'.repeat(64 * 1024)}`, media: true }); @@ -855,6 +855,7 @@ describe('useDownloadStore', () => { expect(state.pendingAddUrls).toBe('https://adult.example/watch/123'); expect(state.pendingAddMediaUrls).toEqual(['https://adult.example/watch/123']); expect(state.pendingAddCookies).toBe(''); + expect(state.pendingAddHeaders).toBe('User-Agent: Firefox Test'); }); it('preserves extension cookies for ordinary captured downloads', async () => { diff --git a/src/store/useDownloadStore.ts b/src/store/useDownloadStore.ts index 07cffdd..094b610 100644 --- a/src/store/useDownloadStore.ts +++ b/src/store/useDownloadStore.ts @@ -67,6 +67,16 @@ const removeStaleBackendDispatch = async (id: string): Promise => { const errorMessage = (error: unknown): string => error instanceof Error ? error.message : String(error); +const stripCookieHeaders = (value: string | null | undefined): string => + (value || '') + .split(/\r?\n/) + .filter(line => { + const separator = line.indexOf(':'); + return separator < 0 || line.slice(0, separator).trim().toLowerCase() !== 'cookie'; + }) + .join('\n') + .trim(); + const speedLimitForDispatch = (itemSpeedLimit: string | undefined, globalSpeedLimit: string): string | null => { const explicitLimit = itemSpeedLimit?.trim(); if (explicitLimit) { @@ -558,12 +568,15 @@ export const useDownloadStore = create((set, get) => ({ // cookie source. Keep this frontend guard for events from older desktop or // extension builds; ordinary captured downloads retain their cookies. const cookies = request.media === true ? null : request.cookies; + const headers = request.media === true + ? stripCookieHeaders(request.headers) || null + : request.headers; get().openAddModalWithUrls( urls.join('\n'), request.referer, urls.length === 1 ? request.filename : null, - request.headers, + headers, cookies, request.media === true ); diff --git a/src/utils/clipboard.test.ts b/src/utils/clipboard.test.ts index e05e112..f2ade62 100644 --- a/src/utils/clipboard.test.ts +++ b/src/utils/clipboard.test.ts @@ -13,12 +13,13 @@ describe('clipboard URL extraction', () => { it('reads only supported, unique download URLs from clipboard text', async () => { vi.mocked(readText).mockResolvedValue( - 'https://example.com/file.zip\nhttps://example.com/file.zip ftp://example.com/file.bin mailto:user@example.com' + 'https://example.com/file.zip\nhttps://example.com/file.zip ftp://example.com/file.bin sftp://example.com/file.iso mailto:user@example.com' ); await expect(readClipboardDownloadUrls()).resolves.toEqual([ 'https://example.com/file.zip', 'ftp://example.com/file.bin', + 'sftp://example.com/file.iso', ]); }); diff --git a/src/utils/url.ts b/src/utils/url.ts index bad22da..656bb2e 100644 --- a/src/utils/url.ts +++ b/src/utils/url.ts @@ -11,7 +11,7 @@ export function extractValidDownloadUrls(text: string): string[] { for (const part of parts) { try { const url = new URL(part); - if (url.protocol === 'http:' || url.protocol === 'https:' || url.protocol === 'ftp:') { + if (url.protocol === 'http:' || url.protocol === 'https:' || url.protocol === 'ftp:' || url.protocol === 'sftp:') { urls.push(url.toString()); } } catch (e) {