path: root/src/launcher/download.rs

use std::error::Error;
use std::fmt::{Debug, Display, Formatter};
use std::io::ErrorKind;
use std::path::{Path, PathBuf};
use futures::{stream, StreamExt, TryStream, TryStreamExt};
use log::{debug, warn};
use reqwest::{Client, IntoUrl, Method, RequestBuilder};
use sha1_smol::{Digest, Sha1};
use tokio::fs;
use tokio::fs::File;
use tokio::io::{self, AsyncWriteExt};
use crate::launcher::constants::USER_AGENT;
use crate::util;
use crate::util::{FileVerifyError, IntegrityError};

pub trait Download: Debug + Display {
    // return Ok(None) to skip downloading this file
    fn get_url(&self) -> impl IntoUrl;

    async fn prepare(&mut self, req: RequestBuilder) -> Result<Option<RequestBuilder>, Box<dyn Error>>;
    async fn handle_chunk(&mut self, chunk: &[u8]) -> Result<(), Box<dyn Error>>;
    async fn finish(&mut self) -> Result<(), Box<dyn Error>>;
}

pub struct MultiDownloader<'j, T: Download + 'j, I: Iterator<Item = &'j mut T>> {
    jobs: I,
    nconcurrent: usize
}

#[derive(Debug, Clone, Copy)]
pub enum Phase {
    Prepare,
    Send,
    Receive,
    HandleChunk,
    Finish
}

impl Display for Phase {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            /* an error occurred while (present participle) ... */
            Self::Prepare => f.write_str("preparing the request"),
            Self::Send => f.write_str("sending the request"),
            Self::Receive => f.write_str("receiving response data"),
            Self::HandleChunk => f.write_str("handling response data"),
            Self::Finish => f.write_str("finishing the request"),
        }
    }
}

pub struct PhaseDownloadError<'j, T: Download> {
    phase: Phase,
    inner: Box<dyn Error>,
    job: &'j T
}

impl<'j, T: Download> Debug for PhaseDownloadError<'j, T> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PhaseDownloadError")
            .field("phase", &self.phase)
            .field("inner", &self.inner)
            .field("job", &self.job)
            .finish()
    }
}

impl<'j, T: Download> Display for PhaseDownloadError<'j, T> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "error while {} ({}): {}", self.phase, self.job, self.inner)
    }
}

impl<'j, T: Download> Error for PhaseDownloadError<'j, T> {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        Some(&*self.inner)
    }
}

impl<'j, T: Download> PhaseDownloadError<'j, T> {
    fn new(phase: Phase, inner: Box<dyn Error>, job: &'j T) -> Self {
        PhaseDownloadError {
            phase, inner, job
        }
    }
}

impl<'j, T: Download + 'j, I: Iterator<Item = &'j mut T>> MultiDownloader<'j, T, I> {
    pub fn new(jobs: I) -> MultiDownloader<'j, T, I> {
        Self::with_concurrent(jobs, 24)
    }

    pub fn with_concurrent(jobs: I, n: usize) -> MultiDownloader<'j, T, I> {
        assert!(n > 0);

        MultiDownloader {
            jobs,
            nconcurrent: n
        }
    }

    pub async fn perform(self, client: &'j Client) -> impl TryStream<Ok = (), Error = PhaseDownloadError<'j, T>> {
        stream::iter(self.jobs.into_iter()).map(move |job| Ok(async move {
            macro_rules! map_err {
                ($result:expr, $phase:expr, $job:expr) => {
                    match $result {
                        Ok(v) => v,
                        Err(e) => return Err(PhaseDownloadError::new($phase, e.into(), $job))
                    }
                }
            }

            let Some(rq) = map_err!(
            job.prepare(client.request(Method::GET, job.get_url())
                .header(reqwest::header::USER_AGENT, USER_AGENT)).await, Phase::Prepare, job) else {
                return Ok(())
            };

            let mut data = map_err!(map_err!(rq.send().await, Phase::Send, job).error_for_status(), Phase::Send, job).bytes_stream();

            while let Some(bytes) = data.next().await {
                let bytes = map_err!(bytes, Phase::Receive, job);

                map_err!(job.handle_chunk(bytes.as_ref()).await, Phase::HandleChunk, job);
            }

            job.finish().await.map_err(|e| PhaseDownloadError::new(Phase::Finish, e.into(), job))?;

            Ok(())
        })).try_buffer_unordered(self.nconcurrent)
    }
}

pub struct VerifiedDownload {
    url: String,
    expect_size: Option<usize>,
    expect_sha1: Option<Digest>,

    path: PathBuf,
    file: Option<File>,
    sha1: Sha1,
    tally: usize
}

