set response headers

Peergos · Jul 11, 2023 · 6242c48 · 6242c48
1 parent 40ac96b
commit 6242c48
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 5 deletions.
diff --git a/src/main/java/org/peergos/HttpProxyService.java b/src/main/java/org/peergos/HttpProxyService.java
@@ -70,7 +70,10 @@ public ProxyResponse proxyRequest(Multihash targetNodeId, ProxyRequest request)
         resp.content().readBytes(bout, contentLength);
         Map<String, String> headers = new HashMap<>();
         for (Map.Entry<String, String> entry: resp.headers().entries()) {
-            headers.put(entry.getKey(), entry.getValue());
+            String key = entry.getKey();
+            if (key != null) {
+                headers.put(key, entry.getValue());
+            }
         }
         int code = resp.status().code();
         resp.release();

diff --git a/src/main/java/org/peergos/client/RequestSender.java b/src/main/java/org/peergos/client/RequestSender.java
@@ -13,6 +13,7 @@
 import java.net.URL;
 import java.net.URLEncoder;
 import java.nio.charset.StandardCharsets;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -34,15 +35,22 @@ public static FullHttpResponse proxy(MultiAddress proxyTargetAddress, FullHttpRe
         FullHttpResponse httpResponse = new DefaultFullHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.valueOf(reply.statusCode),
                 reply.body.length > 0 ?
                         Unpooled.wrappedBuffer(reply.body) : Unpooled.buffer(0));
-        httpResponse.headers().set(HttpHeaderNames.CONTENT_LENGTH, reply.body.length);
+        for (Map.Entry<String, List<String>> entry : reply.responseHeaders.entrySet()) {
+            String key = entry.getKey();
+            if (key != null) {
+                httpResponse.headers().set(key, entry.getValue());
+            }
+        }
         return httpResponse;
     }
     public static class Response {
         public final byte[] body;
+        public final Map<String, List<String>> responseHeaders;
         public final int statusCode;
 
-        public Response(byte[] body, int statusCode) {
+        public Response(byte[] body, Map<String, List<String>> responseHeaders, int statusCode) {
             this.body = body;
+            this.responseHeaders = responseHeaders;
             this.statusCode = statusCode;
         }
     }
@@ -65,6 +73,7 @@ public static Response send(URL target, String method, byte[] body, Map<String,
         int r;
         while ((r = in.read(buf)) >= 0)
             resp.write(buf, 0, r);
-        return new Response(resp.toByteArray(), conn.getResponseCode());
+        Map<String, List<String>> map = Collections.unmodifiableMap(conn.getHeaderFields());
+        return new Response(resp.toByteArray(), map, conn.getResponseCode());
     }
 }
diff --git a/src/main/java/org/peergos/net/HttpProxyHandler.java b/src/main/java/org/peergos/net/HttpProxyHandler.java
@@ -58,7 +58,7 @@ public void handleCallToAPI(HttpExchange httpExchange) {
                 ProxyResponse response = service.proxyRequest(targetNodeId, request);
                 Headers reponseHeaders = httpExchange.getResponseHeaders();
                 for (Map.Entry<String, String> entry : response.headers.entrySet()) {
-                    reponseHeaders.replace(entry.getKey(), List.of(entry.getValue()));
+                    reponseHeaders.set(entry.getKey(), entry.getValue());
                 }
                 httpExchange.sendResponseHeaders(response.statusCode, response.body.length);
                 httpExchange.getResponseBody().write(response.body);