From e7e93210809c7ebf99ae4981d992a939a5ff1131 Mon Sep 17 00:00:00 2001
From: Quaternions <krakow20@gmail.com>
Date: Thu, 10 Apr 2025 15:55:09 -0700
Subject: [PATCH] rbx_asset: change api to save intermediate allocation

---
 rbx_asset/Cargo.toml    |  6 ++++-
 rbx_asset/src/cloud.rs  |  9 +++----
 rbx_asset/src/cookie.rs | 17 ++++++-------
 rbx_asset/src/util.rs   | 56 ++++++++++++++++++++++++++++++++++-------
 src/main.rs             | 26 +++++++++----------
 5 files changed, 77 insertions(+), 37 deletions(-)

diff --git a/rbx_asset/Cargo.toml b/rbx_asset/Cargo.toml
index 2b0ad78..0120817 100644
--- a/rbx_asset/Cargo.toml
+++ b/rbx_asset/Cargo.toml
@@ -10,10 +10,14 @@ authors = ["Rhys Lloyd <krakow20@gmail.com>"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
+[features]
+default = ["gzip"]
+gzip = ["dep:flate2"]
+
 [dependencies]
 bytes = "1.10.1"
 chrono = { version = "0.4.38", features = ["serde"] }
-flate2 = "1.0.29"
+flate2 = { version = "1.0.29", optional = true }
 reqwest = { version = "0.12.4", features = ["json","multipart"] }
 serde = { version = "1.0.199", features = ["derive"] }
 serde_json = "1.0.111"
diff --git a/rbx_asset/src/cloud.rs b/rbx_asset/src/cloud.rs
index 9965e13..18e985b 100644
--- a/rbx_asset/src/cloud.rs
+++ b/rbx_asset/src/cloud.rs
@@ -1,4 +1,4 @@
-use crate::util::{serialize_u64,deserialize_u64,response_ok,ResponseError,maybe_gzip_decode};
+use crate::util::{serialize_u64,deserialize_u64,response_ok,ResponseError,MaybeGzippedBytes};
 
 #[derive(Debug,serde::Deserialize,serde::Serialize)]
 #[allow(nonstandard_style,dead_code)]
@@ -174,7 +174,6 @@ pub enum GetError{
 	ParseError(url::ParseError),
 	Response(ResponseError),
 	Reqwest(reqwest::Error),
-	IO(std::io::Error)
 }
 impl std::fmt::Display for GetError{
 	fn fmt(&self,f:&mut std::fmt::Formatter<'_>)->std::fmt::Result{
@@ -453,15 +452,15 @@ impl Context{
 		).await.map_err(GetError::Response)?
 		.json().await.map_err(GetError::Reqwest)
 	}
-	pub async fn get_asset(&self,config:&AssetLocation)->Result<Vec<u8>,GetError>{
+	pub async fn get_asset(&self,config:&AssetLocation)->Result<MaybeGzippedBytes,GetError>{
 		let url=reqwest::Url::parse(config.location()).map_err(GetError::ParseError)?;
 
-		let body=response_ok(
+		let bytes=response_ok(
 			self.get(url).await.map_err(GetError::Reqwest)?
 		).await.map_err(GetError::Response)?
 		.bytes().await.map_err(GetError::Reqwest)?;
 
-		maybe_gzip_decode(body).map_err(GetError::IO)
+		Ok(MaybeGzippedBytes::new(bytes))
 	}
 	pub async fn get_asset_versions(&self,config:AssetVersionsRequest)->Result<AssetVersionsResponse,AssetVersionsError>{
 		let raw_url=format!("https://apis.roblox.com/assets/v1/assets/{}/versions",config.asset_id);
diff --git a/rbx_asset/src/cookie.rs b/rbx_asset/src/cookie.rs
index 34e83e7..cbc56d3 100644
--- a/rbx_asset/src/cookie.rs
+++ b/rbx_asset/src/cookie.rs
@@ -1,4 +1,4 @@
-use crate::util::{response_ok,ResponseError,maybe_gzip_decode};
+use crate::util::{response_ok,ResponseError,MaybeGzippedBytes};
 
 #[derive(Debug)]
 pub enum PostError{
@@ -91,7 +91,6 @@ pub enum GetError{
 	ParseError(url::ParseError),
 	Response(ResponseError),
 	Reqwest(reqwest::Error),
-	IO(std::io::Error)
 }
 impl std::fmt::Display for GetError{
 	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
@@ -449,7 +448,7 @@ impl Context{
 			})
 		}
 	}
-	pub async fn get_asset(&self,config:GetAssetRequest)->Result<Vec<u8>,GetError>{
+	pub async fn get_asset(&self,config:GetAssetRequest)->Result<MaybeGzippedBytes,GetError>{
 		let mut url=reqwest::Url::parse("https://assetdelivery.roblox.com/v1/asset/").map_err(GetError::ParseError)?;
 		//url borrow scope
 		{
@@ -459,13 +458,13 @@ impl Context{
 				query.append_pair("version",version.to_string().as_str());
 			}
 		}
-		let body=response_ok(
+
+		let bytes=response_ok(
 			self.get(url).await.map_err(GetError::Reqwest)?
 		).await.map_err(GetError::Response)?
 		.bytes().await.map_err(GetError::Reqwest)?;
 
-
-		maybe_gzip_decode(body).map_err(GetError::IO)
+		Ok(MaybeGzippedBytes::new(bytes))
 	}
 	pub async fn get_asset_v2(&self,config:GetAssetRequest)->Result<GetAssetV2,GetAssetV2Error>{
 		let mut url=reqwest::Url::parse("https://assetdelivery.roblox.com/v2/asset").map_err(GetAssetV2Error::ParseError)?;
@@ -497,15 +496,15 @@ impl Context{
 			info,
 		})
 	}
-	pub async fn get_asset_v2_download(&self,config:&GetAssetV2Location)->Result<Vec<u8>,GetError>{
+	pub async fn get_asset_v2_download(&self,config:&GetAssetV2Location)->Result<MaybeGzippedBytes,GetError>{
 		let url=reqwest::Url::parse(config.location.as_str()).map_err(GetError::ParseError)?;
 
-		let body=response_ok(
+		let bytes=response_ok(
 			self.get(url).await.map_err(GetError::Reqwest)?
 		).await.map_err(GetError::Response)?
 		.bytes().await.map_err(GetError::Reqwest)?;
 
-		maybe_gzip_decode(body).map_err(GetError::IO)
+		Ok(MaybeGzippedBytes::new(bytes))
 	}
 	pub async fn get_asset_details(&self,config:GetAssetDetailsRequest)->Result<AssetDetails,GetError>{
 		let url=reqwest::Url::parse(format!("https://economy.roblox.com/v2/assets/{}/details",config.asset_id).as_str()).map_err(GetError::ParseError)?;
diff --git a/rbx_asset/src/util.rs b/rbx_asset/src/util.rs
index 523743e..3c37784 100644
--- a/rbx_asset/src/util.rs
+++ b/rbx_asset/src/util.rs
@@ -33,15 +33,53 @@ pub(crate) async fn response_ok(response:reqwest::Response)->Result<reqwest::Res
 	}
 }
 
-pub(crate) fn maybe_gzip_decode(data:bytes::Bytes)->std::io::Result<Vec<u8>>{
-	match data.get(0..2){
-		Some(b"\x1f\x8b")=>{
-			use std::io::Read;
-			let mut buf=Vec::new();
-			flate2::read::GzDecoder::new(std::io::Cursor::new(data)).read_to_end(&mut buf)?;
-			Ok(buf)
-		},
-		_=>Ok(data.to_vec()),
+#[cfg(feature="gzip")]
+use std::io::Cursor;
+#[cfg(feature="gzip")]
+use flate2::read::GzDecoder;
+
+/// Some bytes that might be gzipped.  Use the read_with or to_vec methods to transparently decode gzip.
+pub struct MaybeGzippedBytes{
+	bytes:bytes::Bytes,
+}
+impl MaybeGzippedBytes{
+	pub(crate) fn new(bytes:bytes::Bytes)->Self{
+		Self{bytes}
+	}
+	pub fn into_inner(self)->bytes::Bytes{
+		self.bytes
+	}
+	/// get a reference to the bytes, ignoring gzip decoding
+	pub fn as_raw_ref(&self)->&[u8]{
+		self.bytes.as_ref()
+	}
+	/// Transparently decode gzip data, if present (intermediate allocation)
+	#[cfg(feature="gzip")]
+	pub fn to_vec(&self)->std::io::Result<Vec<u8>>{
+		use std::io::Read;
+		match self.bytes.get(0..2){
+			Some(b"\x1f\x8b")=>{
+				let mut buf=Vec::new();
+				GzDecoder::new(Cursor::new(self.bytes.as_ref())).read_to_end(&mut buf)?;
+				Ok(buf)
+			},
+			_=>Ok(self.bytes.to_vec())
+		}
+	}
+	/// Read the bytes with the provided decoders.
+	/// The idea is to make a function that is generic over std::io::Read
+	/// and pass the same function to both closures.
+	/// This two closure hack must be done because of the different concrete types.
+	#[cfg(feature="gzip")]
+	pub fn read_with<'a,ReadGzip,ReadRaw,T>(&'a self,read_gzip:ReadGzip,read_raw:ReadRaw)->T
+		where
+			ReadGzip:Fn(GzDecoder<Cursor<&'a [u8]>>)->T,
+			ReadRaw:Fn(Cursor<&'a [u8]>)->T,
+	{
+		match self.bytes.get(0..2){
+			Some(b"\x1f\x8b")=>read_gzip(GzDecoder::new(Cursor::new(self.bytes.as_ref()))),
+			_=>read_raw(Cursor::new(self.bytes.as_ref()))
+		}
 	}
 }
 
diff --git a/src/main.rs b/src/main.rs
index b30aa28..e8703bf 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -903,11 +903,11 @@ async fn create_asset_medias(config:CreateAssetMediasConfig)->AResult<()>{
 		async move{(path,
 		async move{
 			let asset_response=asset_response_result.map_err(DownloadDecalError::PollOperation)?;
-			let file=cookie_context.get_asset(rbx_asset::cookie::GetAssetRequest{
+			let maybe_gzip=cookie_context.get_asset(rbx_asset::cookie::GetAssetRequest{
 				asset_id:asset_response.assetId,
 				version:None,
 			}).await.map_err(DownloadDecalError::Get)?;
-			let dom=load_dom(std::io::Cursor::new(file)).map_err(DownloadDecalError::LoadDom)?;
+			let dom=maybe_gzip.read_with(load_dom,load_dom).map_err(DownloadDecalError::LoadDom)?;
 			let instance=dom.get_by_ref(
 				*dom.root().children().first().ok_or(DownloadDecalError::NoFirstInstance)?
 			).ok_or(DownloadDecalError::NoFirstInstance)?;
@@ -993,8 +993,8 @@ async fn asset_details(cookie:Cookie,asset_id:AssetID)->AResult<()>{
 
 async fn download_version(cookie:Cookie,asset_id:AssetID,version:Option<u64>,dest:PathBuf)->AResult<()>{
 	let context=CookieContext::new(cookie);
-	let data=context.get_asset(rbx_asset::cookie::GetAssetRequest{asset_id,version}).await?;
-	tokio::fs::write(dest,data).await?;
+	let maybe_gzip=context.get_asset(rbx_asset::cookie::GetAssetRequest{asset_id,version}).await?;
+	tokio::fs::write(dest,maybe_gzip.to_vec()?).await?;
 	Ok(())
 }
 
@@ -1006,9 +1006,9 @@ async fn download_version_v2(cookie:Cookie,asset_id:AssetID,version:Option<u64>,
 	println!("version:{}",info.version);
 
 	let location=info.info.locations.first().ok_or(anyhow::Error::msg("No locations"))?;
-	let data=context.get_asset_v2_download(location).await?;
+	let maybe_gzip=context.get_asset_v2_download(location).await?;
 
-	tokio::fs::write(dest,data).await?;
+	tokio::fs::write(dest,maybe_gzip.to_vec()?).await?;
 	Ok(())
 }
 
@@ -1024,7 +1024,7 @@ async fn download_list(cookie:Cookie,asset_id_file_map:AssetIDFileMap)->AResult<
 	.buffer_unordered(CONCURRENT_REQUESTS)
 	.for_each(|b:AResult<_>|async{
 			match b{
-				Ok((dest,data))=>if let Err(e)=tokio::fs::write(dest,data).await{
+				Ok((dest,maybe_gzip))=>if let Err(e)=(async||{tokio::fs::write(dest,maybe_gzip.to_vec()?).await})().await{
 					eprintln!("fs error: {}",e);
 				},
 				Err(e)=>eprintln!("dl error: {}",e),
@@ -1228,9 +1228,9 @@ async fn download_history(mut config:DownloadHistoryConfig)->AResult<()>{
 			let mut path=output_folder.clone();
 			path.push(format!("{}_v{}.rbxl",config.asset_id,version_number));
 			join_set.spawn(async move{
-				let file=context.get_asset(rbx_asset::cookie::GetAssetRequest{asset_id:config.asset_id,version:Some(version_number)}).await?;
+				let maybe_gzip=context.get_asset(rbx_asset::cookie::GetAssetRequest{asset_id:config.asset_id,version:Some(version_number)}).await?;
 
-				tokio::fs::write(path,file).await?;
+				tokio::fs::write(path,maybe_gzip.to_vec()?).await?;
 
 				Ok::<_,anyhow::Error>(())
 			});
@@ -1350,9 +1350,9 @@ struct DownloadDecompileConfig{
 
 async fn download_decompile(config:DownloadDecompileConfig)->AResult<()>{
 	let context=CookieContext::new(config.cookie);
-	let file=context.get_asset(rbx_asset::cookie::GetAssetRequest{asset_id:config.asset_id,version:None}).await?;
+	let maybe_gzip=context.get_asset(rbx_asset::cookie::GetAssetRequest{asset_id:config.asset_id,version:None}).await?;
 
-	let dom=load_dom(std::io::Cursor::new(file))?;
+	let dom=maybe_gzip.read_with(load_dom,load_dom)?;
 	let context=rox_compiler::DecompiledContext::from_dom(dom);
 
 	context.write_files(rox_compiler::WriteConfig{
@@ -1532,8 +1532,8 @@ async fn download_and_decompile_history_into_git(config:DownloadAndDecompileHist
 	.map(|asset_version|{
 		let context=context.clone();
 		tokio::task::spawn(async move{
-			let file=context.get_asset(rbx_asset::cookie::GetAssetRequest{asset_id,version:Some(asset_version.assetVersionNumber)}).await?;
-			let dom=load_dom(std::io::Cursor::new(file))?;
+			let maybe_gzip=context.get_asset(rbx_asset::cookie::GetAssetRequest{asset_id,version:Some(asset_version.assetVersionNumber)}).await?;
+			let dom=maybe_gzip.read_with(load_dom,load_dom)?;
 			Ok::<_,anyhow::Error>((asset_version,rox_compiler::DecompiledContext::from_dom(dom)))
 		})
 	}))