From c348eb31ea6eefea919f46cce188434e2edba346 Mon Sep 17 00:00:00 2001
From: Fabian-Lars <fabianlars@fabianlars.de>
Date: Wed, 15 Feb 2023 13:06:43 +0100
Subject: [PATCH] feat(upload): Add function to download files (#89)

---
 Cargo.lock                       |  2 +-
 plugins/upload/Cargo.toml        |  2 +-
 plugins/upload/guest-js/index.ts | 35 ++++++++++++++++++++---
 plugins/upload/src/lib.rs        | 48 ++++++++++++++++++++++++++++++--
 4 files changed, 78 insertions(+), 9 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 75568883..ebd4ec1a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4422,7 +4422,7 @@ dependencies = [
 name = "tauri-plugin-upload"
 version = "0.1.0"
 dependencies = [
- "futures",
+ "futures-util",
  "log",
  "read-progress-stream",
  "reqwest",
diff --git a/plugins/upload/Cargo.toml b/plugins/upload/Cargo.toml
index 26ec517e..985645d6 100644
--- a/plugins/upload/Cargo.toml
+++ b/plugins/upload/Cargo.toml
@@ -18,5 +18,5 @@ thiserror.workspace = true
 tokio = { version = "1", features = [ "fs" ] }
 tokio-util = { version = "0.7", features = [ "codec" ] }
 reqwest = { version = "0.11", features = [ "json", "stream" ] }
-futures = "0.3"
+futures-util = "0.3"
 read-progress-stream = "1.0.0"
\ No newline at end of file
diff --git a/plugins/upload/guest-js/index.ts b/plugins/upload/guest-js/index.ts
index 88f5202e..7b4a94c7 100644
--- a/plugins/upload/guest-js/index.ts
+++ b/plugins/upload/guest-js/index.ts
@@ -11,12 +11,12 @@ type ProgressHandler = (progress: number, total: number) => void;
 const handlers: Map<number, ProgressHandler> = new Map();
 let listening = false;
 
-async function listenToUploadEventIfNeeded(): Promise<void> {
+async function listenToEventIfNeeded(event: string): Promise<void> {
   if (listening) {
     return await Promise.resolve();
   }
   return await appWindow
-    .listen<ProgressPayload>("upload://progress", ({ payload }) => {
+    .listen<ProgressPayload>(event, ({ payload }) => {
       const handler = handlers.get(payload.id);
       if (handler != null) {
         handler(payload.progress, payload.total);
@@ -27,7 +27,7 @@ async function listenToUploadEventIfNeeded(): Promise<void> {
     });
 }
 
-export default async function upload(
+async function upload(
   url: string,
   filePath: string,
   progressHandler?: ProgressHandler,
@@ -41,7 +41,7 @@ export default async function upload(
     handlers.set(id, progressHandler);
   }
 
-  await listenToUploadEventIfNeeded();
+  await listenToEventIfNeeded("upload://progress");
 
   await invoke("plugin:upload|upload", {
     id,
@@ -50,3 +50,30 @@ export default async function upload(
     headers: headers ?? {},
   });
 }
+
+async function download(
+  url: string,
+  filePath: string,
+  progressHandler?: ProgressHandler,
+  headers?: Map<string, string>
+): Promise<void> {
+  const ids = new Uint32Array(1);
+  window.crypto.getRandomValues(ids);
+  const id = ids[0];
+
+  if (progressHandler != null) {
+    handlers.set(id, progressHandler);
+  }
+
+  await listenToEventIfNeeded("download://progress");
+
+  await invoke("plugin:upload|upload", {
+    id,
+    url,
+    filePath,
+    headers: headers ?? {},
+  });
+}
+
+export default upload;
+export { download, upload };
diff --git a/plugins/upload/src/lib.rs b/plugins/upload/src/lib.rs
index a83a87ff..f57683c6 100644
--- a/plugins/upload/src/lib.rs
+++ b/plugins/upload/src/lib.rs
@@ -2,14 +2,14 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-License-Identifier: MIT
 
-use futures::TryStreamExt;
+use futures_util::TryStreamExt;
 use serde::{ser::Serializer, Serialize};
 use tauri::{
     command,
     plugin::{Builder as PluginBuilder, TauriPlugin},
     Runtime, Window,
 };
-use tokio::fs::File;
+use tokio::{fs::File, io::AsyncWriteExt};
 use tokio_util::codec::{BytesCodec, FramedRead};
 
 use read_progress_stream::ReadProgressStream;
@@ -24,6 +24,8 @@ pub enum Error {
     Io(#[from] std::io::Error),
     #[error(transparent)]
     Request(#[from] reqwest::Error),
+    #[error("{0}")]
+    ContentLength(String),
 }
 
 impl Serialize for Error {
@@ -42,6 +44,46 @@ struct ProgressPayload {
     total: u64,
 }
 
+#[command]
+async fn download<R: Runtime>(
+    window: Window<R>,
+    id: u32,
+    url: &str,
+    file_path: &str,
+    headers: HashMap<String, String>,
+) -> Result<u32> {
+    let client = reqwest::Client::new();
+
+    let mut request = client.get(url);
+    // Loop trought the headers keys and values
+    // and add them to the request object.
+    for (key, value) in headers {
+        request = request.header(&key, value);
+    }
+
+    let response = request.send().await?;
+    let total = response.content_length().ok_or_else(|| {
+        Error::ContentLength(format!("Failed to get content length from '{}'", url))
+    })?;
+
+    let mut file = File::create(file_path).await?;
+    let mut stream = response.bytes_stream();
+
+    while let Some(chunk) = stream.try_next().await? {
+        file.write_all(&chunk).await?;
+        let _ = window.emit(
+            "download://progress",
+            ProgressPayload {
+                id,
+                progress: chunk.len() as u64,
+                total,
+            },
+        );
+    }
+
+    Ok(id)
+}
+
 #[command]
 async fn upload<R: Runtime>(
     window: Window<R>,
@@ -88,6 +130,6 @@ fn file_to_body<R: Runtime>(id: u32, window: Window<R>, file: File) -> reqwest::
 
 pub fn init<R: Runtime>() -> TauriPlugin<R> {
     PluginBuilder::new("upload")
-        .invoke_handler(tauri::generate_handler![upload])
+        .invoke_handler(tauri::generate_handler![download, upload])
         .build()
 }