Skip to content

Commit 315ec42

Browse files
Varun NuthalapatiVarun Nuthalapati
authored andcommitted
fixup: sanitize API errors, add --dry-run, stream response to file
- Sanitize API error body strings with sanitize_for_terminal() before embedding them in GwsError::Api to prevent terminal escape injection - Add --dry-run support: after metadata fetch and path resolution, print what would be downloaded and return without network/disk I/O - Stream response bytes to file via tokio::fs::File + bytes_stream() instead of resp.bytes().await to avoid OOM on large Drive files
1 parent 51f65d2 commit 315ec42

1 file changed

Lines changed: 61 additions & 25 deletions

File tree

  • crates/google-workspace-cli/src/helpers

crates/google-workspace-cli/src/helpers/drive.rs

Lines changed: 61 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -178,10 +178,14 @@ TIPS:
178178
}
179179

180180
async fn handle_download(matches: &ArgMatches) -> Result<(), GwsError> {
181+
use futures_util::StreamExt;
182+
use tokio::io::AsyncWriteExt;
183+
181184
let file_id =
182185
crate::validate::validate_resource_name(matches.get_one::<String>("file").unwrap())?;
183186
let output_arg = matches.get_one::<String>("output");
184187
let export_mime = matches.get_one::<String>("mime-type");
188+
let dry_run = matches.get_flag("dry-run");
185189

186190
// Validate export mime-type for dangerous characters if provided
187191
if let Some(mime) = export_mime {
@@ -214,7 +218,7 @@ async fn handle_download(matches: &ArgMatches) -> Result<(), GwsError> {
214218
let body = meta_resp.text().await.unwrap_or_default();
215219
return Err(GwsError::Api {
216220
code: status.into(),
217-
message: body,
221+
message: crate::output::sanitize_for_terminal(&body),
218222
reason: "files_get_metadata_failed".to_string(),
219223
enable_url: None,
220224
});
@@ -238,10 +242,27 @@ async fn handle_download(matches: &ArgMatches) -> Result<(), GwsError> {
238242
let out_str = output_arg.map(|s| s.as_str()).unwrap_or(drive_name);
239243
let out_path = crate::validate::validate_safe_file_path(out_str, "--output")?;
240244

241-
// 3. Fetch file content — native Google Workspace files require export,
242-
// everything else uses alt=media.
245+
// 3. Dry-run: print what would be done and exit without network or disk I/O
246+
if dry_run {
247+
println!(
248+
"{}",
249+
serde_json::to_string_pretty(&json!({
250+
"dryRun": true,
251+
"fileId": file_id,
252+
"driveName": drive_name,
253+
"mimeType": mime_type,
254+
"output": out_path.display().to_string(),
255+
"exportMimeType": export_mime,
256+
}))
257+
.unwrap_or_default()
258+
);
259+
return Ok(());
260+
}
261+
262+
// 4. Fetch file content and stream it to disk — native Google Workspace
263+
// files require export; everything else uses alt=media.
243264
let is_google_native = mime_type.starts_with("application/vnd.google-apps.");
244-
let bytes = if is_google_native {
265+
let resp = if is_google_native {
245266
let mime = export_mime.ok_or_else(|| {
246267
GwsError::Validation(format!(
247268
"The file is a Google Workspace native file ({mime_type}). \
@@ -253,7 +274,7 @@ async fn handle_download(matches: &ArgMatches) -> Result<(), GwsError> {
253274
"https://www.googleapis.com/drive/v3/files/{}/export",
254275
crate::validate::encode_path_segment(file_id),
255276
);
256-
let resp = crate::client::send_with_retry(|| {
277+
let r = crate::client::send_with_retry(|| {
257278
client
258279
.get(&export_url)
259280
.query(&[("mimeType", mime.as_str())])
@@ -262,21 +283,19 @@ async fn handle_download(matches: &ArgMatches) -> Result<(), GwsError> {
262283
.await
263284
.map_err(|e| GwsError::Other(anyhow::anyhow!("Drive export request failed: {e}")))?;
264285

265-
if !resp.status().is_success() {
266-
let status = resp.status().as_u16();
267-
let body = resp.text().await.unwrap_or_default();
286+
if !r.status().is_success() {
287+
let status = r.status().as_u16();
288+
let body = r.text().await.unwrap_or_default();
268289
return Err(GwsError::Api {
269290
code: status.into(),
270-
message: body,
291+
message: crate::output::sanitize_for_terminal(&body),
271292
reason: "files_export_failed".to_string(),
272293
enable_url: None,
273294
});
274295
}
275-
resp.bytes()
276-
.await
277-
.map_err(|e| GwsError::Other(anyhow::anyhow!("Failed to read export response: {e}")))?
296+
r
278297
} else {
279-
let resp = crate::client::send_with_retry(|| {
298+
let r = crate::client::send_with_retry(|| {
280299
client
281300
.get(&metadata_url)
282301
.query(&[("alt", "media")])
@@ -285,35 +304,52 @@ async fn handle_download(matches: &ArgMatches) -> Result<(), GwsError> {
285304
.await
286305
.map_err(|e| GwsError::Other(anyhow::anyhow!("Drive download request failed: {e}")))?;
287306

288-
if !resp.status().is_success() {
289-
let status = resp.status().as_u16();
290-
let body = resp.text().await.unwrap_or_default();
307+
if !r.status().is_success() {
308+
let status = r.status().as_u16();
309+
let body = r.text().await.unwrap_or_default();
291310
return Err(GwsError::Api {
292311
code: status.into(),
293-
message: body,
312+
message: crate::output::sanitize_for_terminal(&body),
294313
reason: "files_get_media_failed".to_string(),
295314
enable_url: None,
296315
});
297316
}
298-
resp.bytes()
299-
.await
300-
.map_err(|e| GwsError::Other(anyhow::anyhow!("Failed to read download response: {e}")))?
317+
r
301318
};
302319

303-
// 4. Write to the output file
304-
std::fs::write(&out_path, &bytes).map_err(|e| {
320+
// 5. Stream response body to file to avoid OOM on large downloads
321+
let mut file = tokio::fs::File::create(&out_path).await.map_err(|e| {
322+
GwsError::Other(anyhow::anyhow!(
323+
"Failed to create '{}': {e}",
324+
out_path.display()
325+
))
326+
})?;
327+
let mut byte_count = 0usize;
328+
let mut stream = resp.bytes_stream();
329+
while let Some(chunk) = stream.next().await {
330+
let chunk =
331+
chunk.map_err(|e| GwsError::Other(anyhow::anyhow!("Download stream error: {e}")))?;
332+
byte_count += chunk.len();
333+
file.write_all(&chunk).await.map_err(|e| {
334+
GwsError::Other(anyhow::anyhow!(
335+
"Failed to write to '{}': {e}",
336+
out_path.display()
337+
))
338+
})?;
339+
}
340+
file.flush().await.map_err(|e| {
305341
GwsError::Other(anyhow::anyhow!(
306-
"Failed to write '{}': {e}",
342+
"Failed to flush '{}': {e}",
307343
out_path.display()
308344
))
309345
})?;
310346

311-
// 5. Print result as JSON (consistent with other helper output)
347+
// 6. Print result as JSON (consistent with other helper output)
312348
println!(
313349
"{}",
314350
serde_json::to_string_pretty(&json!({
315351
"file": out_path.display().to_string(),
316-
"bytes": bytes.len(),
352+
"bytes": byte_count,
317353
"mimeType": mime_type,
318354
}))
319355
.unwrap_or_default()

0 commit comments

Comments
 (0)