From 66dc4e3009e8355faf6bf61aee364c7df73a9b7a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Cortier?= <bcortier@proton.me>
Date: Mon, 8 May 2023 13:21:34 -0400
Subject: [PATCH] fix(dgw): ensure mime-type is set in /jet/jrec/pull (#431)

---
 Cargo.lock                          | 26 ++++++++++++++++++++++++++
 devolutions-gateway/Cargo.toml      |  2 +-
 devolutions-gateway/src/api/jrec.rs | 20 ++++++++++++++------
 3 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 1194696aa..2db27e5c2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1877,6 +1877,16 @@ version = "0.3.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
 
+[[package]]
+name = "mime_guess"
+version = "2.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4192263c238a5f0d0c6bfd21f336a313a4ce1c450542449ca191bb657b4642ef"
+dependencies = [
+ "mime",
+ "unicase",
+]
+
 [[package]]
 name = "minimal-lexical"
 version = "0.2.1"
@@ -3863,9 +3873,16 @@ dependencies = [
  "http",
  "http-body",
  "http-range-header",
+ "httpdate",
+ "mime",
+ "mime_guess",
+ "percent-encoding",
  "pin-project-lite",
+ "tokio",
+ "tokio-util",
  "tower-layer",
  "tower-service",
+ "tracing",
 ]
 
 [[package]]
@@ -4021,6 +4038,15 @@ version = "0.1.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94"
 
+[[package]]
+name = "unicase"
+version = "2.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "50f37be617794602aabbeee0be4f259dc1778fabe05e2d67ee8f79326d5cb4f6"
+dependencies = [
+ "version_check",
+]
+
 [[package]]
 name = "unicode-bidi"
 version = "0.3.13"
diff --git a/devolutions-gateway/Cargo.toml b/devolutions-gateway/Cargo.toml
index 023573f2f..776995d59 100644
--- a/devolutions-gateway/Cargo.toml
+++ b/devolutions-gateway/Cargo.toml
@@ -79,7 +79,7 @@ ngrok = "0.11.3"
 hyper = "0.14.26"
 axum = { version = "0.6.18", default-features = false, features = ["http1", "json", "ws", "query", "tracing", "tower-log", "headers"] }
 axum-extra = { version = "0.7.4", features = ["query", "async-read-body"] }
-tower-http = { version = "0.4.0", features = ["cors"] }
+tower-http = { version = "0.4.0", features = ["cors", "fs"] }
 
 # OpenAPI generator
 utoipa = { version = "3.3.0", default-features = false, features = ["uuid", "chrono"], optional = true }
diff --git a/devolutions-gateway/src/api/jrec.rs b/devolutions-gateway/src/api/jrec.rs
index 398c41327..2baeb6d55 100644
--- a/devolutions-gateway/src/api/jrec.rs
+++ b/devolutions-gateway/src/api/jrec.rs
@@ -6,10 +6,9 @@ use std::sync::Arc;
 use anyhow::Context as _;
 use axum::extract::ws::WebSocket;
 use axum::extract::{self, ConnectInfo, Query, State, WebSocketUpgrade};
-use axum::response::{IntoResponse as _, Response};
+use axum::response::Response;
 use axum::routing::get;
 use axum::{Json, Router};
-use tokio::fs::File;
 use tracing::Instrument as _;
 use uuid::Uuid;
 
@@ -141,11 +140,17 @@ pub(crate) async fn list_recordings(
     ),
     security(("jrec_token" = ["pull"])),
 ))]
-pub(crate) async fn pull_recording_file(
+pub(crate) async fn pull_recording_file<ReqBody>(
     State(DgwState { conf_handle, .. }): State<DgwState>,
     extract::Path((id, filename)): extract::Path<(Uuid, String)>,
     JrecToken(claims): JrecToken,
-) -> Result<Response, HttpError> {
+    request: axum::http::Request<ReqBody>,
+) -> Result<Response<tower_http::services::fs::ServeFileSystemResponseBody>, HttpError>
+where
+    ReqBody: Send + 'static,
+{
+    use tower::ServiceExt as _;
+
     if filename.contains("..") || filename.contains('/') || filename.contains('\\') {
         return Err(HttpError::bad_request().msg("invalid file name"));
     }
@@ -164,7 +169,10 @@ pub(crate) async fn pull_recording_file(
         return Err(HttpError::not_found().msg("requested file does not exist"));
     }
 
-    let file = File::open(path).await.map_err(HttpError::internal().err())?;
+    let response = tower_http::services::ServeFile::new(path)
+        .oneshot(request)
+        .await
+        .map_err(HttpError::internal().err())?;
 
-    Ok(axum_extra::body::AsyncReadBody::new(file).into_response())
+    Ok(response)
 }