1 files changed, 152 insertions, 3 deletions

diff --git a/src/launcher/download.rs b/src/launcher/download.rs
index 3598976..ca1c57d 100644
--- a/src/launcher/download.rs
+++ b/src/launcher/download.rs
@@ -1,14 +1,18 @@
 use std::error::Error;
 use std::fmt::{Debug, Display, Formatter};
+use std::io::ErrorKind;
+use std::path::{Path, PathBuf};
 use futures::{stream, Stream, StreamExt};
 use reqwest::{Client, IntoUrl, Method, RequestBuilder};
+use sha1_smol::{Digest, Sha1};
+use tokio::fs;
+use tokio::fs::File;
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use crate::launcher::constants::USER_AGENT;
 
 pub trait Download: Debug + Display {
-    type URLType: IntoUrl;
     // return Ok(None) to skip downloading this file
-
-    fn get_url(&self) -> Self::URLType;
+    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>>;
@@ -134,3 +138,148 @@ impl<T: Download> MultiDownloader<T> {
         }).buffer_unordered(self.nconcurrent)
     }
 }
+
+#[derive(Debug)]
+pub enum IntegrityError {
+    SizeMismatch{ expect: usize, actual: usize },
+    Sha1Mismatch{ expect: Digest, actual: Digest }
+}
+
+impl Display for IntegrityError {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            IntegrityError::SizeMismatch{ expect, actual } =>
+                write!(f, "size mismatch (expect {expect} bytes, got {actual} bytes)"),
+            IntegrityError::Sha1Mismatch {expect, actual} =>
+                write!(f, "sha1 mismatch (expect {expect}, got {actual})")
+        }
+    }
+}
+
+impl Error for IntegrityError {}
+
+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.to_string_lossy())
+    }
+}
+
+impl VerifiedDownload {
+    pub fn new(url: &str, path: &Path) -> VerifiedDownload {
+        VerifiedDownload {
+            url: url.to_owned(),
+            path: path.to_owned(),
+            
+            expect_size: None,
+            expect_sha1: None,
+            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
+    }
+    
+    async fn open_output(&mut self) -> Result<(), tokio::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>> {
+        let mut file = match File::open(&self.path).await {
+            Ok(file) => file,
+            Err(e) => return if e.kind() == ErrorKind::NotFound {
+                // assume the parent folder exists (responsibility of the caller to ensure this)
+                self.open_output().await?;
+                Ok(Some(req))
+            } else {
+                Err(e.into())
+            }
+        };
+
+        let mut tally = 0usize;
+        let mut sha1 = Sha1::new();
+
+        let mut buf = [0u8; 4096];
+        loop {
+            let n = file.read(&mut buf).await?;
+            if n == 0 { break; }
+
+            tally += n;
+            sha1.update(&buf[..n]);
+        }
+
+        if self.expect_sha1.is_none_or(|d| d == sha1.digest())
+            && self.expect_size.is_none_or(|s| s == tally) {
+            return Ok(None);
+        }
+
+        drop(file);
+
+        // 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 {
+                return Err(IntegrityError::Sha1Mismatch { expect: d, actual: digest }.into());
+            }
+        } else if let Some(s) = self.expect_size {
+            if s != self.tally {
+                return Err(IntegrityError::SizeMismatch { expect: s, actual: self.tally }.into());
+            }
+        }
+
+        // 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(())
+    }
+}