From 1bbb4c71b7dc553f782353438405f5f867f81f36 Mon Sep 17 00:00:00 2001
From: James <namnh0122@gmail.com>
Date: Fri, 20 Sep 2024 14:17:47 +0700
Subject: [PATCH] feat: add download resume

---
 engine/services/download_service.cc | 65 ++++++++++++++++++++++++-----
 engine/services/download_service.h  |  6 ++-
 2 files changed, 59 insertions(+), 12 deletions(-)

diff --git a/engine/services/download_service.cc b/engine/services/download_service.cc
index 08e37958c..1cf8b68c4 100644
--- a/engine/services/download_service.cc
+++ b/engine/services/download_service.cc
@@ -1,14 +1,15 @@
+#include "download_service.h"
 #include <curl/curl.h>
 #include <httplib.h>
 #include <stdio.h>
 #include <trantor/utils/Logger.h>
 #include <filesystem>
+#include <ostream>
 #include <thread>
-
-#include "download_service.h"
 #include "exceptions/failed_curl_exception.h"
 #include "exceptions/failed_init_curl_exception.h"
 #include "exceptions/failed_open_file_exception.h"
+#include "utils/format_utils.h"
 #include "utils/logging_utils.h"
 
 namespace {
@@ -19,14 +20,15 @@ size_t WriteCallback(void* ptr, size_t size, size_t nmemb, FILE* stream) {
 }  // namespace
 
 void DownloadService::AddDownloadTask(
-    const DownloadTask& task,
-    std::optional<OnDownloadTaskSuccessfully> callback) {
+    DownloadTask& task, std::optional<OnDownloadTaskSuccessfully> callback) {
   CLI_LOG("Validating download items, please wait..");
   // preprocess to check if all the item are valid
   auto total_download_size{0};
-  for (const auto& item : task.items) {
+  for (auto& item : task.items) {
     try {
-      total_download_size += GetFileSize(item.downloadUrl);
+      auto size = GetFileSize(item.downloadUrl);
+      item.bytes = size;
+      total_download_size += size;
     } catch (const FailedCurlException& e) {
       CTL_ERR("Found invalid download item: " << item.downloadUrl << " - "
                                               << e.what());
@@ -37,7 +39,7 @@ void DownloadService::AddDownloadTask(
   // all items are valid, start downloading
   for (const auto& item : task.items) {
     CLI_LOG("Start downloading: " + item.localPath.filename().string());
-    Download(task.id, item);
+    Download(task.id, item, true);
   }
 
   if (callback.has_value()) {
@@ -76,7 +78,7 @@ void DownloadService::AddAsyncDownloadTask(
 
   for (const auto& item : task.items) {
     std::thread([this, task, &callback, item]() {
-      this->Download(task.id, item);
+      this->Download(task.id, item, false);
     }).detach();
   }
 
@@ -84,7 +86,8 @@ void DownloadService::AddAsyncDownloadTask(
 }
 
 void DownloadService::Download(const std::string& download_id,
-                               const DownloadItem& download_item) {
+                               const DownloadItem& download_item,
+                               bool allow_resume) {
   CTL_INF("Absolute file output: " << download_item.localPath.string());
 
   CURL* curl;
@@ -96,7 +99,43 @@ void DownloadService::Download(const std::string& download_id,
     throw FailedInitCurlException();
   }
 
-  file = fopen(download_item.localPath.string().c_str(), "wb");
+  std::string mode = "wb";
+  if (allow_resume && std::filesystem::exists(download_item.localPath) &&
+      download_item.bytes.has_value()) {
+    FILE* existing_file = fopen(download_item.localPath.string().c_str(), "r");
+    fseek(existing_file, 0, SEEK_END);
+    curl_off_t existing_file_size = ftell(existing_file);
+    fclose(existing_file);
+    auto missing_bytes = download_item.bytes.value() - existing_file_size;
+    if (missing_bytes > 0) {
+      CLI_LOG("Found unfinished download! Additional "
+              << format_utils::BytesToHumanReadable(missing_bytes)
+              << " need to be downloaded.");
+      std::cout << "Continue download [Y/n]: " << std::flush;
+      std::string answer{""};
+      std::cin >> answer;
+      if (answer == "Y" || answer == "y" || answer.empty()) {
+        mode = "ab";
+        CLI_LOG("Resuming download..");
+      } else {
+        CLI_LOG("Start over..");
+      }
+    } else {
+      CLI_LOG(download_item.localPath.filename().string()
+              << " is already downloaded!");
+      std::cout << "Re-download? [Y/n]: " << std::flush;
+
+      std::string answer = "";
+      std::cin >> answer;
+      if (answer == "Y" || answer == "y" || answer.empty()) {
+        CLI_LOG("Re-downloading..");
+      } else {
+        return;
+      }
+    }
+  }
+
+  file = fopen(download_item.localPath.string().c_str(), mode.c_str());
   if (!file) {
     auto err_msg{"Failed to open output file " +
                  download_item.localPath.string()};
@@ -109,6 +148,12 @@ void DownloadService::Download(const std::string& download_id,
   curl_easy_setopt(curl, CURLOPT_NOPROGRESS, 0L);
   curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
 
+  if (mode == "ab") {
+    fseek(file, 0, SEEK_END);
+    curl_off_t local_file_size = ftell(file);
+    curl_easy_setopt(curl, CURLOPT_RESUME_FROM_LARGE, local_file_size);
+  }
+
   res = curl_easy_perform(curl);
 
   if (res != CURLE_OK) {
diff --git a/engine/services/download_service.h b/engine/services/download_service.h
index 5015ec29c..7063be74c 100644
--- a/engine/services/download_service.h
+++ b/engine/services/download_service.h
@@ -20,6 +20,8 @@ struct DownloadItem {
 
   std::optional<std::string> checksum;
 
+  std::optional<uint64_t> bytes;
+
   std::string ToString() const {
     std::ostringstream output;
     output << "DownloadItem{id: " << id << ", downloadUrl: " << downloadUrl
@@ -54,7 +56,7 @@ class DownloadService {
       std::function<void(const DownloadTask& task)>;
 
   void AddDownloadTask(
-      const DownloadTask& task,
+      DownloadTask& task,
       std::optional<OnDownloadTaskSuccessfully> callback = std::nullopt);
 
   void AddAsyncDownloadTask(
@@ -70,5 +72,5 @@ class DownloadService {
 
  private:
   void Download(const std::string& download_id,
-                const DownloadItem& download_item);
+                const DownloadItem& download_item, bool allow_resume);
 };