fix content disposition header

Snd-R · Snd-R · commit fbfd60c4eebe · 2023-01-03T22:01:23.000+03:00
diff --git a/Cargo.toml b/Cargo.toml
@@ -24,3 +24,4 @@ serde = "1.0.147"
 serde_derive = "1.0.147"
 once_cell = "1.10"
 async-compression = { version = "0.3.15", features = ["tokio", "brotli", "gzip", "deflate"] }
+regex = "1"
diff --git a/src/routes.rs b/src/routes.rs
@@ -1,6 +1,9 @@
+use std::path::Path;
+
 use bytes::Bytes;
 use image::ImageFormat;
 use log::info;
+use regex::Regex;
 use tokio::task;
 use unicase::Ascii;
 use warp::{Filter, http, Rejection};
@@ -49,7 +52,7 @@ async fn proxy_upscale_and_forward(
     headers: HeaderMap,
     body: Bytes,
 ) -> Result<Response<Body>, Rejection> {
-    let uri_str = format!("{} {}", method, uri.as_str());
+    let uri_str = format!("{} {}{}", method, uri.as_str(), params.as_deref().map(|query| "?".to_string() + query).unwrap_or_default());
     // TODO do not request compressed data to avoid decode and re-encode
     let response = proxy_to_and_forward_response(proxy_address, base_path, uri, params, method, headers, body).await?;
     let status = response.status();
@@ -104,17 +107,26 @@ async fn response_to_upscaled_reply(
     format: ImageFormat,
 ) -> Result<Response<Body>, errors::Error> {
     let mime_type = match format {
-        ImageFormat::Png => { Some("image/png") }
-        ImageFormat::Jpeg => { Some("image/jpeg") }
-        ImageFormat::WebP => { Some("image/webp") }
+        ImageFormat::Png => { Some(("image/png", "png")) }
+        ImageFormat::Jpeg => { Some(("image/jpeg", "jpeg")) }
+        ImageFormat::WebP => { Some(("image/webp", "webp")) }
         _ => { None }
     };
     let mut builder = http::Response::builder();
     for (k, v) in headers {
         if Ascii::new("Content-Length") == k {
             builder = builder.header("Content-Length", bytes.len())
         } else if Ascii::new("Content-Type") == k && mime_type.is_some() {
-            builder = builder.header("Content-Type", mime_type.unwrap())
+            builder = builder.header("Content-Type", mime_type.unwrap().0)
+        } else if Ascii::new("Content-Disposition") == k && mime_type.is_some() {
+            let new_value: String = v.to_str().unwrap().split("; ")
+                .map(|param| if param.starts_with("filename=") {
+                    with_new_file_extension(param, mime_type.unwrap().1)
+                } else if param.starts_with("filename*=") {
+                    with_new_file_extension(param, mime_type.unwrap().1)
+                } else { param.to_string() })
+                .collect::<Vec<String>>().join("; ");
+            builder = builder.header("Content-Disposition", new_value)
         } else {
             builder = builder.header(k, v);
         }
@@ -125,6 +137,16 @@ async fn response_to_upscaled_reply(
         .map_err(errors::Error::Http)
 }
 
+fn with_new_file_extension(name: &str, extension: &str) -> String {
+    let regex = Regex::new(r"(filename\*=UTF-8''|filename=)(.+\b)").unwrap();
+    let captures = regex.captures(name).unwrap();
+    let param_name = captures.get(1).unwrap().as_str();
+    let filename = captures.get(2).unwrap().as_str();
+    let new_filename = Path::new(filename).with_extension(extension)
+        .into_os_string().into_string().unwrap();
+    format!("{}{}", param_name, new_filename)
+}
+
 fn unwrap_encoding_header(encoding: &HeaderValue) -> Algorithm {
     let encoding = encoding.to_str().unwrap();
     match encoding {
