path: root/ozone/src/launcher/download.rs

use std::borrow::{Borrow, BorrowMut};
use std::error::Error;
use std::ffi::{OsStr, OsString};
use std::fmt::{Debug, Display, Formatter};
use std::io::{ErrorKind, Read};
use std::marker::PhantomData;
use std::path::{Path, PathBuf};
use async_trait::async_trait;
use futures::{stream, StreamExt, TryFutureExt, TryStreamExt};
use log::{debug, warn};
use reqwest::{Client, RequestBuilder, Response, StatusCode};
use sha1_smol::{Digest, Sha1};
use tokio::fs::File;
use tokio::io::AsyncWriteExt;
use crate::util;
use crate::util::IntegrityError;

#[async_trait]
pub trait Download: Debug + Display {
    // return Ok(None) to skip downloading this file
    async fn prepare(&mut self, client: &Client) -> 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>>;

    async fn verify_offline(&self) -> Result<(), Box<dyn Error>>;
}

pub trait FileDownload: Download {
    fn get_path(&self) -> &Path;
}

pub struct MultiDownloader<'d, D, DR, I>
where
    D: Download + ?Sized + 'd,
    DR: BorrowMut<D> + 'd,
    I: Iterator<Item = &'d mut DR>
{
    jobs: I,
    nconcurrent: usize,
    _phantom: PhantomData<&'d D>
}

#[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 {
    phase: Phase,
    inner: Box<dyn Error>,
    job: String
}

impl Debug for PhaseDownloadError {
    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 Display for PhaseDownloadError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "error while {} ({}): {}", self.phase, self.job, self.inner)
    }
}

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

impl PhaseDownloadError {
    fn new(phase: Phase, inner: Box<dyn Error>, job: &(impl Download + ?Sized)) -> Self {
        PhaseDownloadError {
            phase, inner, job: job.to_string()
        }
    }
}

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

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

        MultiDownloader {
            jobs,
            nconcurrent: n,
            _phantom: PhantomData
        }
    }

    pub async fn perform(self, client: &'j Client) -> Result<(), PhaseDownloadError> {
        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))
                }
            }
        }

        stream::iter(self.jobs.map(Ok))
            .try_for_each_concurrent(self.nconcurrent, |job_| async move {
                let job = job_.borrow_mut();
                
                let Some(rq) = map_err!(job.prepare(client).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, job))?;

                Ok(())
            }).await
    }
}

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: impl Into<String>, path: impl Into<PathBuf>, expect_size: Option<usize>, expect_sha1: Option<Digest>) -> VerifiedDownload {
        VerifiedDownload {
            url: url.into(),
            path: path.into(),
            
            expect_size,
            expect_sha1,
            
            file: None,
            sha1: Sha1::new(),
            tally: 0
        }
    }

    #[allow(dead_code)] // these are my emotional support functions
    pub fn with_size(mut self, expect: usize) -> VerifiedDownload {
        self.expect_size = Some(expect);
        self
    }

    #[allow(dead_code)]
    pub fn with_sha1(mut self, expect: Digest) -> VerifiedDownload {
        self.expect_sha1.replace(expect);
        self
    }

    #[allow(dead_code)]
    pub fn get_url(&self) -> &str {
        &self.url
    }

    #[allow(dead_code)]
    pub fn get_expect_size(&self) -> Option<usize> {
        self.expect_size
    }

    #[allow(dead_code)]
    pub fn get_expect_sha1(&self) -> Option<Digest> {
        self.expect_sha1
    }
}

#[async_trait]
impl Download for VerifiedDownload {
    async fn prepare(&mut self, client: &Client) -> Result<Option<RequestBuilder>, Box<dyn Error>> {
        if !util::should_download(&self.path, self.expect_size, self.expect_sha1).await? {
            return Ok(None)
        }

        // potentially racy to close the file and reopen it... :/
        self.file = Some(File::create(&self.path).await?);

        Ok(Some(client.get(&self.url)))
    }

    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(())
    }

    async fn verify_offline(&self) -> Result<(), Box<dyn Error>> {
        debug!("Verifying download {}", self.path.display());
        util::verify_file(&self.path, self.expect_size, self.expect_sha1).err_into().await
    }
}

impl FileDownload for VerifiedDownload {
    fn get_path(&self) -> &Path {
        &self.path
    }
}

pub struct RemoteChecksumDownload {
    url: String,
    check_url: String,

