Set correct Content-Type in responses from filter

5 files changed, 47 insertions, 43 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/filter/AuthorizationFilter.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/filter/AuthorizationFilter.java
index 64af3d5c0ca..41f6d0ef4bc 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/filter/AuthorizationFilter.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/filter/AuthorizationFilter.java
@@ -3,8 +3,6 @@ package com.yahoo.vespa.hosted.provision.restapi.v2.filter;
 
 import com.google.inject.Inject;
 import com.yahoo.config.provision.Zone;
-import com.yahoo.jdisc.handler.FastContentWriter;
-import com.yahoo.jdisc.handler.ResponseDispatch;
 import com.yahoo.jdisc.handler.ResponseHandler;
 import com.yahoo.jdisc.http.filter.DiscFilterRequest;
 import com.yahoo.jdisc.http.filter.SecurityRequestFilter;
@@ -17,9 +15,6 @@ import org.bouncycastle.asn1.x500.style.BCStyle;
 import org.bouncycastle.asn1.x500.style.IETFUtils;
 import org.bouncycastle.cert.jcajce.JcaX509CertificateHolder;
 
-import java.io.ByteArrayOutputStream;
-import java.io.IOException;
-import java.io.UncheckedIOException;
 import java.net.URI;
 import java.security.Principal;
 import java.security.cert.CertificateEncodingException;
@@ -70,20 +65,6 @@ public class AuthorizationFilter implements SecurityRequestFilter {
         }
     }
 
