1use 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
62const MIN_CHUNK_SIZE: u64 = 4 * 1024 * 1024;
64const MAX_RETRIES: usize = 3;
66
67pub struct AdvancedDownloader {
121 client: Client,
122 url: String,
123 output_path: String,
124 quiet_mode: bool,
125 #[allow(dead_code)]
126 proxy: ProxyConfig,
127 optimizer: Optimizer,
128 progress_callback: Option<Arc<dyn Fn(f32) + Send + Sync>>,
129 status_callback: Option<Arc<dyn Fn(String) + Send + Sync>>,
130 cancel_token: Arc<AtomicBool>,
131 expected_sha256: Option<String>,
132}
133
134impl AdvancedDownloader {
135 pub fn new(
159 url: String,
160 output_path: String,
161 quiet_mode: bool,
162 proxy_config: ProxyConfig,
163 optimizer: Optimizer,
164 ) -> Self {
165 let _is_iso = url.to_lowercase().ends_with(".iso");
166
167 let mut client_builder = Client::builder()
168 .timeout(std::time::Duration::from_secs(300))
169 .connect_timeout(std::time::Duration::from_secs(20))
170 .user_agent(concat!("KGet/", env!("CARGO_PKG_VERSION")))
171 .no_gzip()
172 .no_deflate();
173
174 if proxy_config.enabled {
175 if let Some(proxy_url) = &proxy_config.url {
176 let proxy = match proxy_config.proxy_type {
177 crate::config::ProxyType::Http => reqwest::Proxy::http(proxy_url),
178 crate::config::ProxyType::Https => reqwest::Proxy::https(proxy_url),
179 crate::config::ProxyType::Socks5 => reqwest::Proxy::all(proxy_url),
180 };
181
182 if let Ok(mut proxy) = proxy {
183 if let (Some(username), Some(password)) =
184 (&proxy_config.username, &proxy_config.password)
185 {
186 proxy = proxy.basic_auth(username, password);
187 }
188 client_builder = client_builder.proxy(proxy);
189 }
190 }
191 }
192
193 let client = client_builder
194 .build()
195 .expect("Failed to create HTTP client");
196
197 Self {
198 client,
199 url,
200 output_path,
201 quiet_mode,
202 proxy: proxy_config,
203 optimizer,
204 progress_callback: None,
205 status_callback: None,
206 cancel_token: Arc::new(AtomicBool::new(false)),
207 expected_sha256: None,
208 }
209 }
210
211 pub fn set_cancel_token(&mut self, token: Arc<AtomicBool>) {
233 self.cancel_token = token;
234 }
235
236 pub fn is_cancelled(&self) -> bool {
238 self.cancel_token.load(Ordering::Relaxed)
239 }
240
241 pub fn set_progress_callback(&mut self, callback: impl Fn(f32) + Send + Sync + 'static) {
255 self.progress_callback = Some(Arc::new(callback));
256 }
257
258 pub fn set_status_callback(&mut self, callback: impl Fn(String) + Send + Sync + 'static) {
272 self.status_callback = Some(Arc::new(callback));
273 }
274
275 pub fn set_expected_sha256(&mut self, expected_sha256: impl Into<String>) {
277 self.expected_sha256 = Some(expected_sha256.into());
278 }
279
280 fn send_status(&self, msg: &str) {
281 if let Some(cb) = &self.status_callback {
282 cb(msg.to_string());
283 }
284 if !self.quiet_mode {
285 println!("{}", msg);
286 }
287 }
288
289 pub fn download(&self) -> Result<(), Box<dyn Error + Send + Sync>> {
308 let is_iso = self.url.to_lowercase().ends_with(".iso");
309 if !self.quiet_mode {
310 println!("Starting advanced download for: {}", self.url);
311 if is_iso {
312 println!(
313 "Warning: ISO mode active. Disabling optimizations that could corrupt binary data."
314 );
315 }
316 }
317
318 let existing_size = if Path::new(&self.output_path).exists() {
320 let size = std::fs::metadata(&self.output_path)?.len();
321 if !self.quiet_mode {
322 println!("Existing file found with size: {} bytes", size);
323 }
324 Some(size)
325 } else {
326 if !self.quiet_mode {
327 println!("Output file does not exist, starting fresh download");
328 }
329 None
330 };
331
332 if !self.quiet_mode {
334 println!("Querying server for file size and range support...");
335 }
336 let (total_size, supports_range) = self.get_file_size_and_range()?;
337 if !self.quiet_mode {
338 println!("Total file size: {} bytes", total_size);
339 println!("Server supports range requests: {}", supports_range);
340 }
341
342 if let Some(size) = existing_size {
343 if size > total_size {
344 return Err("Existing file is larger than remote; aborting".into());
345 }
346 if !self.quiet_mode {
347 println!("Resuming download from byte: {}", size);
348 }
349 }
350
351 let progress = if !self.quiet_mode || self.progress_callback.is_some() {
353 let bar = ProgressBar::new(total_size);
354 if let Some(size) = existing_size {
355 bar.set_position(size);
356 }
357 if self.quiet_mode {
358 bar.set_draw_target(indicatif::ProgressDrawTarget::hidden());
359 } else {
360 bar.set_style(ProgressStyle::with_template(
361 "{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})"
362 ).unwrap().progress_chars("#>-"));
363 }
364 Some(Arc::new(Mutex::new(bar)))
365 } else {
366 None
367 };
368
369 if !self.quiet_mode {
371 println!("Preparing output file: {}", self.output_path);
372 }
373 let file = if existing_size.is_some() {
374 File::options()
375 .read(true)
376 .write(true)
377 .open(&self.output_path)?
378 } else {
379 File::create(&self.output_path)?
380 };
381 file.set_len(total_size)?;
382 if !self.quiet_mode {
383 println!("File preallocated to {} bytes", total_size);
384 }
385
386 if !supports_range {
388 if !self.quiet_mode {
389 println!("Range requests not supported, falling back to single-threaded download");
390 }
391 self.download_whole(&file, existing_size.unwrap_or(0), progress.clone())?;
392 if let Some(ref bar) = progress {
393 bar.lock()
394 .expect("Progress bar mutex was poisoned")
395 .finish_with_message("Download completed");
396 }
397 if !self.quiet_mode {
398 println!("Single-threaded download completed");
399 }
400 return Ok(());
401 }
402
403 if !self.quiet_mode {
405 println!("Calculating download chunks...");
406 }
407 let chunks = self.calculate_chunks(total_size, existing_size)?;
408 if !self.quiet_mode {
409 println!("Download will be split into {} chunks", chunks.len());
410 }
411
412 if !self.quiet_mode {
414 println!("Starting parallel chunk downloads...");
415 }
416 self.download_chunks_parallel(
417 chunks,
418 &file,
419 progress.clone(),
420 total_size,
421 existing_size.unwrap_or(0),
422 )?;
423
424 if let Some(ref bar) = progress {
425 bar.lock()
426 .expect("Progress bar mutex was poisoned")
427 .finish_with_message("Download completed");
428 }
429
430 if !self.quiet_mode || self.status_callback.is_some() {
432 if is_iso || self.expected_sha256.is_some() {
433 let should_verify = if self.status_callback.is_some() {
434 true
435 } else if self.expected_sha256.is_some() {
436 true
437 } else {
438 println!(
439 "\nThis is an ISO file. Would you like to verify its integrity? (y/N)"
440 );
441 let mut input = String::new();
442 std::io::stdin().read_line(&mut input).is_ok()
443 && input.trim().to_lowercase() == "y"
444 };
445
446 if should_verify {
447 self.verify_integrity(total_size)?;
448 }
449 } else {
450 let metadata = std::fs::metadata(&self.output_path)?;
451 if metadata.len() != total_size {
452 return Err(format!(
453 "File size mismatch: expected {} bytes, got {} bytes",
454 total_size,
455 metadata.len()
456 )
457 .into());
458 }
459 }
460 self.send_status("Advanced download completed successfully!");
461 }
462
463 Ok(())
464 }
465
466 fn get_file_size_and_range(&self) -> Result<(u64, bool), Box<dyn Error + Send + Sync>> {
467 let head_response = self.client.head(&self.url).send();
468 let Ok(response) = head_response else {
469 return self.get_file_size_with_range_probe();
470 };
471
472 if !response.status().is_success() {
473 return self.get_file_size_with_range_probe();
474 }
475
476 let content_length = response
477 .headers()
478 .get(reqwest::header::CONTENT_LENGTH)
479 .and_then(|v| v.to_str().ok())
480 .and_then(|s| s.parse::<u64>().ok());
481
482 let accepts_range = response
483 .headers()
484 .get(reqwest::header::ACCEPT_RANGES)
485 .and_then(|v| v.to_str().ok())
486 .map(|s| s.eq_ignore_ascii_case("bytes"))
487 .unwrap_or(false);
488
489 if let Some(content_length) = content_length {
490 Ok((content_length, accepts_range))
491 } else {
492 self.get_file_size_with_range_probe()
493 }
494 }
495
496 fn get_file_size_with_range_probe(&self) -> Result<(u64, bool), Box<dyn Error + Send + Sync>> {
497 let response = self
498 .client
499 .get(&self.url)
500 .header(reqwest::header::RANGE, "bytes=0-0")
501 .send()?;
502
503 if response.status() == reqwest::StatusCode::PARTIAL_CONTENT {
504 if let Some(total) = response
505 .headers()
506 .get(reqwest::header::CONTENT_RANGE)
507 .and_then(|v| v.to_str().ok())
508 .and_then(parse_content_range_total)
509 {
510 return Ok((total, true));
511 }
512 }
513
514 if response.status().is_success() {
515 if let Some(total) = response
516 .headers()
517 .get(reqwest::header::CONTENT_LENGTH)
518 .and_then(|v| v.to_str().ok())
519 .and_then(|s| s.parse::<u64>().ok())
520 {
521 return Ok((total, false));
522 }
523 }
524
525 Err("Could not determine file size".into())
526 }
527
528 fn calculate_chunks(
529 &self,
530 total_size: u64,
531 existing_size: Option<u64>,
532 ) -> Result<Vec<(u64, u64)>, Box<dyn Error + Send + Sync>> {
533 let mut chunks = Vec::new();
534 let start_from = existing_size.unwrap_or(0);
535
536 let configured_parallelism = self.optimizer.max_connections() as u64;
537 let runtime_parallelism = rayon::current_num_threads() as u64;
538 let parallelism = configured_parallelism.min(runtime_parallelism).max(1);
539 let target_chunks = parallelism.saturating_mul(2).max(2); let chunk_size = ((total_size / target_chunks).max(MIN_CHUNK_SIZE)).min(64 * 1024 * 1024);
541
542 let mut start = start_from;
543 while start < total_size {
544 let end = (start + chunk_size).min(total_size);
545 chunks.push((start, end));
546 start = end;
547 }
548
549 Ok(chunks)
550 }
551
552 fn download_whole(
553 &self,
554 file: &File,
555 offset: u64,
556 progress: Option<Arc<Mutex<ProgressBar>>>,
557 ) -> Result<(), Box<dyn Error + Send + Sync>> {
558 let response = self.client.get(&self.url).send()?;
559 if offset > 0 {
560 return Err("Server does not support range; cannot resume partial file".into());
562 }
563
564 let mut reader = BufReader::new(response);
565 let mut f = file.try_clone()?;
566 f.seek(SeekFrom::Start(0))?;
567
568 struct ProgressWriter<'a, W> {
569 inner: W,
570 progress: Option<Arc<Mutex<ProgressBar>>>,
571 callback: Option<&'a Arc<dyn Fn(f32) + Send + Sync>>,
572 }
573
574 impl<'a, W: Write> Write for ProgressWriter<'a, W> {
575 fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
576 let n = self.inner.write(buf)?;
577 if let Some(ref bar) = self.progress {
578 let guard = bar.lock().expect("Progress bar mutex was poisoned");
579 guard.inc(n as u64);
580 if let Some(cb) = self.callback {
581 let pos = guard.position();
582 let len = guard.length().unwrap_or(1);
583 drop(guard);
584 (cb)(pos as f32 / len as f32);
585 }
586 }
587 Ok(n)
588 }
589
590 fn flush(&mut self) -> std::io::Result<()> {
591 self.inner.flush()
592 }
593 }
594
595 let mut writer = ProgressWriter {
596 inner: f,
597 progress,
598 callback: self.progress_callback.as_ref(),
599 };
600 std::io::copy(&mut reader, &mut writer)?;
601
602 Ok(())
603 }
604
605 fn download_chunks_parallel(
606 &self,
607 chunks: Vec<(u64, u64)>,
608 file: &File,
609 progress: Option<Arc<Mutex<ProgressBar>>>,
610 total_size: u64,
611 initial_downloaded: u64,
612 ) -> Result<(), Box<dyn Error + Send + Sync>> {
613 let file = Arc::new(file);
614 let client = Arc::new(self.client.clone());
615 let url = Arc::new(self.url.clone());
616 let _optimizer = Arc::new(self.optimizer.clone());
617 let progress_callback = self.progress_callback.clone();
618 let cancel_token = self.cancel_token.clone();
619
620 let downloaded_bytes = Arc::new(AtomicU64::new(initial_downloaded));
622 let last_print_time = Arc::new(Mutex::new(Instant::now()));
623 let started_at = Arc::new(Instant::now());
624 let speed_limit = self.optimizer.speed_limit;
625 let quiet_mode = self.quiet_mode;
626
627 chunks.par_iter().try_for_each(|&(start, end)| {
628 if cancel_token.load(Ordering::Relaxed) {
630 return Err("Download cancelled".into());
631 }
632
633 let range = format!("bytes={}-{}", start, end - 1);
634 let range_header = reqwest::header::HeaderValue::from_str(&range)
635 .map_err(|e| format!("Invalid range header {}: {}", range, e))?;
636
637 for retry in 0..=MAX_RETRIES {
638 if cancel_token.load(Ordering::Relaxed) {
640 return Err("Download cancelled".into());
641 }
642
643 let request = client.get(url.as_str());
644 let request = request.header(reqwest::header::RANGE, range_header.clone());
645
646 match request.send() {
647 Ok(mut response) => {
648 let status = response.status();
649 if status == reqwest::StatusCode::PARTIAL_CONTENT {
650 let mut current_pos = start;
654 let mut buffer = [0u8; 16384];
655
656 while current_pos < end {
657 if cancel_token.load(Ordering::Relaxed) {
659 return Err("Download cancelled".into());
660 }
661
662 let limit = (end - current_pos).min(buffer.len() as u64);
663 let n = response.read(&mut buffer[..limit as usize])?;
664 if n == 0 {
665 break;
666 }
667
668 #[cfg(target_family = "unix")]
669 file.write_at(&buffer[..n], current_pos)?;
670
671 #[cfg(target_family = "windows")]
672 file.seek_write(&buffer[..n], current_pos)?;
673
674 current_pos += n as u64;
675
676 let new_downloaded = downloaded_bytes
678 .fetch_add(n as u64, Ordering::Relaxed)
679 + n as u64;
680
681 {
683 let mut last_time = last_print_time.lock().expect("Timer mutex was poisoned");
684 if !quiet_mode && last_time.elapsed() >= Duration::from_millis(200) {
685 let percent = (new_downloaded as f64 / total_size.max(1) as f64
686 * 100.0)
687 .min(100.0);
688 println!(
690 "PROGRESS: {:.1}% ({}/{})",
691 percent, new_downloaded, total_size
692 );
693 *last_time = Instant::now();
694 }
695 }
696
697 throttle_parallel_download(
698 new_downloaded.saturating_sub(initial_downloaded),
699 *started_at,
700 speed_limit,
701 );
702
703 if let Some(ref bar) = progress {
704 let guard = bar.lock().expect("Progress bar mutex was poisoned");
705 guard.inc(n as u64);
706 if let Some(ref cb) = progress_callback {
707 let pos = guard.position();
708 let len = guard.length().unwrap_or(1);
709 drop(guard);
710 (cb)(pos as f32 / len as f32);
711 }
712 }
713 }
714
715 return Ok::<(), Box<dyn Error + Send + Sync>>(());
716 } else if status == reqwest::StatusCode::OK {
717 return Err(format!(
718 "Server ignored range request for chunk {}-{}; refusing to write mismatched data",
719 start, end
720 )
721 .into());
722 } else if status.as_u16() == 416 {
723 if retry == MAX_RETRIES {
724 return Err(format!(
725 "Failed to download chunk {}-{}: HTTP {}",
726 start, end, status
727 )
728 .into());
729 }
730 std::thread::sleep(Duration::from_millis(250 * (retry as u64 + 1)));
731 }
732 }
733 Err(e) => {
734 if retry == MAX_RETRIES {
735 return Err(format!(
736 "Failed to download chunk {}-{}: {}",
737 start, end, e
738 )
739 .into());
740 }
741 std::thread::sleep(Duration::from_millis(250 * (retry as u64 + 1)));
742 }
743 }
744 }
745 Err(format!("Failed to download chunk {}-{} after retries", start, end).into())
746 })?;
747
748 Ok(())
749 }
750
751 fn verify_integrity(&self, expected_size: u64) -> Result<(), Box<dyn Error + Send + Sync>> {
752 let metadata = std::fs::metadata(&self.output_path)?;
753 let actual_size = metadata.len();
754
755 if actual_size != expected_size {
756 return Err(format!(
757 "File size mismatch: expected {} bytes, got {} bytes",
758 expected_size, actual_size
759 )
760 .into());
761 }
762
763 self.send_status(&format!("File size verified: {} bytes", actual_size));
764
765 self.send_status("Calculating SHA256 hash...");
767
768 let mut file = File::open(&self.output_path)?;
769 let mut hasher = Sha256::new();
770 let mut buffer = [0; 8192];
771 loop {
772 let n = file.read(&mut buffer)?;
773 if n == 0 {
774 break;
775 }
776 hasher.update(&buffer[..n]);
777 }
778 let hash = hasher.finalize();
779 let hash_hex = hex::encode(hash);
780
781 self.send_status(&format!("SHA256 hash: {}", hash_hex));
782 if let Some(expected_sha256) = &self.expected_sha256 {
783 let expected_sha256 = expected_sha256.trim().to_ascii_lowercase();
784 if hash_hex != expected_sha256 {
785 return Err(format!(
786 "SHA256 mismatch: expected {}, got {}",
787 expected_sha256, hash_hex
788 )
789 .into());
790 }
791 self.send_status("SHA256 matches expected hash.");
792 }
793 self.send_status("Integrity check passed - file is not corrupted");
794
795 Ok(())
796 }
797}
798
799fn parse_content_range_total(value: &str) -> Option<u64> {
800 let (_, total) = value.rsplit_once('/')?;
801 if total == "*" {
802 return None;
803 }
804 total.parse::<u64>().ok()
805}
806
807fn throttle_parallel_download(downloaded_this_run: u64, started_at: Instant, speed_limit: Option<u64>) {
808 let Some(limit) = speed_limit else { return };
809 if limit == 0 {
810 return;
811 }
812
813 let expected_elapsed = Duration::from_secs_f64(downloaded_this_run as f64 / limit as f64);
814 let actual_elapsed = started_at.elapsed();
815 if expected_elapsed > actual_elapsed {
816 std::thread::sleep(expected_elapsed - actual_elapsed);
817 }
818}