fix(download): verify etags for resumed requests

Signed-off-by: Marcel Guzik <[email protected]>

Author: Bravo555

Cargo.lock

@@ -13,6 +13,7 @@ repository = { workspace = true }
 anyhow = { workspace = true, features = ["backtrace"] }
 backoff = { workspace = true }
 certificate = { workspace = true, features = ["reqwest"] }
+http.workspace = true
 hyper = { workspace = true }
 log = { workspace = true }
 nix = { workspace = true }

crates/common/download/Cargo.toml

@@ -5,6 +5,7 @@ use anyhow::anyhow;
 use backoff::future::retry_notify;
 use backoff::ExponentialBackoff;
 use certificate::CloudHttpConfig;
+use hyper::header;
 use log::debug;
 use log::info;
 use log::warn;
@@ -199,6 +200,11 @@ impl Downloader {
         loop {
             let offset = next_request_offset(&prev_response, file);
             let mut response = self.request_range_from(url, offset).await?;
+            if was_resource_modified(&response, &prev_response) {
+                file.seek(SeekFrom::Start(0)).unwrap();
+                prev_response = response;
+                continue;
+            }
             let offset = partial_response::response_range_start(&response)?;
 
             if offset != 0 {
@@ -211,9 +217,7 @@ impl Downloader {
                 Ok(()) => break,
 
                 Err(SaveChunksError::Network(err)) => {
-                    prev_response = response;
                     warn!("Error while downloading response: {err}.\nRetrying...");
-                    continue;
                 }
 
                 Err(SaveChunksError::Io(err)) => {
@@ -223,6 +227,7 @@ impl Downloader {
                     })
                 }
             }
+            prev_response = response;
         }
 
         Ok(())
@@ -352,6 +357,46 @@ impl Downloader {
     }
 }
 