-    /** Write error response */
-    static void write(ErrorResponse response, ResponseHandler handler) {
-        try (FastContentWriter writer = ResponseDispatch.newInstance(response.getJdiscResponse())
-                                                        .connectFastWriter(handler)) {
-            ByteArrayOutputStream out = new ByteArrayOutputStream();
-            try {
-                response.render(out);
-            } catch (IOException e) {
-                throw new UncheckedIOException(e);
-            }
-            writer.write(out.toByteArray());
-        }
-    }
-
     /** Log error response without writing anything */
     private static void log(ErrorResponse response, @SuppressWarnings("unused") ResponseHandler handler) {
         log.warning("Would reject request: " + response.getStatus() + " - " + response.message());
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/filter/FilterUtils.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/filter/FilterUtils.java
new file mode 100644
index 00000000000..82d82f4694e
--- /dev/null
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/filter/FilterUtils.java
@@ -0,0 +1,35 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.provision.restapi.v2.filter;
+
+import com.yahoo.container.jdisc.HttpResponse;
+import com.yahoo.jdisc.handler.FastContentWriter;
+import com.yahoo.jdisc.handler.ResponseDispatch;
+import com.yahoo.jdisc.handler.ResponseHandler;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+
+/**
+ * @author mpolden
+ */
+public class FilterUtils {
+
+    private FilterUtils() {}
+
+    /** Write HTTP response using given handler */
+    public static void write(HttpResponse response, ResponseHandler handler) {
+        response.headers().put("Content-Type", response.getContentType());
+        try (FastContentWriter writer = ResponseDispatch.newInstance(response.getJdiscResponse())
+                                                        .connectFastWriter(handler)) {
+            ByteArrayOutputStream out = new ByteArrayOutputStream();
+            try {
+                response.render(out);
+            } catch (IOException e) {
+                throw new UncheckedIOException(e);
+            }
+            writer.write(out.toByteArray());
+        }
+    }
+
+}
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/filter/LocalhostFilter.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/filter/LocalhostFilter.java
index 0bd398ce1f2..8276299a225 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/filter/LocalhostFilter.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/filter/LocalhostFilter.java
@@ -1,17 +1,11 @@
 // Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.provision.restapi.v2.filter;
 
-import com.yahoo.jdisc.handler.FastContentWriter;
-import com.yahoo.jdisc.handler.ResponseDispatch;
 import com.yahoo.jdisc.handler.ResponseHandler;
 import com.yahoo.jdisc.http.filter.DiscFilterRequest;
 import com.yahoo.jdisc.http.filter.SecurityRequestFilter;
 import com.yahoo.vespa.hosted.provision.restapi.v2.ErrorResponse;
 
-import java.io.ByteArrayOutputStream;
-import java.io.IOException;
-import java.io.UncheckedIOException;
-
 /**
  * A security filter that rejects requests not originating from localhost.
  *
@@ -30,22 +24,10 @@ public class LocalhostFilter implements SecurityRequestFilter {
             case inet6Loopback:
                 break; // Allowed
             default:
-                write(ErrorResponse.unauthorized(String.format("%s %s denied for %s: Unauthorized host",
-                                                               request.getMethod(), request.getUri().getPath(),
-                                                               request.getRemoteAddr())), handler);
+                FilterUtils.write(ErrorResponse.unauthorized(
+                        String.format("%s %s denied for %s: Unauthorized host", request.getMethod(),
+                                      request.getUri().getPath(), request.getRemoteAddr())), handler);
         }
     }
 
-    private static void write(ErrorResponse response, ResponseHandler handler) {
-        try (FastContentWriter writer = ResponseDispatch.newInstance(response.getJdiscResponse())
-                                                        .connectFastWriter(handler)) {
-            ByteArrayOutputStream out = new ByteArrayOutputStream();
-            try {
-                response.render(out);
-            } catch (IOException e) {
-                throw new UncheckedIOException(e);
-            }
-            writer.write(out.toByteArray());
-        }
-    }
 }
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/v2/filter/AuthorizationFilterTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/v2/filter/AuthorizationFilterTest.java
index b14cc570a75..b245f2b29d0 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/v2/filter/AuthorizationFilterTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/v2/filter/AuthorizationFilterTest.java
@@ -23,7 +23,7 @@ public class AuthorizationFilterTest {
         tester = new FilterTester(new AuthorizationFilter(new Authorizer(SystemName.main,
                                                                          new MockNodeRepository(new MockCurator(),
                                                                                                 new MockNodeFlavors())),
-                                                          AuthorizationFilter::write));
+                                                          FilterUtils::write));
     }
 
     @Test
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/v2/filter/FilterTester.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/v2/filter/FilterTester.java
index 3916be9d826..d3e98ae2c84 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/v2/filter/FilterTester.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/v2/filter/FilterTester.java
@@ -28,6 +28,8 @@ import java.time.Duration;
 import java.time.Instant;
 import java.util.Collections;
 import java.util.Date;
+import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 
 import static org.junit.Assert.assertEquals;
@@ -55,6 +57,8 @@ public class FilterTester {
         Optional<Response> response = getResponse(request);
         assertTrue("Expected response from filter", response.isPresent());
         assertEquals("Response body", body, response.get().body);
+        assertEquals("Content type", "application/json",
+                     response.get().headers.get("Content-Type").get(0));
         assertEquals("Status code", status, response.get().status);
     }
 
@@ -62,7 +66,7 @@ public class FilterTester {
         RequestHandlerTestDriver.MockResponseHandler handler = new RequestHandlerTestDriver.MockResponseHandler();
         filter.filter(toDiscFilterRequest(request), handler);
         return Optional.ofNullable(handler.getResponse())
-                       .map(response -> new Response(response.getStatus(), handler.readAll()));
+                       .map(response -> new Response(response.getStatus(), response.headers(), handler.readAll()));
     }
 
     private static DiscFilterRequest toDiscFilterRequest(Request request) {
@@ -113,10 +117,12 @@ public class FilterTester {
     private static class Response {
 
         private final int status;
+        private final Map<String, List<String>> headers;
         private final String body;
 
-        private Response(int status, String body) {
+        private Response(int status, Map<String, List<String>> headers, String body) {
             this.status = status;
+            this.headers = headers;
             this.body = body;
         }