combine: support combining both xz and gz tarballs

3 files changed, 69 insertions, 19 deletions

diff --git a/src/combiner.rs b/src/combiner.rs
index 653b120..b25063c 100644
--- a/src/combiner.rs
+++ b/src/combiner.rs
@@ -1,8 +1,7 @@
 use super::Scripter;
 use super::Tarballer;
-use crate::util::*;
+use crate::{compression::CompressionFormat, util::*};
 use anyhow::{bail, Context, Result};
-use flate2::read::GzDecoder;
 use std::io::{Read, Write};
 use std::path::Path;
 use tar::Archive;
@@ -59,15 +58,21 @@ impl Combiner {
             .filter(|s| !s.is_empty())
         {
             // Extract the input tarballs
-            let tar = GzDecoder::new(open_file(&input_tarball)?);
-            Archive::new(tar).unpack(&self.work_dir).with_context(|| {
-                format!(
-                    "unable to extract '{}' into '{}'",
-                    &input_tarball, self.work_dir
-                )
-            })?;
-
-            let pkg_name = input_tarball.trim_end_matches(".tar.gz");
+            let compression =
+                CompressionFormat::detect_from_path(input_tarball).ok_or_else(|| {
+                    anyhow::anyhow!("couldn't figure out the format of {}", input_tarball)
+                })?;
+            Archive::new(compression.decode(input_tarball)?)
+                .unpack(&self.work_dir)
+                .with_context(|| {
+                    format!(
+                        "unable to extract '{}' into '{}'",
+                        &input_tarball, self.work_dir
+                    )
+                })?;
+
+            let pkg_name =
+                input_tarball.trim_end_matches(&format!(".tar.{}", compression.extension()));
             let pkg_name = Path::new(pkg_name).file_name().unwrap();
             let pkg_dir = Path::new(&self.work_dir).join(&pkg_name);
 
diff --git a/src/compression.rs b/src/compression.rs
index 949c0a2..b8cfd9a 100644
--- a/src/compression.rs
+++ b/src/compression.rs
@@ -1,8 +1,8 @@
 use anyhow::{Context, Error};
-use flate2::write::GzEncoder;
+use flate2::{read::GzDecoder, write::GzEncoder};
 use rayon::prelude::*;
-use std::{io::Write, path::Path};
-use xz2::write::XzEncoder;
+use std::{io::Read, io::Write, path::Path};
+use xz2::{read::XzDecoder, write::XzEncoder};
 
 pub(crate) enum CompressionFormat {
     Gz,
@@ -10,13 +10,24 @@ pub(crate) enum CompressionFormat {
 }
 
 impl CompressionFormat {
+    pub(crate) fn detect_from_path(path: impl AsRef<Path>) -> Option<Self> {
+        match path.as_ref().extension().and_then(|e| e.to_str()) {
+            Some("gz") => Some(CompressionFormat::Gz),
+            Some("xz") => Some(CompressionFormat::Xz),
+            _ => None,
+        }
+    }
+
+    pub(crate) fn extension(&self) -> &'static str {
+        match self {
+            CompressionFormat::Gz => "gz",
+            CompressionFormat::Xz => "xz",
+        }
+    }
+
     pub(crate) fn encode(&self, path: impl AsRef<Path>) -> Result<Box<dyn Encoder>, Error> {
-        let extension = match self {
-            CompressionFormat::Gz => ".gz",
-            CompressionFormat::Xz => ".xz",
-        };
         let mut os = path.as_ref().as_os_str().to_os_string();
-        os.push(extension);
+        os.push(format!(".{}", self.extension()));
         let path = Path::new(&os);
 
         if path.exists() {
@@ -38,6 +49,14 @@ impl CompressionFormat {
             }
         })
     }
+
+    pub(crate) fn decode(&self, path: impl AsRef<Path>) -> Result<Box<dyn Read>, Error> {
+        let file = crate::util::open_file(path.as_ref())?;
+        Ok(match self {
+            CompressionFormat::Gz => Box::new(GzDecoder::new(file)),
+            CompressionFormat::Xz => Box::new(XzDecoder::new(file)),
+        })
+    }
 }
 
 pub(crate) trait Encoder: Send + Write {
diff --git a/test.sh b/test.sh
index dc7275b..3dc5e49 100755
--- a/test.sh
+++ b/test.sh
@@ -1164,6 +1164,32 @@ docdir_combined() {
 }
 runtest docdir_combined
 
+combine_installers_different_input_compression_formats() {
+    try sh "$S/gen-installer.sh" \
+        --image-dir="$TEST_DIR/image1" \
+        --work-dir="$WORK_DIR" \
+        --output-dir="$OUT_DIR" \
+        --package-name=rustc \
+        --component-name=rustc \
+        --compression-formats=xz
+    try sh "$S/gen-installer.sh" \
+        --image-dir="$TEST_DIR/image3" \
+        --work-dir="$WORK_DIR" \
+        --output-dir="$OUT_DIR" \
+        --package-name=cargo \
+        --component-name=cargo \
+        --compression-formats=gz
+    try sh "$S/combine-installers.sh" \
+        --work-dir="$WORK_DIR" \
+        --output-dir="$OUT_DIR" \
+        --package-name=rust \
+        --input-tarballs="$OUT_DIR/rustc.tar.xz,$OUT_DIR/cargo.tar.gz"
+
+    try test -e "${OUT_DIR}/rust.tar.gz"
+    try test -e "${OUT_DIR}/rust.tar.xz"
+}
+runtest combine_installers_different_input_compression_formats
+
 echo
 echo "TOTAL SUCCESS!"
 echo