fix(core): set correct mimetype for asset protocol streams, #5203 (#5536)

lucasfernog · web-flow · commit 9b1a6a1c02b8 · 2022-11-04T12:35:45.000-03:00
diff --git a/.changes/asset-protocol-streaming-mime-type.md b/.changes/asset-protocol-streaming-mime-type.md
@@ -0,0 +1,5 @@
+---
+"tauri": "patch"
+---
+
+Set the correct mimetype when streaming files through `asset:` protocol
diff --git a/core/tauri/src/manager.rs b/core/tauri/src/manager.rs
@@ -541,23 +541,39 @@ impl<R: Runtime> WindowManager<R> {
           .get("range")
           .and_then(|r| r.to_str().map(|r| r.to_string()).ok())
         {
-          let (headers, status_code, data) = crate::async_runtime::safe_block_on(async move {
-            let mut headers = HashMap::new();
-            let mut buf = Vec::new();
+          #[derive(Default)]
+          struct RangeMetadata {
+            file: Option<tokio::fs::File>,
+            range: Option<crate::runtime::http::HttpRange>,
+            metadata: Option<std::fs::Metadata>,
+            headers: HashMap<&'static str, String>,
+            status_code: u16,
+            body: Vec<u8>,
+          }
+
+          let mut range_metadata = crate::async_runtime::safe_block_on(async move {
+            let mut data = RangeMetadata::default();
             // open the file
             let mut file = match tokio::fs::File::open(path_.clone()).await {
               Ok(file) => file,
               Err(e) => {
                 debug_eprintln!("Failed to open asset: {}", e);
-                return (headers, 404, buf);
+                data.status_code = 404;
+                return data;
               }
             };
             // Get the file size
             let file_size = match file.metadata().await {
-              Ok(metadata) => metadata.len(),
+              Ok(metadata) => {
+                let len = metadata.len();
+                data.metadata.replace(metadata);
+                len
+              }
               Err(e) => {
                 debug_eprintln!("Failed to read asset metadata: {}", e);
-                return (headers, 404, buf);
+                data.file.replace(file);
+                data.status_code = 404;
+                return data;
               }
             };
             // parse the range
@@ -572,13 +588,16 @@ impl<R: Runtime> WindowManager<R> {
               Ok(r) => r,
               Err(e) => {
                 debug_eprintln!("Failed to parse range {}: {:?}", range, e);
-                return (headers, 400, buf);
+                data.file.replace(file);
+                data.status_code = 400;
+                return data;
               }
             };
 
             // FIXME: Support multiple ranges
             // let support only 1 range for now
-            let status_code = if let Some(range) = range.first() {
+            if let Some(range) = range.first() {
+              data.range.replace(*range);
               let mut real_length = range.length;
               // prevent max_length;
               // specially on webview2
@@ -592,38 +611,84 @@ impl<R: Runtime> WindowManager<R> {
               // who should be skipped on the header
               let last_byte = range.start + real_length - 1;
 
-              headers.insert("Connection", "Keep-Alive".into());
-              headers.insert("Accept-Ranges", "bytes".into());
-              headers.insert("Content-Length", real_length.to_string());
-              headers.insert(
+              data.headers.insert("Connection", "Keep-Alive".into());
+              data.headers.insert("Accept-Ranges", "bytes".into());
+              data
+                .headers
+                .insert("Content-Length", real_length.to_string());
+              data.headers.insert(
                 "Content-Range",
                 format!("bytes {}-{}/{}", range.start, last_byte, file_size),
               );
 
               if let Err(e) = file.seek(std::io::SeekFrom::Start(range.start)).await {
                 debug_eprintln!("Failed to seek file to {}: {}", range.start, e);
-                return (headers, 422, buf);
+                data.file.replace(file);
+                data.status_code = 422;
+                return data;
               }
 
-              if let Err(e) = file.take(real_length).read_to_end(&mut buf).await {
+              let mut f = file.take(real_length);
+              let r = f.read_to_end(&mut data.body).await;
+              file = f.into_inner();
+              data.file.replace(file);
+
+              if let Err(e) = r {
                 debug_eprintln!("Failed read file: {}", e);
-                return (headers, 422, buf);
+                data.status_code = 422;
+                return data;
               }
               // partial content
-              206
+              data.status_code = 206;
             } else {
-              200
-            };
+              data.status_code = 200;
+            }
 
-            (headers, status_code, buf)
+            data
           });
 
-          for (k, v) in headers {
+          for (k, v) in range_metadata.headers {
             response = response.header(k, v);
           }
 
-          let mime_type = MimeType::parse(&data, &path);
-          response.mimetype(&mime_type).status(status_code).body(data)
+          let mime_type = if let (Some(mut file), Some(metadata), Some(range)) = (
+            range_metadata.file,
+            range_metadata.metadata,
+            range_metadata.range,
+          ) {
+            // if we're already reading the beginning of the file, we do not need to re-read it
+            if range.start == 0 {
+              MimeType::parse(&range_metadata.body, &path)
+            } else {
+              let (status, bytes) = crate::async_runtime::safe_block_on(async move {
+                let mut status = None;
+                if let Err(e) = file.rewind().await {
+                  debug_eprintln!("Failed to rewind file: {}", e);
+                  status.replace(422);
+                  (status, Vec::with_capacity(0))
+                } else {
+                  // taken from https://docs.rs/infer/0.9.0/src/infer/lib.rs.html#240-251
+                  let limit = std::cmp::min(metadata.len(), 8192) as usize + 1;
+                  let mut bytes = Vec::with_capacity(limit);
+                  if let Err(e) = file.take(8192).read_to_end(&mut bytes).await {
+                    debug_eprintln!("Failed read file: {}", e);
+                    status.replace(422);
+                  }
+                  (status, bytes)
+                }
+              });
+              if let Some(s) = status {
+                range_metadata.status_code = s;
+              }
+              MimeType::parse(&bytes, &path)
+            }
+          } else {
+            MimeType::parse(&range_metadata.body, &path)
+          };
+          response
+            .mimetype(&mime_type)
+            .status(range_metadata.status_code)
+            .body(range_metadata.body)
         } else {
           match crate::async_runtime::safe_block_on(async move { tokio::fs::read(path_).await }) {
             Ok(data) => {
diff --git a/examples/streaming/README.md b/examples/streaming/README.md
@@ -3,3 +3,5 @@
 A simple Tauri Application showcase the streaming functionality.
 
 To execute run the following on the root directory of the repository: `cargo run --example streaming`.
+
+By default the example uses a custom URI scheme protocol. To use the builtin `asset` protocol, run `cargo run --example streaming --features protocol-asset`.
diff --git a/examples/streaming/index.html b/examples/streaming/index.html
@@ -1,28 +1,32 @@
 <!DOCTYPE html>
 <html lang="en">
-  <head>
-    <meta charset="UTF-8" />
-    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-    <style>
-      video {
-        width: 100vw;
-        height: 100vh;
-      }
-    </style>
-  </head>
 
-  <body>
-    <video id="video_source" controls="" autoplay="" name="media">
-      <source type="video/mp4" />
-    </video>
-    <script>
-      const { convertFileSrc } = window.__TAURI__.tauri
-      const video = document.getElementById('video_source')
-      const source = document.createElement('source')
+<head>
+  <meta charset="UTF-8" />
+  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+  <style>
+    video {
+      width: 100vw;
+      height: 100vh;
+    }
+  </style>
+</head>
+
+<body>
+  <video id="video_source" controls="" autoplay="" name="media">
+    <source type="video/mp4" />
+  </video>
+  <script>
+    const { invoke, convertFileSrc } = window.__TAURI__.tauri
+    const video = document.getElementById('video_source')
+    const source = document.createElement('source')
+    invoke('video_uri').then(([scheme, path]) => {
       source.type = 'video/mp4'
-      source.src = convertFileSrc('example/test_video.mp4', 'stream')
+      source.src = convertFileSrc(path, scheme)
       video.appendChild(source)
       video.load()
-    </script>
-  </body>
-</html>
+    })
+  </script>
+</body>
+
+</html>
diff --git a/examples/streaming/main.rs b/examples/streaming/main.rs
@@ -39,14 +39,15 @@ fn main() {
   }
 
   tauri::Builder::default()
+    .invoke_handler(tauri::generate_handler![video_uri])
     .register_uri_scheme_protocol("stream", move |_app, request| {
       // prepare our response
       let mut response = ResponseBuilder::new();
       // get the wanted path
       #[cfg(target_os = "windows")]
       let path = request.uri().strip_prefix("stream://localhost/").unwrap();
       #[cfg(not(target_os = "windows"))]
-      let path = request.uri().strip_prefix("stream://").unwrap();
+      let path = request.uri().strip_prefix("stream://localhost/").unwrap();
       let path = percent_encoding::percent_decode(path.as_bytes())
         .decode_utf8_lossy()
         .to_string();
@@ -117,3 +118,18 @@ fn main() {
     ))
     .expect("error while running tauri application");
 }
+
+// returns the scheme and the path of the video file
+// we're using this just to allow using the custom `stream` protocol or tauri built-in `asset` protocol
+#[tauri::command]
+fn video_uri() -> (&'static str, std::path::PathBuf) {
+  #[cfg(feature = "protocol-asset")]
+  {
+    let mut path = std::env::current_dir().unwrap();
+    path.push("test_video.mp4");
+    ("asset", path)
+  }
+
+  #[cfg(not(feature = "protocol-asset"))]
+  ("stream", "example/test_video.mp4".into())
+}
diff --git a/examples/streaming/tauri.conf.json b/examples/streaming/tauri.conf.json
@@ -38,7 +38,10 @@
       }
     },
     "allowlist": {
-      "all": false
+      "all": false,
+      "protocol": {
+        "assetScope": ["**/test_video.mp4"]
+      }
     },
     "windows": [
       {
@@ -50,7 +53,7 @@
       }
     ],
     "security": {
-      "csp": "default-src 'self'; media-src stream: https://stream.localhost"
+      "csp": "default-src 'self'; media-src stream: https://stream.localhost asset: https://asset.localhost"
     },
     "updater": {
       "active": false
