From 1887843fe2494fd072f4f6c8d7090b3f2cdd1e3e Mon Sep 17 00:00:00 2001
From: Robin Appelman <robin@icewind.nl>
Date: Mon, 9 May 2022 23:30:01 +0200
Subject: [PATCH] double validate downloaded demo

---
 Cargo.lock    |  1 +
 Cargo.toml    |  1 +
 src/backup.rs | 17 ++++++++++++++---
 src/main.rs   |  2 ++
 src/store.rs  | 24 ++++++++++++++++++++++++
 5 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 564a8cd..3b0cc47 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -34,6 +34,7 @@ dependencies = [
  "dotenv",
  "futures-util",
  "main_error",
+ "md5",
  "thiserror",
  "tokio",
  "tracing",
diff --git a/Cargo.toml b/Cargo.toml
index 651416d..0396c8e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -13,6 +13,7 @@ tokio = { version = "1.17.0", features = ["rt-multi-thread", "macros"] }
 tracing = "0.1.33"
 tracing-subscriber = "0.3.11"
 futures-util = "0.3.21"
+md5 = "0.7.0"
 
 [profile.release]
 lto = true
\ No newline at end of file
diff --git a/src/backup.rs b/src/backup.rs
index bd9476d..cd63885 100644
--- a/src/backup.rs
+++ b/src/backup.rs
@@ -20,10 +20,21 @@ impl Backup {
     async fn backup_demo(&self, name: &str, demo: &Demo) -> Result<(), Error> {
         info!("backing up");
 
-        let mut file = self.store.create(name).await?;
+        {
+            let file = self.store.create(name).await?;
+            demo.save(&self.client, file).await?;
+        }
 
-        demo.save(&self.client, &mut file).await?;
-        Ok(())
+        let digest = self.store.hash(name)?;
+        if digest == demo.hash || digest == [0; 16] {
+            Ok(())
+        } else {
+            let _ = self.store.remove(name);
+            Err(Error::DigestMismatch {
+                expected: demo.hash,
+                got: digest,
+            })
+        }
     }
 
     #[instrument(skip(self))]
diff --git a/src/main.rs b/src/main.rs
index 7effa1b..aa84c19 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -16,6 +16,8 @@ pub enum Error {
     Request(#[from] std::io::Error),
     #[error(transparent)]
     Api(#[from] demostf_client::Error),
+    #[error("MD5 digest mismatch for downloaded demo, expected {expected:?}, received {got:?}")]
+    DigestMismatch { expected: [u8; 16], got: [u8; 16] },
 }
 
 #[tokio::main]
diff --git a/src/store.rs b/src/store.rs
index e8f110b..ebc646a 100644
--- a/src/store.rs
+++ b/src/store.rs
@@ -1,6 +1,8 @@
 use crate::Error;
+use md5::Context;
 use std::fs;
 use std::fs::{File, Permissions};
+use std::io::Read;
 use std::os::unix::fs::PermissionsExt;
 use std::path::{Path, PathBuf};
 
@@ -33,6 +35,28 @@ impl Store {
         fs::remove_file(self.generate_path(name))
     }
 
+    pub fn hash(&self, name: &str) -> Result<[u8; 16], Error> {
+        let path = self.generate_path(name);
+        let mut file = File::open(path)?;
+
+        let mut hash = Context::new();
+
+        let mut buff = vec![0; 1024 * 1024];
+
+        loop {
+            let read = file.read(&mut buff)?;
+
+            if read == 0 {
+                break;
+            }
+
+            let data = &buff[0..read];
+            hash.consume(&data);
+        }
+
+        Ok(hash.compute().0)
+    }
+
     fn generate_path(&self, name: &str) -> PathBuf {
         let mut path = self.basedir.clone();
         path.push(&name[0..2]);