+/// Checks if the resource was modified between the current and previous response.
+///
+/// If the resource was updated, we should restart download and request full range of the new
+/// resource. Otherwise, a partial request can be used to resume the download.
+fn was_resource_modified(response: &Response, prev_response: &Response) -> bool {
+    // etags in current and previous request must match
+    let etag = response
+        .headers()
+        .get(header::ETAG)
+        .and_then(|h| h.to_str().ok());
+    let prev_etag = prev_response
+        .headers()
+        .get(header::ETAG)
+        .and_then(|h| h.to_str().ok());
+
+    match (etag, prev_etag) {
+        (None, None) => {
+            // no etags in either request, assume resource is unchanged
+            false
+        }
+        (None, Some(_)) | (Some(_), None) => {
+            // previous request didn't have etag and this does or vice versa, abort
+            true
+        }
+        (Some(etag), Some(prev_etag)) => {
+            // Examples:
+            // ETag: "xyzzy"
+            // ETag: W/"xyzzy"
+            // ETag: ""
+            if etag.starts_with("W/") {
+                // validator is weak, but in range requests tags must match using strong comparison
+                // https://www.rfc-editor.org/rfc/rfc9110#entity.tag.comparison
+                return true;
+            }
+
+            etag != prev_etag
+        }
+    }
+}
+
 /// Decides whether HTTP request error is retryable.
 fn reqwest_err_to_backoff(err: reqwest::Error) -> backoff::Error<reqwest::Error> {
     if err.is_timeout() || err.is_connect() {

crates/common/download/src/download.rs

@@ -1,5 +1,6 @@
 use super::*;
 use axum::Router;
+use http::StatusCode;
 use hyper::header::AUTHORIZATION;
 use rustls::pki_types::pem::PemObject;
 use rustls::pki_types::PrivateKeyDer;
@@ -433,6 +434,29 @@ async fn downloader_error_shows_certificate_required_error_when_appropriate() {
     assert!(dbg!(format!("{err:#}")).contains("received fatal alert: CertificateRequired"));
 }
 
+#[test_case::test_case(Some(r#""xyzzy"#), Some(r#""xyzzy"#), false)]
+#[test_case::test_case(Some(r#"W/"xyzzy"#), Some(r#""xyzzy"#), true)]
+#[test_case::test_case(Some(r#""xyzzy"#), Some(r#"W/"xyzzy"#), true)]
+#[test_case::test_case(Some(r#""xyzzy1"#), Some(r#""xyzzy2"#), true)]
+#[test_case::test_case(None, None, false)]
+#[test_case::test_case(Some(r#""xyzzy1"#), None, true)]
+#[test_case::test_case(None, Some(r#""xyzzy2"#), true)]
+fn verifies_etags(etag1: Option<&'static str>, etag2: Option<&'static str>, modified: bool) {
+    let mut response1 = http::Response::builder().status(StatusCode::PARTIAL_CONTENT);
+    if let Some(etag) = etag1 {
+        response1 = response1.header(http::header::ETAG, etag);
+    }
+    let response1 = response1.body("").unwrap().into();
+
+    let mut response2 = http::Response::builder().status(StatusCode::PARTIAL_CONTENT);
+    if let Some(etag) = etag2 {
+        response2 = response2.header(http::header::ETAG, etag);
+    }
+    let response2 = response2.body("").unwrap().into();
+
+    assert_eq!(was_resource_modified(&response1, &response2), modified);
+}
+
 fn create_file_with_size(size: usize) -> Result<NamedTempFile, anyhow::Error> {
     let mut file = NamedTempFile::new().unwrap();
     let data: String = "Some data!".into();

crates/common/download/src/download/tests.rs

@@ -102,3 +102,109 @@ async fn resume_download_when_disconnected() {
 
     server_task.abort();
 }
+
+/// If after retrying ETag of the resource is different, we should download it
+/// from scratch again.
+#[tokio::test]
+async fn resume_download_with_etag_changed() {
+    let file_v1 = "AAAABBBBCCCCDDDD";
+    let file_v2 = "XXXXYYYYZZZZWWWW";
+    let chunk_size = 4;
+
+    let listener = TcpListener::bind("localhost:0").await.unwrap();
+    let port = listener.local_addr().unwrap().port();
+
+    // server should do 3 things:
+    // - only send a portion of the first request
+    // - after the first request update the resource
+    // - serve 2nd request normally (but we expect it to be a range request)
+    let server_task = tokio::spawn(async move {
+        let mut request_count = 0;
+        while let Ok((mut stream, _addr)) = listener.accept().await {
+            let response_task = async move {
+                let (reader, mut writer) = stream.split();
+                let mut lines = BufReader::new(reader).lines();
+                let mut range: Option<std::ops::Range<usize>> = None;
+
+                // We got an HTTP request, read the lines of the request
+                while let Ok(Some(line)) = lines.next_line().await {
+                    if line.to_ascii_lowercase().contains("range:") {
+                        let (_, bytes) = line.split_once('=').unwrap();
+                        let (start, end) = bytes.split_once('-').unwrap();
+                        let start = start.parse().unwrap_or(0);
+                        let end = end.parse().unwrap_or(file_v2.len());
+                        range = Some(start..end)
+                    }
+                    // On `\r\n\r\n` (empty line) stop reading the request
+                    // and start responding
+                    if line.is_empty() {
+                        break;
+                    }
+                }
+
+                let file = if request_count == 0 { file_v1 } else { file_v2 };
+                let etag = if request_count == 0 {
+                    "v1-initial"
+                } else {
+                    "v2-changed"
+                };
+
+                if let Some(range) = range {
+                    // Return range for both first and subsequent requests
+                    let start = range.start;
+                    let end = range.end.min(file.len());
+                    let header = format!(
+                        "HTTP/1.1 206 Partial Content\r\n\
+                        transfer-encoding: chunked\r\n\
+                        connection: close\r\n\
+                        content-type: application/octet-stream\r\n\
+                        content-range: bytes {start}-{end}/*\r\n\
+                        accept-ranges: bytes\r\n\
+                        etag: \"{etag}\"\r\n"
+                    );
+                    let body = &file[start..end];
+                    let size = body.len();
+                    let msg = format!("{header}\r\n{size:x}\r\n{body}\r\n0\r\n\r\n");
+                    writer.write_all(msg.as_bytes()).await.unwrap();
+                    writer.flush().await.unwrap();
+                } else {
+                    let header = format!(
+                        "HTTP/1.1 200 OK\r\n\
+                        transfer-encoding: chunked\r\n\
+                        connection: close\r\n\
+                        content-type: application/octet-stream\r\n\
+                        accept-ranges: bytes\r\n\
+                        etag: \"{etag}\"\r\n"
+                    );
+
+                    let body = &file[0..chunk_size];
+                    let size = body.len();
+                    let msg = format!("{header}\r\n{size:x}\r\n{body}\r\n");
+                    writer.write_all(msg.as_bytes()).await.unwrap();
+                    writer.flush().await.unwrap();
+                    // Connection drops here without sending final chunk
+                }
+            };
+            request_count += 1;
+            tokio::spawn(response_task);
+        }
+    });
+
+    // Wait until task binds a listener on the TCP port
+    tokio::time::sleep(std::time::Duration::from_millis(50)).await;
+
+    let tmpdir = TempDir::new().unwrap();
+    let target_path = tmpdir.path().join("partial_download_etag");
+
+    let downloader = Downloader::new(target_path, None, CloudHttpConfig::test_value());
+    let url = DownloadInfo::new(&format!("http://localhost:{port}/"));
+
+    downloader.download(&url).await.unwrap();
+    let saved_file = std::fs::read_to_string(downloader.filename()).unwrap();
+    // Should have the complete new file content since ETag changed
+    assert_eq!(saved_file, file_v2);
+
+    downloader.cleanup().await.unwrap();
+
+    server_task.abort();
+}