    expect_sha1: Option<Digest>,

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

impl RemoteChecksumDownload {
    pub fn new(url: impl Into<String>, check_url: impl Into<String>, path: impl Into<PathBuf>) -> RemoteChecksumDownload {
        RemoteChecksumDownload {
            url: url.into(),
            check_url: check_url.into(),
            path: path.into(),

            expect_sha1: None,

            file: None,
            sha1: Sha1::new(),
            tally: 0
        }
    }

    fn get_sha1_path(&self) -> PathBuf {
        [self.path.as_os_str(), OsStr::new(".sha1")].into_iter().collect::<OsString>().into()
    }
}

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

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

const MAX_DIGEST_RESPONSE_LEN: usize = 1024; // really we only need 40 bytes lol
async fn read_response_digest(res: Response) -> Result<Digest, Box<dyn Error>> {
    res.bytes_stream().err_into::<Box<dyn Error>>().try_fold(String::new(), |mut dig_string, bytes| async move {
        if bytes.len() > MAX_DIGEST_RESPONSE_LEN - dig_string.len() {
            Err(format!("oversize digest response ({} > {MAX_DIGEST_RESPONSE_LEN})", dig_string.len() + bytes.len()).into())
        } else {
            (&*bytes).read_to_string(&mut dig_string)?;
            Ok(dig_string)
        }
    }).await?.parse().map_err(Into::into)
}

#[async_trait]
impl Download for RemoteChecksumDownload {
    async fn prepare(&mut self, client: &Client) -> Result<Option<RequestBuilder>, Box<dyn Error>> {
        debug!("Downloading SHA1 fingerprint for {} from {}", self.path.display(), self.check_url);
        let res = client.get(self.check_url.as_str()).send().await
            .inspect_err(|e| warn!("Failed to download SHA1 fingerprint for {} (from {}): {}", self.path.display(), self.check_url, e))?;

        self.expect_sha1 = if res.status() == StatusCode::OK {
            read_response_digest(res).await
                .inspect_err(|e| warn!("Error downloading SHA1 fingerprint for {} (from {}): {}", self.path.display(), self.check_url, e))
                .ok()
        } else {
            warn!("Cannot download SHA1 fingerprint for {} (from {}): received status code {}", self.path.display(), self.check_url, res.status());
            None
        };

        if let Some(dig) = self.expect_sha1 {
            let _ = tokio::fs::write(self.get_sha1_path(), dig.to_string()).await
                .inspect_err(|e| warn!("Failed to save SHA1 fingerprint for {} - integrity will not be checked while offline: {e}", self.path.display()));
        }

        if !util::should_download(&self.path, None, self.expect_sha1).await? {
            return Ok(None);
        }

        self.file = Some(File::create(&self.path).await?);

        Ok(Some(client.get(self.url.as_str())))
    }

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

        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!("Failed to download {}: sha1 mismatch (exp {}, got {}).", self.path.display(), digest, d);
                return Err(IntegrityError::Sha1Mismatch { expect: d, actual: digest }.into());
            }
        }

        debug!("Successfully downloaded {} ({} bytes).", self.path.display(), self.tally);

        drop(self.file.take().unwrap());

        Ok(())
    }

    async fn verify_offline(&self) -> Result<(), Box<dyn Error>> {
        // try to load fingerprint from file
        let digest_str = match tokio::fs::read_to_string(self.get_sha1_path()).await {
            Ok(s) => s,
            Err(e) if e.kind() == ErrorKind::NotFound => {
                warn!("Unable to read sha1 fingerprint to verify {} - have to assume it's good: {e}", self.path.display());
                return Ok(());
            },
            Err(e) => return Err(e.into())
        };

        let digest: Digest = digest_str.parse()
            .inspect_err(|e| warn!("Malformed sha1 fingerprint for file {}: {e}", self.path.display()))
            .map_err::<Box<dyn Error>, _>(Into::into)?;

        util::verify_file(&self.path, None, Some(digest)).err_into().await
    }
}

impl FileDownload for RemoteChecksumDownload {
    fn get_path(&self) -> &Path {
        &self.path
    }
}

pub async fn verify_files<'d, D, DR>(files: impl Iterator<Item = &'d DR>) -> Result<(), Box<dyn Error>>
where
    D: Download + ?Sized,
    DR: Borrow<D> + 'd
{
    stream::iter(files.map(Ok))
        .try_for_each_concurrent(32, |d| async move {
            d.borrow().verify_offline().await
        })
        .await
}