kget/
advanced_download.rs

1//! Advanced parallel download functionality with resume support.
2//!
3//! The [`AdvancedDownloader`] provides high-performance downloads using:
4//! - **Parallel connections**: Split files into chunks downloaded simultaneously
5//! - **Resume support**: Continue interrupted downloads from where they left off
6//! - **Progress callbacks**: Real-time progress and status updates
7//! - **Cancellation**: Graceful download cancellation via atomic tokens
8//!
9//! # Example
10//!
11//! ```rust,no_run
12//! use kget::{AdvancedDownloader, ProxyConfig, Optimizer};
13//!
14//! let mut downloader = AdvancedDownloader::new(
15//!     "https://releases.ubuntu.com/22.04/ubuntu-22.04-desktop-amd64.iso".to_string(),
16//!     "ubuntu.iso".to_string(),
17//!     false,
18//!     ProxyConfig::default(),
19//!     Optimizer::new(),
20//! );
21//!
22//! // Set progress callback (0.0 to 1.0)
23//! downloader.set_progress_callback(|progress| {
24//!     println!("Progress: {:.1}%", progress * 100.0);
25//! });
26//!
27//! // Set status callback for messages
28//! downloader.set_status_callback(|msg| {
29//!     println!("Status: {}", msg);
30//! });
31//!
32//! // Start download
33//! downloader.download().unwrap();
34//! ```
35//!
36//! # Parallel Downloads
37//!
38//! The downloader automatically determines the optimal number of connections
39//! based on the [`Optimizer`] configuration. For large files,
40//! this can provide significant speed improvements.
41
42use crate::config::ProxyConfig;
43use crate::optimization::Optimizer;
44use hex;
45use indicatif::{ProgressBar, ProgressStyle};
46use rayon::prelude::*;
47use reqwest::blocking::Client;
48use sha2::{Digest, Sha256};
49use std::error::Error;
50use std::fs::File;
51use std::io::{BufReader, Read, Seek, SeekFrom, Write};
52use std::path::Path;
53use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
54use std::sync::{Arc, Mutex};
55use std::time::{Duration, Instant};
56
57#[cfg(target_family = "unix")]
58use std::os::unix::fs::FileExt;
59#[cfg(target_family = "windows")]
60use std::os::windows::fs::FileExt;
61
62/// Minimum chunk size for parallel downloads (4 MB)
63const MIN_CHUNK_SIZE: u64 = 4 * 1024 * 1024;
64/// Maximum retry attempts per chunk
65const MAX_RETRIES: usize = 3;
66
67/// Controls how [`AdvancedDownloader`] behaves when user interaction would otherwise be required.
68///
69/// Non-interactive callers (library, `--jsonl`, automation) must use
70/// [`AlwaysResume`](ResumePolicy::AlwaysResume) or [`AlwaysRestart`](ResumePolicy::AlwaysRestart)
71/// to avoid blocking on stdin.
72#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
73pub enum ResumePolicy {
74    /// Block on stdin for user confirmation (interactive CLI only).
75    #[default]
76    Ask,
77    /// Proceed automatically without prompting; verifies integrity when a hash is set.
78    AlwaysResume,
79    /// Always restart from scratch, discarding any existing partial file.
80    AlwaysRestart,
81}
82
83/// Token bucket for global download rate limiting across parallel threads.
84///
85/// Unlike per-thread throttling, this bounds aggregate throughput regardless of
86/// how many threads are active.
87struct TokenBucket {
88    limit_bps: u64,
89    tokens: f64,
90    last_refill: Instant,
91}
92
93impl TokenBucket {
94    fn new(limit_bps: u64) -> Self {
95        Self {
96            limit_bps,
97            tokens: limit_bps as f64,
98            last_refill: Instant::now(),
99        }
100    }
101
102    /// Consume `n` bytes from the bucket and return how long to sleep before writing.
103    ///
104    /// The caller must release the lock before sleeping.
105    fn consume(&mut self, n: u64) -> Option<Duration> {
106        let elapsed = self.last_refill.elapsed();
107        let refill = elapsed.as_secs_f64() * self.limit_bps as f64;
108        self.tokens = (self.tokens + refill).min(self.limit_bps as f64);
109        self.last_refill = Instant::now();
110
111        if self.tokens >= n as f64 {
112            self.tokens -= n as f64;
113            None
114        } else {
115            let deficit = n as f64 - self.tokens;
116            self.tokens = 0.0;
117            Some(Duration::from_secs_f64(deficit / self.limit_bps as f64))
118        }
119    }
120}
121
122/// High-performance downloader with parallel connections and resume support.
123///
124/// `AdvancedDownloader` is the recommended way to download large files. It provides:
125///
126/// - **Parallel chunk downloads**: Splits files into segments downloaded simultaneously
127/// - **Automatic resume**: Detects existing partial files and resumes from last position
128/// - **Server compatibility**: Falls back to single-stream if server doesn't support ranges
129/// - **ISO optimization**: Disables compression for binary files to prevent corruption
130/// - **Progress tracking**: Real-time callbacks for UI integration
131/// - **Cancellation support**: Stop downloads gracefully via atomic cancel token
132///
133/// # Example
134///
135/// ```rust,no_run
136/// use kget::{AdvancedDownloader, ProxyConfig, Optimizer};
137///
138/// let downloader = AdvancedDownloader::new(
139///     "https://example.com/large-file.zip".to_string(),
140///     "large-file.zip".to_string(),
141///     false,  // quiet_mode
142///     ProxyConfig::default(),
143///     Optimizer::new(),
144/// );
145///
146/// downloader.download().expect("Download failed");
147/// ```
148///
149/// # With Progress Tracking
150///
151/// ```rust,no_run
152/// use kget::{AdvancedDownloader, ProxyConfig, Optimizer};
153///
154/// let mut dl = AdvancedDownloader::new(
155///     "https://example.com/file.iso".to_string(),
156///     "file.iso".to_string(),
157///     true, // quiet mode (no stdout)
158///     ProxyConfig::default(),
159///     Optimizer::new(),
160/// );
161///
162/// dl.set_progress_callback(|p| {
163///     // p is 0.0 to 1.0
164///     update_ui_progress(p);
165/// });
166///
167/// dl.set_status_callback(|msg| println!("{}", msg));
168///
169/// dl.download().unwrap();
170///
171/// fn update_ui_progress(p: f32) {
172///     // Update your UI here
173/// }
174/// ```
175pub struct AdvancedDownloader {
176    client: Client,
177    url: String,
178    output_path: String,
179    quiet_mode: bool,
180    #[allow(dead_code)]
181    proxy: ProxyConfig,
182    optimizer: Optimizer,
183    progress_callback: Option<Arc<dyn Fn(f32) + Send + Sync>>,
184    status_callback: Option<Arc<dyn Fn(String) + Send + Sync>>,
185    cancel_token: Arc<AtomicBool>,
186    expected_sha256: Option<String>,
187    extra_headers: Vec<(String, String)>,
188    resume_policy: ResumePolicy,
189}
190
191impl AdvancedDownloader {
192    /// Create a new `AdvancedDownloader` instance.
193    ///
194    /// # Arguments
195    ///
196    /// * `url` - URL to download from
197    /// * `output_path` - Local path for the downloaded file
198    /// * `quiet_mode` - If true, suppress console output
199    /// * `proxy_config` - Proxy settings (use `ProxyConfig::default()` for direct connection)
200    /// * `optimizer` - Optimizer for connection settings
201    ///
202    /// # Example
203    ///
204    /// ```rust,no_run
205    /// use kget::{AdvancedDownloader, ProxyConfig, Optimizer};
206    ///
207    /// let dl = AdvancedDownloader::new(
208    ///     "https://example.com/file.zip".to_string(),
209    ///     "./downloads/file.zip".to_string(),
210    ///     false,
211    ///     ProxyConfig::default(),
212    ///     Optimizer::new(),
213    /// );
214    /// ```
215    pub fn new(
216        url: String,
217        output_path: String,
218        quiet_mode: bool,
219        proxy_config: ProxyConfig,
220        optimizer: Optimizer,
221    ) -> Result<Self, Box<dyn Error + Send + Sync>> {
222        let _is_iso = url.to_lowercase().ends_with(".iso");
223
224        let mut client_builder = Client::builder()
225            .timeout(std::time::Duration::from_secs(300))
226            .connect_timeout(std::time::Duration::from_secs(20))
227            .user_agent(concat!("KGet/", env!("CARGO_PKG_VERSION")))
228            .no_gzip()
229            .no_deflate();
230
231        if proxy_config.enabled {
232            if let Some(proxy_url) = &proxy_config.url {
233                let proxy = match proxy_config.proxy_type {
234                    crate::config::ProxyType::Http => reqwest::Proxy::http(proxy_url),
235                    crate::config::ProxyType::Https => reqwest::Proxy::https(proxy_url),
236                    crate::config::ProxyType::Socks5 => reqwest::Proxy::all(proxy_url),
237                };
238
239                if let Ok(mut proxy) = proxy {
240                    if let (Some(username), Some(password)) =
241                        (&proxy_config.username, &proxy_config.password)
242                    {
243                        proxy = proxy.basic_auth(username, password);
244                    }
245                    client_builder = client_builder.proxy(proxy);
246                }
247            }
248        }
249
250        let client = client_builder
251            .build()
252            .map_err(|e| format!("Failed to create HTTP client: {}", e))?;
253
254        Ok(Self {
255            client,
256            url,
257            output_path,
258            quiet_mode,
259            proxy: proxy_config,
260            optimizer,
261            progress_callback: None,
262            status_callback: None,
263            cancel_token: Arc::new(AtomicBool::new(false)),
264            expected_sha256: None,
265            extra_headers: Vec::new(),
266            resume_policy: ResumePolicy::default(),
267        })
268    }
269
270    /// Set a custom cancellation token for graceful download interruption.
271    ///
272    /// When the token is set to `true`, the download will stop at the next
273    /// checkpoint and return an error.
274    ///
275    /// # Example
276    ///
277    /// ```rust,no_run
278    /// use kget::{AdvancedDownloader, ProxyConfig, Optimizer};
279    /// use std::sync::Arc;
280    /// use std::sync::atomic::AtomicBool;
281    ///
282    /// let cancel = Arc::new(AtomicBool::new(false));
283    /// let mut dl = AdvancedDownloader::new(/* ... */
284    /// #    "".to_string(), "".to_string(), false, ProxyConfig::default(), Optimizer::new()
285    /// );
286    /// dl.set_cancel_token(cancel.clone());
287    ///
288    /// // In another thread:
289    /// // cancel.store(true, std::sync::atomic::Ordering::Relaxed);
290    /// ```
291    pub fn set_cancel_token(&mut self, token: Arc<AtomicBool>) {
292        self.cancel_token = token;
293    }
294
295    /// Check if the download has been cancelled.
296    pub fn is_cancelled(&self) -> bool {
297        self.cancel_token.load(Ordering::Relaxed)
298    }
299
300    /// Set a callback for progress updates.
301    ///
302    /// The callback receives a value from 0.0 (start) to 1.0 (complete).
303    ///
304    /// # Example
305    ///
306    /// ```rust,no_run
307    /// # use kget::{AdvancedDownloader, ProxyConfig, Optimizer};
308    /// # let mut dl = AdvancedDownloader::new("".to_string(), "".to_string(), false, ProxyConfig::default(), Optimizer::new());
309    /// dl.set_progress_callback(|progress| {
310    ///     println!("Downloaded: {:.1}%", progress * 100.0);
311    /// });
312    /// ```
313    pub fn set_progress_callback(&mut self, callback: impl Fn(f32) + Send + Sync + 'static) {
314        self.progress_callback = Some(Arc::new(callback));
315    }
316
317    /// Set a callback for status messages.
318    ///
319    /// Receives human-readable status updates during the download.
320    ///
321    /// # Example
322    ///
323    /// ```rust,no_run
324    /// # use kget::{AdvancedDownloader, ProxyConfig, Optimizer};
325    /// # let mut dl = AdvancedDownloader::new("".to_string(), "".to_string(), false, ProxyConfig::default(), Optimizer::new());
326    /// dl.set_status_callback(|msg| {
327    ///     println!("Download status: {}", msg);
328    /// });
329    /// ```
330    pub fn set_status_callback(&mut self, callback: impl Fn(String) + Send + Sync + 'static) {
331        self.status_callback = Some(Arc::new(callback));
332    }
333
334    /// Set an expected SHA-256 hash for automatic verification after download.
335    pub fn set_expected_sha256(&mut self, expected_sha256: impl Into<String>) {
336        self.expected_sha256 = Some(expected_sha256.into());
337    }
338
339    /// Set extra HTTP headers sent with every request (e.g. `Referer` or `Cookie`).
340    pub fn set_extra_headers(&mut self, headers: Vec<(String, String)>) {
341        self.extra_headers = headers;
342    }
343
344    /// Set how the downloader handles interactive prompts.
345    ///
346    /// Library and automation callers must set [`ResumePolicy::AlwaysResume`] to
347    /// avoid blocking on stdin.
348    pub fn set_resume_policy(&mut self, policy: ResumePolicy) {
349        self.resume_policy = policy;
350    }
351
352    /// Apply `self.extra_headers` to a request builder.
353    fn apply_headers(&self, mut req: reqwest::blocking::RequestBuilder) -> reqwest::blocking::RequestBuilder {
354        for (name, value) in &self.extra_headers {
355            if let (Ok(n), Ok(v)) = (
356                reqwest::header::HeaderName::from_bytes(name.as_bytes()),
357                reqwest::header::HeaderValue::from_str(value),
358            ) {
359                req = req.header(n, v);
360            }
361        }
362        req
363    }
364
365    fn send_status(&self, msg: &str) {
366        if let Some(cb) = &self.status_callback {
367            cb(msg.to_string());
368        }
369        if !self.quiet_mode {
370            println!("{}", msg);
371        }
372    }
373
374    /// Start the download.
375    ///
376    /// This method:
377    /// 1. Checks for existing partial file (resume support)
378    /// 2. Queries server for file size and range support
379    /// 3. Downloads using parallel connections if supported
380    /// 4. Falls back to single-stream if server doesn't support ranges
381    ///
382    /// # Returns
383    ///
384    /// Returns `Ok(())` on successful download, or an error if the download fails.
385    ///
386    /// # Errors
387    ///
388    /// - Network connection failures
389    /// - Existing file larger than remote (corrupted state)
390    /// - Cancellation via cancel token
391    /// - Disk I/O errors
392    pub fn download(&self) -> Result<(), Box<dyn Error + Send + Sync>> {
393        let is_iso = self.url.to_lowercase().ends_with(".iso");
394        if !self.quiet_mode {
395            println!("Starting advanced download for: {}", self.url);
396            if is_iso {
397                println!(
398                    "Warning: ISO mode active. Disabling optimizations that could corrupt binary data."
399                );
400            }
401        }
402
403        // Verify if the output path is valid
404        let existing_size = if Path::new(&self.output_path).exists() {
405            let size = std::fs::metadata(&self.output_path)?.len();
406            if !self.quiet_mode {
407                println!("Existing file found with size: {} bytes", size);
408            }
409            Some(size)
410        } else {
411            if !self.quiet_mode {
412                println!("Output file does not exist, starting fresh download");
413            }
414            None
415        };
416
417        // Get the total file size and range support
418        if !self.quiet_mode {
419            println!("Querying server for file size and range support...");
420        }
421        let (total_size, supports_range) = self.get_file_size_and_range()?;
422        if !self.quiet_mode {
423            println!("Total file size: {} bytes", total_size);
424            println!("Server supports range requests: {}", supports_range);
425        }
426
427        if let Some(size) = existing_size {
428            if size > total_size {
429                return Err("Existing file is larger than remote; aborting".into());
430            }
431            if !self.quiet_mode {
432                println!("Resuming download from byte: {}", size);
433            }
434        }
435
436        // AlwaysRestart: discard any existing partial file and start fresh.
437        let existing_size = if self.resume_policy == ResumePolicy::AlwaysRestart {
438            None
439        } else {
440            existing_size
441        };
442
443        // Create a progress bar if not quiet or if we have a callback
444        let progress = if !self.quiet_mode || self.progress_callback.is_some() {
445            let bar = ProgressBar::new(total_size);
446            if let Some(size) = existing_size {
447                bar.set_position(size);
448            }
449            if self.quiet_mode {
450                bar.set_draw_target(indicatif::ProgressDrawTarget::hidden());
451            } else {
452                bar.set_style(ProgressStyle::with_template(
453                    "{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})"
454                ).unwrap().progress_chars("#>-"));
455            }
456            Some(Arc::new(Mutex::new(bar)))
457        } else {
458            None
459        };
460
461        // Create or open the output file and preallocate
462        if !self.quiet_mode {
463            println!("Preparing output file: {}", self.output_path);
464        }
465        let file = if existing_size.is_some() {
466            File::options()
467                .read(true)
468                .write(true)
469                .open(&self.output_path)?
470        } else {
471            File::create(&self.output_path)?
472        };
473
474        // If range not supported, do a single download (no preallocation required)
475        if !supports_range {
476            if !self.quiet_mode {
477                println!("Range requests not supported, falling back to single-threaded download");
478            }
479            self.download_whole(&file, existing_size.unwrap_or(0), progress.clone())?;
480            if let Some(ref bar) = progress {
481                bar.lock()
482                    .expect("Progress bar mutex was poisoned")
483                    .finish_with_message("Download completed");
484            }
485            if !self.quiet_mode {
486                println!("Single-threaded download completed");
487            }
488            return Ok(());
489        }
490
491        // Range supported: preallocate file so parallel chunk writes land at the right offsets.
492        file.set_len(total_size)?;
493        if !self.quiet_mode {
494            println!("File preallocated to {} bytes", total_size);
495        }
496
497        // Calculate chunks for parallel download
498        if !self.quiet_mode {
499            println!("Calculating download chunks...");
500        }
501        let chunks = self.calculate_chunks(total_size, existing_size)?;
502        if !self.quiet_mode {
503            println!("Download will be split into {} chunks", chunks.len());
504        }
505
506        // Build a global token bucket so the aggregate rate across all threads stays at the limit.
507        let throttle_bucket = self.optimizer.speed_limit.map(|limit| {
508            Arc::new(Mutex::new(TokenBucket::new(limit)))
509        });
510
511        // Download parallel chunks
512        if !self.quiet_mode {
513            println!("Starting parallel chunk downloads...");
514        }
515        self.download_chunks_parallel(
516            chunks,
517            &file,
518            progress.clone(),
519            total_size,
520            existing_size.unwrap_or(0),
521            throttle_bucket,
522        )?;
523
524        if let Some(ref bar) = progress {
525            bar.lock()
526                .expect("Progress bar mutex was poisoned")
527                .finish_with_message("Download completed");
528        }
529
530        // Verify download integrity
531        if !self.quiet_mode || self.status_callback.is_some() {
532            if is_iso || self.expected_sha256.is_some() {
533                let should_verify = if self.status_callback.is_some()
534                    || self.expected_sha256.is_some()
535                {
536                    true
537                } else {
538                    match self.resume_policy {
539                        ResumePolicy::Ask => {
540                            println!(
541                                "\nThis is an ISO file. Would you like to verify its integrity? (y/N)"
542                            );
543                            let mut input = String::new();
544                            std::io::stdin().read_line(&mut input).is_ok()
545                                && input.trim().to_lowercase() == "y"
546                        }
547                        ResumePolicy::AlwaysResume => true,
548                        ResumePolicy::AlwaysRestart => false,
549                    }
550                };
551
552                if should_verify {
553                    self.verify_integrity(total_size)?;
554                }
555            } else {
556                let metadata = std::fs::metadata(&self.output_path)?;
557                if metadata.len() != total_size {
558                    return Err(format!(
559                        "File size mismatch: expected {} bytes, got {} bytes",
560                        total_size,
561                        metadata.len()
562                    )
563                    .into());
564                }
565            }
566            self.send_status("Advanced download completed successfully!");
567        }
568
569        Ok(())
570    }
571
572    fn get_file_size_and_range(&self) -> Result<(u64, bool), Box<dyn Error + Send + Sync>> {
573        let head_response = self.apply_headers(self.client.head(&self.url)).send();
574        let Ok(response) = head_response else {
575            return self.get_file_size_with_range_probe();
576        };
577
578        if !response.status().is_success() {
579            return self.get_file_size_with_range_probe();
580        }
581
582        let content_length = response
583            .headers()
584            .get(reqwest::header::CONTENT_LENGTH)
585            .and_then(|v| v.to_str().ok())
586            .and_then(|s| s.parse::<u64>().ok());
587
588        let accepts_range = response
589            .headers()
590            .get(reqwest::header::ACCEPT_RANGES)
591            .and_then(|v| v.to_str().ok())
592            .map(|s| s.eq_ignore_ascii_case("bytes"))
593            .unwrap_or(false);
594
595        if let Some(content_length) = content_length {
596            Ok((content_length, accepts_range))
597        } else {
598            self.get_file_size_with_range_probe()
599        }
600    }
601
602    fn get_file_size_with_range_probe(&self) -> Result<(u64, bool), Box<dyn Error + Send + Sync>> {
603        let response = self
604            .apply_headers(self.client.get(&self.url))
605            .header(reqwest::header::RANGE, "bytes=0-0")
606            .send()?;
607
608        if response.status() == reqwest::StatusCode::PARTIAL_CONTENT {
609            if let Some(total) = response
610                .headers()
611                .get(reqwest::header::CONTENT_RANGE)
612                .and_then(|v| v.to_str().ok())
613                .and_then(parse_content_range_total)
614            {
615                return Ok((total, true));
616            }
617        }
618
619        if response.status().is_success() {
620            if let Some(total) = response
621                .headers()
622                .get(reqwest::header::CONTENT_LENGTH)
623                .and_then(|v| v.to_str().ok())
624                .and_then(|s| s.parse::<u64>().ok())
625            {
626                return Ok((total, false));
627            }
628        }
629
630        Err("Could not determine file size".into())
631    }
632
633    fn calculate_chunks(
634        &self,
635        total_size: u64,
636        existing_size: Option<u64>,
637    ) -> Result<Vec<(u64, u64)>, Box<dyn Error + Send + Sync>> {
638        let mut chunks = Vec::new();
639        let start_from = existing_size.unwrap_or(0);
640
641        let configured_parallelism = self.optimizer.max_connections() as u64;
642        let runtime_parallelism = rayon::current_num_threads() as u64;
643        let parallelism = configured_parallelism.min(runtime_parallelism).max(1);
644        let target_chunks = parallelism.saturating_mul(2).max(2); // Keep workers fed without overwhelming servers.
645        let chunk_size = ((total_size / target_chunks).max(MIN_CHUNK_SIZE)).min(64 * 1024 * 1024);
646
647        let mut start = start_from;
648        while start < total_size {
649            let end = (start + chunk_size).min(total_size);
650            chunks.push((start, end));
651            start = end;
652        }
653
654        Ok(chunks)
655    }
656
657    fn download_whole(
658        &self,
659        file: &File,
660        offset: u64,
661        progress: Option<Arc<Mutex<ProgressBar>>>,
662    ) -> Result<(), Box<dyn Error + Send + Sync>> {
663        let response = self.apply_headers(self.client.get(&self.url)).send()?;
664        if offset > 0 {
665            // Resume not possible without range; warn
666            return Err("Server does not support range; cannot resume partial file".into());
667        }
668
669        let mut reader = BufReader::new(response);
670        let mut f = file.try_clone()?;
671        f.seek(SeekFrom::Start(0))?;
672
673        struct ProgressWriter<'a, W> {
674            inner: W,
675            progress: Option<Arc<Mutex<ProgressBar>>>,
676            callback: Option<&'a Arc<dyn Fn(f32) + Send + Sync>>,
677        }
678
679        impl<'a, W: Write> Write for ProgressWriter<'a, W> {
680            fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
681                let n = self.inner.write(buf)?;
682                if let Some(ref bar) = self.progress {
683                    let guard = bar.lock().expect("Progress bar mutex was poisoned");
684                    guard.inc(n as u64);
685                    if let Some(cb) = self.callback {
686                        let pos = guard.position();
687                        let len = guard.length().unwrap_or(1);
688                        drop(guard);
689                        (cb)(pos as f32 / len as f32);
690                    }
691                }
692                Ok(n)
693            }
694
695            fn flush(&mut self) -> std::io::Result<()> {
696                self.inner.flush()
697            }
698        }
699
700        let mut writer = ProgressWriter {
701            inner: f,
702            progress,
703            callback: self.progress_callback.as_ref(),
704        };
705        std::io::copy(&mut reader, &mut writer)?;
706
707        Ok(())
708    }
709
710    fn download_chunks_parallel(
711        &self,
712        chunks: Vec<(u64, u64)>,
713        file: &File,
714        progress: Option<Arc<Mutex<ProgressBar>>>,
715        total_size: u64,
716        initial_downloaded: u64,
717        throttle_bucket: Option<Arc<Mutex<TokenBucket>>>,
718    ) -> Result<(), Box<dyn Error + Send + Sync>> {
719        let file = Arc::new(file);
720        let client = Arc::new(self.client.clone());
721        let url = Arc::new(self.url.clone());
722        let progress_callback = self.progress_callback.clone();
723        let cancel_token = self.cancel_token.clone();
724        let extra_headers = Arc::new(self.extra_headers.clone());
725
726        // Shared progress counter for pipe-friendly output
727        let downloaded_bytes = Arc::new(AtomicU64::new(initial_downloaded));
728        let last_print_time = Arc::new(Mutex::new(Instant::now()));
729        let quiet_mode = self.quiet_mode;
730
731        chunks.par_iter().try_for_each(|&(start, end)| {
732            // Check for cancellation before starting chunk
733            if cancel_token.load(Ordering::Relaxed) {
734                return Err("Download cancelled".into());
735            }
736
737            let range = format!("bytes={}-{}", start, end - 1);
738            let range_header = reqwest::header::HeaderValue::from_str(&range)
739                .map_err(|e| format!("Invalid range header {}: {}", range, e))?;
740
741            for retry in 0..=MAX_RETRIES {
742                // Check for cancellation on each retry
743                if cancel_token.load(Ordering::Relaxed) {
744                    return Err("Download cancelled".into());
745                }
746
747                let mut request = client.get(url.as_str());
748                for (name, value) in extra_headers.as_ref() {
749                    if let (Ok(n), Ok(v)) = (
750                        reqwest::header::HeaderName::from_bytes(name.as_bytes()),
751                        reqwest::header::HeaderValue::from_str(value),
752                    ) {
753                        request = request.header(n, v);
754                    }
755                }
756                let request = request.header(reqwest::header::RANGE, range_header.clone());
757
758                match request.send() {
759                    Ok(mut response) => {
760                        let status = response.status();
761                        if status == reqwest::StatusCode::PARTIAL_CONTENT {
762                            // Use FileExt to write at specific offset without seeking shared cursor
763                            // This prevents race conditions when multiple threads write to the same file
764
765                            let mut current_pos = start;
766                            let mut buffer = [0u8; 16384];
767
768                            while current_pos < end {
769                                // Check for cancellation periodically during download
770                                if cancel_token.load(Ordering::Relaxed) {
771                                    return Err("Download cancelled".into());
772                                }
773
774                                let limit = (end - current_pos).min(buffer.len() as u64);
775                                let n = response.read(&mut buffer[..limit as usize])?;
776                                if n == 0 {
777                                    break;
778                                }
779
780                                #[cfg(target_family = "unix")]
781                                file.write_at(&buffer[..n], current_pos)?;
782
783                                #[cfg(target_family = "windows")]
784                                file.seek_write(&buffer[..n], current_pos)?;
785
786                                current_pos += n as u64;
787
788                                // Update shared progress counter
789                                let new_downloaded = downloaded_bytes
790                                    .fetch_add(n as u64, Ordering::Relaxed)
791                                    + n as u64;
792
793                                // Print progress periodically (every 200ms) for pipe-friendly output
794                                {
795                                    let mut last_time = last_print_time.lock().expect("Timer mutex was poisoned");
796                                    if !quiet_mode && last_time.elapsed() >= Duration::from_millis(200) {
797                                let percent = (new_downloaded as f64 / total_size.max(1) as f64
798                                            * 100.0)
799                                            .min(100.0);
800                                        // PROGRESS: format that Swift can parse
801                                        println!(
802                                            "PROGRESS: {:.1}% ({}/{})",
803                                            percent, new_downloaded, total_size
804                                        );
805                                        *last_time = Instant::now();
806                                    }
807                                }
808
809                                if let Some(ref bucket) = throttle_bucket {
810                                    let sleep_dur = {
811                                        let mut guard = bucket
812                                            .lock()
813                                            .expect("throttle bucket poisoned");
814                                        guard.consume(n as u64)
815                                    };
816                                    if let Some(dur) = sleep_dur {
817                                        std::thread::sleep(dur);
818                                    }
819                                }
820
821                                if let Some(ref bar) = progress {
822                                    let guard = bar.lock().expect("Progress bar mutex was poisoned");
823                                    guard.inc(n as u64);
824                                    if let Some(ref cb) = progress_callback {
825                                        let pos = guard.position();
826                                        let len = guard.length().unwrap_or(1);
827                                        drop(guard);
828                                        (cb)(pos as f32 / len as f32);
829                                    }
830                                }
831                            }
832
833                            return Ok::<(), Box<dyn Error + Send + Sync>>(());
834                        } else if status == reqwest::StatusCode::OK {
835                            return Err(format!(
836                                "Server ignored range request for chunk {}-{}; refusing to write mismatched data",
837                                start, end
838                            )
839                            .into());
840                        } else if status.as_u16() == 416 {
841                            if retry == MAX_RETRIES {
842                                return Err(format!(
843                                    "Failed to download chunk {}-{}: HTTP {}",
844                                    start, end, status
845                                )
846                                .into());
847                            }
848                            std::thread::sleep(Duration::from_millis(250 * (retry as u64 + 1)));
849                        }
850                    }
851                    Err(e) => {
852                        if retry == MAX_RETRIES {
853                            return Err(format!(
854                                "Failed to download chunk {}-{}: {}",
855                                start, end, e
856                            )
857                            .into());
858                        }
859                        std::thread::sleep(Duration::from_millis(250 * (retry as u64 + 1)));
860                    }
861                }
862            }
863            Err(format!("Failed to download chunk {}-{} after retries", start, end).into())
864        })?;
865
866        Ok(())
867    }
868
869    fn verify_integrity(&self, expected_size: u64) -> Result<(), Box<dyn Error + Send + Sync>> {
870        let metadata = std::fs::metadata(&self.output_path)?;
871        let actual_size = metadata.len();
872
873        if actual_size != expected_size {
874            return Err(format!(
875                "File size mismatch: expected {} bytes, got {} bytes",
876                expected_size, actual_size
877            )
878            .into());
879        }
880
881        self.send_status(&format!("File size verified: {} bytes", actual_size));
882
883        // Calculate SHA256 hash for corruption check
884        self.send_status("Calculating SHA256 hash...");
885
886        let mut file = File::open(&self.output_path)?;
887        let mut hasher = Sha256::new();
888        let mut buffer = [0; 8192];
889        loop {
890            let n = file.read(&mut buffer)?;
891            if n == 0 {
892                break;
893            }
894            hasher.update(&buffer[..n]);
895        }
896        let hash = hasher.finalize();
897        let hash_hex = hex::encode(hash);
898
899        self.send_status(&format!("SHA256 hash: {}", hash_hex));
900        if let Some(expected_sha256) = &self.expected_sha256 {
901            let expected_sha256 = expected_sha256.trim().to_ascii_lowercase();
902            if hash_hex != expected_sha256 {
903                return Err(format!(
904                    "SHA256 mismatch: expected {}, got {}",
905                    expected_sha256, hash_hex
906                )
907                .into());
908            }
909            self.send_status("SHA256 matches expected hash.");
910        }
911        self.send_status("Integrity check passed - file is not corrupted");
912
913        Ok(())
914    }
915}
916
917fn parse_content_range_total(value: &str) -> Option<u64> {
918    let (_, total) = value.rsplit_once('/')?;
919    if total == "*" {
920        return None;
921    }
922    total.parse::<u64>().ok()
923}
924
kget/advanced_download.rs

kget/
advanced_download.rs
