fix: keep Google Drive upload tails on shared drives

docker-compose.yml

@@ -25,7 +25,7 @@ services:
       LOG: debug
       TZ: "Europe/Paris"
       #      TMPDIR: /scratch
-      EDGE_KEY: "eyJzZXJ2ZXJVcmwiOiJodHRwOi8vbG9jYWxob3N0Ojg4ODciLCJhZ2VudElkIjoiMTM1ZmFiM2QtZGQ2Yy00MjVjLWI1YTMtMzRmMTBhYTBmNDU0IiwibWFzdGVyS2V5QjY0IjoiQlhWM1hvbEM2NTZTVjdkTmdjV1BHUWxrKytycExJNmxHRGk3Q1BCNWllbz0ifQ=="
+      EDGE_KEY: "eyJzZXJ2ZXJVcmwiOiJodHRwOi8vbG9jYWxob3N0Ojg4ODciLCJhZ2VudElkIjoiOGYzOTAyYjgtZmU0Yy00NmY2LWE0ODktNzE5MTNhZTgxMjg3IiwibWFzdGVyS2V5QjY0IjoiMUh0djdtWCtYVkJxL0IzUEV2WDlZZjlQeUdVZW5oRHlXemo5THRqNW90WT0ifQ=="
       #CHUNK_SIZE_MB: "1"
       #POLLING: 1
       #RETRY_ATTEMPTS: 3

src/services/storage/providers/google_drive/helpers.rs

@@ -11,6 +11,11 @@ use reqwest::{Client as ReqwestClient, StatusCode};
 use reqwest::{Client, header};
 use serde_json::{Value, json};
 
+pub const DRIVE_FILES_URL: &str =
+    "https://www.googleapis.com/drive/v3/files?supportsAllDrives=true";
+pub const DRIVE_RESUMABLE_UPLOAD_URL: &str =
+    "https://www.googleapis.com/upload/drive/v3/files?uploadType=resumable&supportsAllDrives=true";
+
 pub async fn get_google_drive_token(config: &GoogleDriveProviderConfig) -> Result<String> {
     let http_client = OAuth2ReqwestClient::new();
 
@@ -80,11 +85,10 @@ pub async fn ensure_folder_path(
             "name": name,
             "mimeType": "application/vnd.google-apps.folder",
             "parents": [parent_id],
-            "supportsAllDrives": true,
         });
 
         let folder = client
-            .post("https://www.googleapis.com/drive/v3/files")
+            .post(DRIVE_FILES_URL)
             .bearer_auth(&token)
             .json(&create_payload)
             .send()
