From f9e181df1b465010190093f5049bf7c849b0b078 Mon Sep 17 00:00:00 2001
From: link2xt <link2xt@testrun.org>
Date: Thu, 24 Oct 2024 22:28:14 +0000
Subject: [PATCH] feat: support importing gzip-compressed backups

---
 Cargo.lock           |  1 +
 Cargo.toml           |  1 +
 src/imex.rs          | 39 ++++++++++++++++++++++-----------------
 src/imex/transfer.rs |  5 +++--
 4 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 14069e547d..941568a40f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1298,6 +1298,7 @@ dependencies = [
  "anyhow",
  "async-broadcast",
  "async-channel 2.3.1",
+ "async-compression",
  "async-imap",
  "async-native-tls",
  "async-smtp",
diff --git a/Cargo.toml b/Cargo.toml
index d9142aa5b0..26affa3d7f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -42,6 +42,7 @@ anyhow = { workspace = true }
 async-broadcast = "0.7.1"
 async-channel = { workspace = true }
 async-imap = { version = "0.10.2", default-features = false, features = ["runtime-tokio", "compress"] }
+async-compression = { version = "0.4.15", default-features = false, features = ["tokio", "gzip"] }
 async-native-tls = { version = "0.5", default-features = false, features = ["runtime-tokio"] }
 async-smtp = { version = "0.9", default-features = false, features = ["runtime-tokio"] }
 async_zip = { version = "0.0.17", default-features = false, features = ["deflate", "tokio-fs"] }
diff --git a/src/imex.rs b/src/imex.rs
index f99115cb6f..2399b86a09 100644
--- a/src/imex.rs
+++ b/src/imex.rs
@@ -269,30 +269,24 @@ async fn import_backup(
         context.get_dbfile().display()
     );
 
-    import_backup_stream(context, backup_file, file_size, passphrase).await?;
+    let backup_file = ProgressReader::new(backup_file, context.clone(), file_size);
+    if backup_to_import.extension() == Some(OsStr::new("gz")) {
+        let backup_file = tokio::io::BufReader::new(backup_file);
+        let backup_file = async_compression::tokio::bufread::GzipDecoder::new(backup_file);
+        import_backup_stream(context, backup_file, passphrase).await?;
+    } else {
+        import_backup_stream(context, backup_file, passphrase).await?;
+    }
     Ok(())
 }
 
 /// Imports backup by reading a tar file from a stream.
-///
-/// `file_size` is used to calculate the progress
-/// and emit progress events.
-/// Ideally it is the sum of the entry
-/// sizes without the header overhead,
-/// but can be estimated as tar file size
-/// in which case the progress is underestimated
-/// and may not reach 99.9% by the end of import.
-/// Underestimating is better than
-/// overestimating because the progress
-/// jumps to 100% instead of getting stuck at 99.9%
-/// for some time.
 pub(crate) async fn import_backup_stream<R: tokio::io::AsyncRead + Unpin>(
     context: &Context,
     backup_file: R,
-    file_size: u64,
     passphrase: String,
 ) -> Result<()> {
-    import_backup_stream_inner(context, backup_file, file_size, passphrase)
+    import_backup_stream_inner(context, backup_file, passphrase)
         .await
         .0
 }
@@ -319,6 +313,19 @@ struct ProgressReader<R> {
 }
 
 impl<R> ProgressReader<R> {
+    /// Creates a new `ProgressReader`.
+    ///
+    /// `file_size` is used to calculate the progress
+    /// and emit progress events.
+    /// Ideally it is the sum of the entry
+    /// sizes without the header overhead,
+    /// but can be estimated as tar file size
+    /// in which case the progress is underestimated
+    /// and may not reach 99.9% by the end of import.
+    /// Underestimating is better than
+    /// overestimating because the progress
+    /// jumps to 100% instead of getting stuck at 99.9%
+    /// for some time.
     fn new(r: R, context: Context, file_size: u64) -> Self {
         Self {
             inner: r,
@@ -358,10 +365,8 @@ where
 async fn import_backup_stream_inner<R: tokio::io::AsyncRead + Unpin>(
     context: &Context,
     backup_file: R,
-    file_size: u64,
     passphrase: String,
 ) -> (Result<()>,) {
-    let backup_file = ProgressReader::new(backup_file, context.clone(), file_size);
     let mut archive = Archive::new(backup_file);
 
     let mut entries = match archive.entries() {
diff --git a/src/imex/transfer.rs b/src/imex/transfer.rs
index d280dd8f88..fb4c94ac6c 100644
--- a/src/imex/transfer.rs
+++ b/src/imex/transfer.rs
@@ -41,7 +41,7 @@ use tokio_util::sync::CancellationToken;
 
 use crate::chat::add_device_msg;
 use crate::context::Context;
-use crate::imex::BlobDirContents;
+use crate::imex::{BlobDirContents, ProgressReader};
 use crate::message::{Message, Viewtype};
 use crate::qr::Qr;
 use crate::stock_str::backup_transfer_msg_body;
@@ -310,7 +310,8 @@ pub async fn get_backup2(
     let mut file_size_buf = [0u8; 8];
     recv_stream.read_exact(&mut file_size_buf).await?;
     let file_size = u64::from_be_bytes(file_size_buf);
-    import_backup_stream(context, recv_stream, file_size, passphrase)
+    let recv_stream = ProgressReader::new(recv_stream, context.clone(), file_size);
+    import_backup_stream(context, recv_stream, passphrase)
         .await
         .context("Failed to import backup from QUIC stream")?;
     info!(context, "Finished importing backup from the stream.");