From 9bfb896812f47c667048090c0aec3c207d162fba Mon Sep 17 00:00:00 2001
From: nek0R1n <37893549+Sayaka617@users.noreply.github.com>
Date: Mon, 17 Jul 2023 08:54:05 +0800
Subject: [PATCH] [#3850] Fix the problem when set content-type by
 ResponseEntity.header() (#3852)

---
 .../demo/springmvc/client/TestDownloadSchema.java | 15 +++++++++++++++
 .../demo/springmvc/server/DownloadSchema.java     | 11 +++++++++++
 .../foundation/vertx/http/DownloadUtils.java      |  6 +++++-
 3 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/demo/demo-springmvc/springmvc-client/src/main/java/org/apache/servicecomb/demo/springmvc/client/TestDownloadSchema.java b/demo/demo-springmvc/springmvc-client/src/main/java/org/apache/servicecomb/demo/springmvc/client/TestDownloadSchema.java
index 37a24116a4..40214c1f5d 100644
--- a/demo/demo-springmvc/springmvc-client/src/main/java/org/apache/servicecomb/demo/springmvc/client/TestDownloadSchema.java
+++ b/demo/demo-springmvc/springmvc-client/src/main/java/org/apache/servicecomb/demo/springmvc/client/TestDownloadSchema.java
@@ -21,9 +21,13 @@
 import org.apache.servicecomb.demo.TestMgr;
 import org.apache.servicecomb.foundation.vertx.http.ReadStreamPart;
 import org.apache.servicecomb.provider.springmvc.reference.RestTemplateBuilder;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.ResponseEntity;
 import org.springframework.stereotype.Component;
 import org.springframework.web.client.RestTemplate;
 
+import java.util.Collections;
+
 @Component
 public class TestDownloadSchema implements CategorizedTestCase {
   @Override
@@ -31,6 +35,7 @@ public void testRestTransport() throws Exception {
     testDownloadFileAndDeleted();
     testDownloadFileNotDeleted();
     testDownloadFileWithNull();
+    testSetContentTypeByResponseEntity();
   }
 
   private void testDownloadFileAndDeleted() throws Exception {
@@ -69,4 +74,14 @@ private void testDownloadFileNotDeleted() throws Exception {
         .getForObject("servicecomb://springmvc/download/assertLastFileDeleted", boolean.class);
     TestMgr.check(exists, true);
   }
+
+  private void testSetContentTypeByResponseEntity() throws Exception {
+    RestTemplate restTemplate = RestTemplateBuilder.create();
+    ResponseEntity<ReadStreamPart> responseEntity = restTemplate
+            .getForEntity("servicecomb://springmvc/download/setContentTypeByResponseEntity?content=hello&contentType=customType",
+                    ReadStreamPart.class);
+    String hello = responseEntity.getBody().saveAsString().get();
+    TestMgr.check(responseEntity.getHeaders().get(HttpHeaders.CONTENT_TYPE), Collections.singletonList("customType"));
+    TestMgr.check(hello, "hello");
+  }
 }
diff --git a/demo/demo-springmvc/springmvc-server/src/main/java/org/apache/servicecomb/demo/springmvc/server/DownloadSchema.java b/demo/demo-springmvc/springmvc-server/src/main/java/org/apache/servicecomb/demo/springmvc/server/DownloadSchema.java
index c75775f24e..8eed21c19e 100644
--- a/demo/demo-springmvc/springmvc-server/src/main/java/org/apache/servicecomb/demo/springmvc/server/DownloadSchema.java
+++ b/demo/demo-springmvc/springmvc-server/src/main/java/org/apache/servicecomb/demo/springmvc/server/DownloadSchema.java
@@ -92,6 +92,17 @@ public ResponseEntity<Part> notDeleteAfterFinished(@RequestParam("content") Stri
         .body(new FilePart(null, file));
   }
 
+  @GetMapping(path = "/setContentTypeByResponseEntity")
+  public ResponseEntity<Part> setContentTypeByResponseEntity(@RequestParam("content") String content, @RequestParam("contentType") String contentType) throws IOException {
+    File file = createTempFile(content);
+
+    return ResponseEntity
+            .ok()
+            .header(HttpHeaders.CONTENT_TYPE, contentType)
+            .header(HttpHeaders.CONTENT_DISPOSITION, "attachment;filename=tempFileEntity.txt")
+            .body(new FilePart(null, file));
+  }
+
   @GetMapping(path = "/assertLastFileDeleted")
   public boolean assertLastFileDeleted() {
     return lastFile.exists();
diff --git a/foundations/foundation-vertx/src/main/java/org/apache/servicecomb/foundation/vertx/http/DownloadUtils.java b/foundations/foundation-vertx/src/main/java/org/apache/servicecomb/foundation/vertx/http/DownloadUtils.java
index f37712e34e..e83e65e647 100644
--- a/foundations/foundation-vertx/src/main/java/org/apache/servicecomb/foundation/vertx/http/DownloadUtils.java
+++ b/foundations/foundation-vertx/src/main/java/org/apache/servicecomb/foundation/vertx/http/DownloadUtils.java
@@ -44,7 +44,11 @@ public static void prepareDownloadHeader(HttpServletResponseEx responseEx, Part
       return;
     }
     if (responseEx.getHeader(HttpHeaders.CONTENT_TYPE.toString()) == null) {
-      responseEx.setHeader(HttpHeaders.CONTENT_TYPE.toString(), part.getContentType());
+      if (responseEx.getContentType() != null) {
+        responseEx.setHeader(HttpHeaders.CONTENT_TYPE.toString(), responseEx.getContentType());
+      } else {
+        responseEx.setHeader(HttpHeaders.CONTENT_TYPE.toString(), part.getContentType());
+      }
     }
 
     if (responseEx.getHeader(javax.ws.rs.core.HttpHeaders.CONTENT_DISPOSITION) == null) {