@@ -141,7 +145,7 @@ pub async fn find_file_by_name(
 pub async fn upload_stream_to_google_drive(
     config: &GoogleDriveProviderConfig,
     full_path: &str,
-    mut content_stream: impl Stream<Item = Result<Bytes, std::io::Error>> + Send + Unpin + 'static,
+    content_stream: impl Stream<Item = Result<Bytes, std::io::Error>> + Send + Unpin + 'static,
     total_size: u64,
     mime_type: Option<&str>,
 ) -> Result<()> {
@@ -171,14 +175,12 @@ pub async fn upload_stream_to_google_drive(
         "name": file_name,
         "parents": [folder_id],
         "mimeType": mime,
-        "supportsAllDrives": true,
     });
 
     let session_res = client
-        .post("https://www.googleapis.com/upload/drive/v3/files?uploadType=resumable")
+        .post(DRIVE_RESUMABLE_UPLOAD_URL)
         .bearer_auth(&token)
         .header("X-Upload-Content-Type", mime)
-        .header("X-Upload-Content-Length", total_size.to_string()) // Helps a lot
         .json(&metadata)
         .send()
         .await
@@ -196,54 +198,61 @@ pub async fn upload_stream_to_google_drive(
         .to_str()?
         .to_string();
 
-    const CHUNK_SIZE: u64 = 8 * 1024 * 1024;
-
-    let mut uploaded: u64 = 0;
+    upload_chunks(&client, &upload_url, content_stream, total_size).await
+}
 
-    while uploaded < total_size {
-        let chunk_size = (total_size - uploaded).min(CHUNK_SIZE);
+pub async fn upload_chunks(
+    client: &Client,
+    upload_url: &str,
+    mut content_stream: impl Stream<Item = Result<Bytes, std::io::Error>> + Unpin,
+    source_size: u64,
+) -> Result<()> {
+    const CHUNK_SIZE: usize = 8 * 1024 * 1024;
 
-        let mut chunk_bytes = Vec::with_capacity(chunk_size as usize);
-        let mut remaining = chunk_size;
+    let mut uploaded: u64 = 0;
+    let mut pending = Bytes::new();
+    let mut stream_ended = false;
 
-        while remaining > 0 {
-            match content_stream.next().await {
-                Some(Ok(bytes)) => {
-                    let to_take = remaining.min(bytes.len() as u64) as usize;
-                    chunk_bytes.extend_from_slice(&bytes[..to_take]);
-                    remaining -= to_take as u64;
+    while !stream_ended {
+        let mut chunk_bytes = Vec::with_capacity(CHUNK_SIZE);
 
-                    if to_take < bytes.len() {
-                        // TODO : Put remainder back
-                    }
-                }
-                Some(Err(e)) => return Err(e).context("Stream error during chunk"),
-                None => {
-                    if uploaded + chunk_bytes.len() as u64 != total_size {
-                        return Err(anyhow::anyhow!("Stream ended early"));
+        loop {
+            if pending.is_empty() {
+                match content_stream.next().await {
+                    Some(Ok(bytes)) => pending = bytes,
+                    Some(Err(e)) => return Err(e).context("Stream error during chunk"),
+                    None => {
+                        stream_ended = true;
+                        break;
                     }
                 }
+            } else if chunk_bytes.len() < CHUNK_SIZE {
+                let to_take = (CHUNK_SIZE - chunk_bytes.len()).min(pending.len());
+                chunk_bytes.extend_from_slice(&pending.split_to(to_take));
+            } else {
+                break;
             }
         }
 
-        if chunk_bytes.is_empty() && uploaded < total_size {
-            return Err(anyhow::anyhow!("Unexpected end of stream"));
+        let chunk_end = uploaded + chunk_bytes.len() as u64;
+        if stream_ended && chunk_end < source_size {
+            return Err(anyhow::anyhow!("Stream ended early"));
         }
 
-        let range_end = uploaded + chunk_bytes.len() as u64 - 1;
-        let content_range = if uploaded + chunk_bytes.len() as u64 == total_size {
-            format!("bytes {}-{}/{}", uploaded, range_end, total_size)
-        } else {
-            format!("bytes {}-{}/*", uploaded, range_end)
+        let content_range = match (chunk_bytes.is_empty(), stream_ended) {
+            (true, _) => format!("bytes */{chunk_end}"),
+            (false, true) => format!("bytes {uploaded}-{}/{chunk_end}", chunk_end - 1),
+            (false, false) => format!("bytes {uploaded}-{}/*", chunk_end - 1),
         };
+        let chunk_bytes = Bytes::from(chunk_bytes);
 
         let mut retries = 0;
         loop {
             let res = client
-                .put(&upload_url)
+                .put(upload_url)
                 .header("Content-Range", &content_range)
                 .header("Content-Length", chunk_bytes.len().to_string())
-                .body(chunk_bytes.clone()) // clone is cheap if small; optimize later if needed
+                .body(chunk_bytes.clone())
                 .send()
                 .await;
 
@@ -253,8 +262,8 @@ pub async fn upload_stream_to_google_drive(
                         || resp.status() == StatusCode::PERMANENT_REDIRECT =>
                 {
                     // 200 or 308 = good
-                    uploaded += chunk_bytes.len() as u64;
-                    tracing::info!("Uploaded {}/{} bytes", uploaded, total_size);
+                    uploaded = chunk_end;
+                    tracing::info!("Uploaded {} bytes", uploaded);
                     break;
                 }
                 Ok(resp) if resp.status() == StatusCode::TOO_MANY_REQUESTS => {

src/services/storage/providers/google_drive/mod.rs

@@ -1,4 +1,4 @@
-mod helpers;
+pub mod helpers;
 pub mod models;
 
 use crate::core::context::Context;

src/tests/storage/google_drive.rs

@@ -0,0 +1,151 @@
+use crate::services::storage::providers::google_drive::helpers::{
+    DRIVE_FILES_URL, DRIVE_RESUMABLE_UPLOAD_URL, upload_chunks,
+};
+use crate::tests::services::local_storage_tests::MASTER_KEY_B64;
+use crate::utils::file::decrypt_file_stream_gcm;
+use crate::utils::stream::build_stream;
+
+use bytes::Bytes;
+use futures::stream;
+use reqwest::Client;
+use tempfile::TempDir;
+use wiremock::matchers::{method, path};
+use wiremock::{Mock, MockServer, ResponseTemplate};
+
+const WINDOW: usize = 8 * 1024 * 1024;
+
+async fn drive_session() -> MockServer {
+    let server = MockServer::start().await;
+    Mock::given(method("PUT"))
+        .and(path("/upload"))
+        .respond_with(ResponseTemplate::new(200))
+        .mount(&server)
+        .await;
+    server
+}
+
+async fn received_puts(server: &MockServer) -> Vec<(String, Vec<u8>)> {
+    server
+        .received_requests()
+        .await
+        .unwrap()
+        .into_iter()
+        .map(|request| {
+            let content_range = request.headers["content-range"]
+                .to_str()
+                .unwrap()
+                .to_string();
+            (content_range, request.body)
+        })
+        .collect()
+}
+
+#[tokio::test]
+async fn sixteen_mib_chunk_is_sent_as_two_eight_mib_windows() {
+    let server = drive_session().await;
+    let payload: Vec<u8> = (0..2 * WINDOW).map(|index| (index % 251) as u8).collect();
+    let content_stream = stream::iter(vec![Ok(Bytes::from(payload.clone()))]);
+
+    upload_chunks(
+        &Client::new(),
+        &format!("{}/upload", server.uri()),
+        content_stream,
+        payload.len() as u64,
+    )
+    .await
+    .unwrap();
+
+    let puts = received_puts(&server).await;
+    assert_eq!(puts.len(), 2);
+    assert_eq!(puts[0].0, "bytes 0-8388607/*");
+    assert_eq!(puts[1].0, "bytes 8388608-16777215/16777216");
+    assert_eq!(puts[0].1.len(), WINDOW);
+    assert_eq!(puts[1].1.len(), WINDOW);
+    assert_eq!(
+        [puts[0].1.as_slice(), puts[1].1.as_slice()].concat(),
+        payload
+    );
+}
+
+#[tokio::test]
+async fn encrypted_dump_uploaded_to_drive_decrypts_back() {
+    let server = drive_session().await;
+    let work = TempDir::new().unwrap();
+    let dump_path = work.path().join("dump.sql");
+    let plaintext: Vec<u8> = (0..WINDOW + 1024 * 1024)
+        .map(|index| (index % 253) as u8)
+        .collect();
+    std::fs::write(&dump_path, &plaintext).unwrap();
+    let upload = build_stream(&dump_path, true, &MASTER_KEY_B64.to_string())
+        .await
+        .unwrap();
+
+    upload_chunks(
+        &Client::new(),
+        &format!("{}/upload", server.uri()),
+        upload.stream,
+        plaintext.len() as u64,
+    )
+    .await
+    .unwrap();
+
+    let puts = received_puts(&server).await;
+    let uploaded: Vec<u8> = puts.iter().flat_map(|(_, body)| body.clone()).collect();
+    let (last_range, _) = puts.last().unwrap();
+    assert!(last_range.ends_with(&format!("/{}", uploaded.len())));
+    let encrypted_path = work.path().join("dump.sql.enc");
+    let decrypted_path = work.path().join("dump.restored.sql");
+    std::fs::write(&encrypted_path, &uploaded).unwrap();
+    decrypt_file_stream_gcm(
+        encrypted_path,
+        decrypted_path.clone(),
+        MASTER_KEY_B64.to_string(),
+    )
+    .await
+    .unwrap();
+    assert_eq!(std::fs::read(decrypted_path).unwrap(), plaintext);
+}
+
+#[tokio::test]
+async fn stream_shorter_than_the_dump_never_finalizes_the_upload() {
+    let server = drive_session().await;
+    let content_stream = stream::iter(vec![Ok(Bytes::from(vec![1u8; 1024]))]);
+
+    let error = upload_chunks(
+        &Client::new(),
+        &format!("{}/upload", server.uri()),
+        content_stream,
+        2048,
+    )
+    .await
+    .unwrap_err();
+
+    assert_eq!(error.to_string(), "Stream ended early");
+    assert!(received_puts(&server).await.is_empty());
+}
+
+#[tokio::test]
+async fn empty_dump_finalizes_a_zero_byte_file() {
+    let server = drive_session().await;
+
+    upload_chunks(
+        &Client::new(),
+        &format!("{}/upload", server.uri()),
+        stream::iter(Vec::<Result<Bytes, std::io::Error>>::new()),
+        0,
+    )
+    .await
+    .unwrap();
+
+    let puts = received_puts(&server).await;
+    assert_eq!(puts.len(), 1);
+    assert_eq!(puts[0].0, "bytes */0");
+    assert!(puts[0].1.is_empty());
+}
+
+#[test]
+fn drive_write_urls_carry_supports_all_drives() {
+    assert!(DRIVE_RESUMABLE_UPLOAD_URL.contains("supportsAllDrives=true"));
+    assert!(DRIVE_RESUMABLE_UPLOAD_URL.contains("uploadType=resumable"));
+    assert!(DRIVE_FILES_URL.contains("supportsAllDrives=true"));
+}

src/tests/storage/mod.rs

@@ -1,5 +1,6 @@
 mod azure_blob;
 mod google_cloud_storage;
+mod google_drive;
 mod rclone;
 mod rclone_target;
 mod sftp;