From 43f17fe44103628ff943e5988f7dfb9a7f14478c Mon Sep 17 00:00:00 2001
From: Connor Peet <connor@peet.io>
Date: Mon, 5 Aug 2024 14:19:36 -0700
Subject: [PATCH] cli: determine archive type based on filename instead of path
 (#224605)

Refs #219632

Seems like PRSS sometimes(?) doesn't return the full archive name in
the response. I don't reproduce this, but others consistently do. This
PR removes the dependency on the URL path and instead checks for the
gzip magic number in the first two bytes of the archive to figure out
what to do.
---
 cli/src/update_service.rs | 15 ++++++++-------
 cli/src/util/tar.rs       | 21 ++++++++++++++-------
 cli/src/util/zipper.rs    |  9 +++------
 3 files changed, 25 insertions(+), 20 deletions(-)

diff --git a/cli/src/update_service.rs b/cli/src/update_service.rs
index 4bec13d6e86a3..cc3caeb446924 100644
--- a/cli/src/update_service.rs
+++ b/cli/src/update_service.rs
@@ -3,7 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-use std::{ffi::OsStr, fmt, path::Path};
+use std::{fmt, path::Path};
 
 use serde::{Deserialize, Serialize};
 
@@ -11,10 +11,11 @@ use crate::{
 	constants::VSCODE_CLI_UPDATE_ENDPOINT,
 	debug, log, options, spanf,
 	util::{
-		errors::{AnyError, CodeError, WrappedError},
+		errors::{wrap, AnyError, CodeError, WrappedError},
 		http::{BoxedHttp, SimpleResponse},
 		io::ReportCopyProgress,
-		tar, zipper,
+		tar::{self, has_gzip_header},
+		zipper,
 	},
 };
 
@@ -178,10 +179,10 @@ pub fn unzip_downloaded_release<T>(
 where
 	T: ReportCopyProgress,
 {
-	if compressed_file.extension() == Some(OsStr::new("zip")) {
-		zipper::unzip_file(compressed_file, target_dir, reporter)
-	} else {
-		tar::decompress_tarball(compressed_file, target_dir, reporter)
+	match has_gzip_header(compressed_file) {
+		Ok((f, true)) => tar::decompress_tarball(f, target_dir, reporter),
+		Ok((f, false)) => zipper::unzip_file(f, target_dir, reporter),
+		Err(e) => Err(wrap(e, "error checking for gzip header")),
 	}
 }
 
diff --git a/cli/src/util/tar.rs b/cli/src/util/tar.rs
index fe4d426970021..10577140325eb 100644
--- a/cli/src/util/tar.rs
+++ b/cli/src/util/tar.rs
@@ -5,8 +5,8 @@
 use crate::util::errors::{wrap, WrappedError};
 
 use flate2::read::GzDecoder;
-use std::fs;
-use std::io::Seek;
+use std::fs::{self, File};
+use std::io::{Read, Seek};
 use std::path::{Path, PathBuf};
 use tar::Archive;
 
@@ -57,16 +57,13 @@ fn should_skip_first_segment(file: &fs::File) -> Result<(bool, u64), WrappedErro
 }
 
 pub fn decompress_tarball<T>(
-	path: &Path,
+	mut tar_gz: File,
 	parent_path: &Path,
 	mut reporter: T,
 ) -> Result<(), WrappedError>
 where
 	T: ReportCopyProgress,
 {
-	let mut tar_gz = fs::File::open(path)
-		.map_err(|e| wrap(e, format!("error opening file {}", path.display())))?;
-
 	let (skip_first, num_entries) = should_skip_first_segment(&tar_gz)?;
 	let report_progress_every = num_entries / 20;
 	let mut entries_so_far = 0;
@@ -81,7 +78,7 @@ where
 	let mut archive = Archive::new(tar);
 	archive
 		.entries()
-		.map_err(|e| wrap(e, format!("error opening archive {}", path.display())))?
+		.map_err(|e| wrap(e, "error opening archive"))?
 		.filter_map(|e| e.ok())
 		.try_for_each::<_, Result<_, WrappedError>>(|mut entry| {
 			// approximate progress based on where we are in the archive:
@@ -118,3 +115,13 @@ where
 
 	Ok(())
 }
+
+pub fn has_gzip_header(path: &Path) -> std::io::Result<(File, bool)> {
+	let mut file = fs::File::open(path)?;
+	let mut header = [0; 2];
+	let _ = file.read_exact(&mut header);
+
+	file.rewind()?;
+
+	Ok((file, header[0] == 0x1f && header[1] == 0x8b))
+}
diff --git a/cli/src/util/zipper.rs b/cli/src/util/zipper.rs
index 69bcf2d23f6d3..5521fa42c5424 100644
--- a/cli/src/util/zipper.rs
+++ b/cli/src/util/zipper.rs
@@ -44,15 +44,12 @@ fn should_skip_first_segment(archive: &mut ZipArchive<File>) -> bool {
 	archive.len() > 1 // prefix removal is invalid if there's only a single file
 }
 
-pub fn unzip_file<T>(path: &Path, parent_path: &Path, mut reporter: T) -> Result<(), WrappedError>
+pub fn unzip_file<T>(file: File, parent_path: &Path, mut reporter: T) -> Result<(), WrappedError>
 where
 	T: ReportCopyProgress,
 {
-	let file = fs::File::open(path)
-		.map_err(|e| wrap(e, format!("unable to open file {}", path.display())))?;
-
-	let mut archive = zip::ZipArchive::new(file)
-		.map_err(|e| wrap(e, format!("failed to open zip archive {}", path.display())))?;
+	let mut archive =
+		zip::ZipArchive::new(file).map_err(|e| wrap(e, "failed to open zip archive"))?;
 
 	let skip_segments_no = usize::from(should_skip_first_segment(&mut archive));
 	let report_progress_every = archive.len() / 20;