impl Debug for VerifiedDownload {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("VerifiedDownload")
            .field("url", &self.url)
            .field("expect_size", &self.expect_size)
            .field("expect_sha1", &self.expect_sha1)
            .field("path", &self.path).finish()
    }
}

impl Display for VerifiedDownload {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "downloading {} to {}", self.url, self.path.display())
    }
}

impl VerifiedDownload {
    pub fn new(url: &str, path: &Path, expect_size: Option<usize>, expect_sha1: Option<Digest>) -> VerifiedDownload {
        VerifiedDownload {
            url: url.to_owned(),
            path: path.to_owned(),
            
            expect_size,
            expect_sha1,
            
            file: None,
            sha1: Sha1::new(),
            tally: 0
        }
    }
    
    pub fn with_size(mut self, expect: usize) -> VerifiedDownload {
        self.expect_size = Some(expect);
        self
    }
    
    pub fn with_sha1(mut self, expect: Digest) -> VerifiedDownload {
        self.expect_sha1.replace(expect);
        self
    }
    
    pub fn get_url(&self) -> &str {
        &self.url
    }
    
    pub fn get_path(&self) -> &Path {
        &self.path
    }
    
    pub fn get_expect_size(&self) -> Option<usize> {
        self.expect_size
    }
    
    pub fn get_expect_sha1(&self) -> Option<Digest> {
        self.expect_sha1
    }
    
    pub async fn make_dirs(&self) -> Result<(), io::Error> {
        fs::create_dir_all(self.path.parent().expect("download created with no containing directory (?)")).await
    }
    
    async fn open_output(&mut self) -> Result<(), io::Error> {
        self.file.replace(File::create(&self.path).await?);
        Ok(())
    }
}

impl Download for VerifiedDownload {
    fn get_url(&self) -> impl IntoUrl {
        &self.url
    }

    async fn prepare(&mut self, req: RequestBuilder) -> Result<Option<RequestBuilder>, Box<dyn Error>> {
        match util::verify_file(&self.path, self.expect_size, self.expect_sha1).await {
            Ok(()) => {
                debug!("Skipping download for file {}, integrity matches.", self.path.display());
                return Ok(None);
            },
            Err(e) => match e {
                FileVerifyError::Integrity(_, _) => {
                    warn!("Integrity error on library: {}", e);
                    
                    // try to delete the file since it's bad
                    let _ = fs::remove_file(&self.path).await
                        .map_err(|e| warn!("Error deleting corrupted/modified file {} (ignoring): {}", self.path.display(), e));
                },
                FileVerifyError::Open(_, e) => match e.kind() {
                    ErrorKind::NotFound => {
                        debug!("File {} is missing, downloading it.", self.path.display());
                    },
                    _ => return Err(e.into())
                },
                _ => return Err(e.into())
            }
        }

        // potentially racy to close the file and reopen it... :/
        self.open_output().await?;

        Ok(Some(req))
    }

    async fn handle_chunk(&mut self, chunk: &[u8]) -> Result<(), Box<dyn Error>> {
        self.file.as_mut().unwrap().write_all(chunk).await?;
        self.tally += chunk.len();
        self.sha1.update(chunk);

        Ok(())
    }

    async fn finish(&mut self) -> Result<(), Box<dyn Error>> {
        let digest = self.sha1.digest();

        if let Some(d) = self.expect_sha1 {
            if d != digest {
                debug!("Could not download {}: sha1 mismatch (exp {}, got {}).", self.path.display(), d, digest);
                return Err(IntegrityError::Sha1Mismatch { expect: d, actual: digest }.into());
            }
        } else if let Some(s) = self.expect_size {
            if s != self.tally {
                debug!("Could not download {}: size mismatch (exp {}, got {}).", self.path.display(), s, self.tally);
                return Err(IntegrityError::SizeMismatch { expect: s, actual: self.tally }.into());
            }
        }
        
        debug!("Successfully downloaded {} ({} bytes).", self.path.display(), self.tally);

        // release the file descriptor (don't want to wait until it's dropped automatically because idk when that would be)
        drop(self.file.take().unwrap());

        Ok(())
    }
}

pub async fn verify_files(files: impl Iterator<Item = &mut VerifiedDownload>) -> Result<(), FileVerifyError> {
    stream::iter(files)
        .map(|dl| Ok(async move {
            debug!("Verifying library {}", dl.get_path().display());
            util::verify_file(dl.get_path(), dl.get_expect_size(), dl.get_expect_sha1()).await
        }))
        .try_buffer_unordered(32)
        .try_fold((), |_, _| async {Ok(())})
        .await
}