From ef72558ec0a63c9be3fb0311c855e50711d99e81 Mon Sep 17 00:00:00 2001
From: xrelkd <46590321+xrelkd@users.noreply.github.com>
Date: Fri, 8 Mar 2024 00:04:55 +0800
Subject: [PATCH] fix(crates/engine): mark progress as completed once received
 bytes >= length of chunk

---
 crates/engine/src/downloader/worker.rs | 18 ++++++++++++++----
 crates/engine/src/fetcher/http.rs      |  9 +++++++--
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/crates/engine/src/downloader/worker.rs b/crates/engine/src/downloader/worker.rs
index 7c47934c..15649f7d 100644
--- a/crates/engine/src/downloader/worker.rs
+++ b/crates/engine/src/downloader/worker.rs
@@ -69,12 +69,22 @@ impl Worker {
                                 .with_context(|_| error::SeekFileSnafu {
                                     file_path: file_path.clone(),
                                 })?;
-                            sink.write_all(bytes).await.with_context(|_| {
-                                error::WriteFileSnafu { file_path: file_path.clone() }
-                            })?;
+                            received += if bytes.len() as u64 + chunk.start <= chunk.end {
+                                sink.write_all(bytes).await.with_context(|_| {
+                                    error::WriteFileSnafu { file_path: file_path.clone() }
+                                })?;
+                                bytes.len() as u64
+                            } else {
+                                let n = usize::try_from(chunk.remaining()).unwrap_or_default();
+                                if n > 0 {
+                                    sink.write_all(&bytes[..n]).await.with_context(|_| {
+                                        error::WriteFileSnafu { file_path: file_path.clone() }
+                                    })?;
+                                }
+                                n as u64
+                            };
                             drop(sink);
                         }
-                        received += bytes.len() as u64;
                         progress_updater.update(id, chunk.start, chunk.end, received);
                     }
                     future::Either::Left((Ok(None), _)) => {
diff --git a/crates/engine/src/fetcher/http.rs b/crates/engine/src/fetcher/http.rs
index fd50fffe..1d8d9d04 100644
--- a/crates/engine/src/fetcher/http.rs
+++ b/crates/engine/src/fetcher/http.rs
@@ -92,8 +92,13 @@ impl ByteStream {
     pub async fn bytes(&mut self) -> Result<Option<&[u8]>> {
         match self.response.chunk().await.context(error::FetchBytesFromHttpSnafu) {
             Ok(Some(mut bytes)) => {
-                std::mem::swap(&mut self.buffer, &mut bytes);
-                Ok(Some(self.buffer.as_ref()))
+                if bytes.is_empty() {
+                    Ok(None)
+                } else {
+                    self.buffer.clear();
+                    std::mem::swap(&mut self.buffer, &mut bytes);
+                    Ok(Some(self.buffer.as_ref()))
+                }
             }
             Ok(None) => Ok(None),
             Err(err) => Err(err),