From a7957a99bbefae5dd0cd00f229b6c3f5c2ad04fe Mon Sep 17 00:00:00 2001
From: Gaius <gaius.qi@gmail.com>
Date: Tue, 30 Jan 2024 12:59:19 +0800
Subject: [PATCH] feat: if download by range, add content-range header to
 resopnse header (#244)

Signed-off-by: Gaius <gaius.qi@gmail.com>
---
 Cargo.lock      | 29 +++++++++++++++++++++++++++--
 Cargo.toml      |  1 +
 src/lib.rs      |  4 ++++
 src/task/mod.rs | 17 ++++++++++++++++-
 4 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 4c8bf52f..88e8cbd5 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -553,6 +553,7 @@ dependencies = [
  "futures",
  "futures-util",
  "hashring",
+ "headers 0.4.0",
  "hex",
  "home",
  "hostname",
@@ -916,13 +917,28 @@ checksum = "06683b93020a07e3dbcf5f8c0f6d40080d725bea7936fc01ad345c01b97dc270"
 dependencies = [
  "base64",
  "bytes",
- "headers-core",
+ "headers-core 0.2.0",
  "http 0.2.11",
  "httpdate",
  "mime",
  "sha1",
 ]
 
+[[package]]
+name = "headers"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "322106e6bd0cba2d5ead589ddb8150a13d7c4217cf80d7c4f682ca994ccc6aa9"
+dependencies = [
+ "base64",
+ "bytes",
+ "headers-core 0.3.0",
+ "http 1.0.0",
+ "httpdate",
+ "mime",
+ "sha1",
+]
+
 [[package]]
 name = "headers-core"
 version = "0.2.0"
@@ -932,6 +948,15 @@ dependencies = [
  "http 0.2.11",
 ]
 
+[[package]]
+name = "headers-core"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "54b4a22553d4242c49fddb9ba998a99962b5cc6f22cb5a3482bec22522403ce4"
+dependencies = [
+ "http 1.0.0",
+]
+
 [[package]]
 name = "heck"
 version = "0.4.1"
@@ -3252,7 +3277,7 @@ dependencies = [
  "bytes",
  "futures-channel",
  "futures-util",
- "headers",
+ "headers 0.3.9",
  "http 0.2.11",
  "hyper 0.14.28",
  "log",
diff --git a/Cargo.toml b/Cargo.toml
index 1aa36bd4..f1204614 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -86,3 +86,4 @@ http-body-util = "0.1.0"
 regex = "1.10.2"
 http-range-header = "0.4.0"
 futures-util = "0.3.30"
+headers = "0.4.0"
diff --git a/src/lib.rs b/src/lib.rs
index 2e78a5c9..d1cb0d9c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -61,6 +61,10 @@ pub enum Error {
     #[error(transparent)]
     TonicStatus(#[from] tonic::Status),
 
+    // HeaderError is the error for headers.
+    #[error(transparent)]
+    HeadersError(#[from] headers::Error),
+
     // Reqwest is the error for reqwest.
     #[error(transparent)]
     Reqwest(#[from] reqwest::Error),
diff --git a/src/task/mod.rs b/src/task/mod.rs
index 71ae8cb2..60ae334e 100644
--- a/src/task/mod.rs
+++ b/src/task/mod.rs
@@ -186,6 +186,21 @@ impl Task {
             return Err(Error::InvalidContentLength());
         };
 
+        // If download range is specified, insert the content range header
+        // to the resopnse header.
+        let mut response_header = task.response_header.clone();
+        if let Some(range) = download.range.clone() {
+            response_header.insert(
+                reqwest::header::CONTENT_RANGE.to_string(),
+                format!(
+                    "bytes {}-{}/{}",
+                    range.start,
+                    range.start + range.length - 1,
+                    content_length
+                ),
+            );
+        }
+
         // Send the download task started request.
         download_progress_tx
             .send_timeout(
@@ -197,7 +212,7 @@ impl Task {
                         download_task_response::Response::DownloadTaskStartedResponse(
                             dfdaemon::v2::DownloadTaskStartedResponse {
                                 content_length,
-                                response_header: task.response_header.clone(),
+                                response_header,
                             },
                         ),
                     ),