Add extra headers for static resources

Allow users to add extra headers when using the
`HttpResourceBinding` for static resources.

Signed-off-by: Yihong Wang <yh.wang@ibm.com>

http-server/src/main/java/com/facebook/airlift/http/server/ClassPathResourceHandler.java

@@ -16,6 +16,7 @@
 package com.facebook.airlift.http.server;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.io.ByteStreams;
 import com.google.common.net.HttpHeaders;
 import org.eclipse.jetty.http.HttpMethod;
@@ -31,6 +32,7 @@
 import java.io.InputStream;
 import java.net.URL;
 import java.util.List;
+import java.util.Map;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static java.util.Objects.requireNonNull;
@@ -54,13 +56,14 @@ public class ClassPathResourceHandler
     private final String baseUri; // "" or "/foo"
     private final String classPathResourceBase;
     private final List<String> welcomeFiles;
+    private final Map<String, String> extraHeaders;
 
     public ClassPathResourceHandler(String baseUri, String classPathResourceBase, String... welcomeFiles)
     {
-        this(baseUri, classPathResourceBase, ImmutableList.copyOf(welcomeFiles));
+        this(baseUri, classPathResourceBase, ImmutableList.copyOf(welcomeFiles), ImmutableMap.of());
     }
 
-    public ClassPathResourceHandler(String baseUri, String classPathResourceBase, List<String> welcomeFiles)
+    public ClassPathResourceHandler(String baseUri, String classPathResourceBase, List<String> welcomeFiles, Map<String, String> extraHeaders)
     {
         requireNonNull(baseUri, "baseUri is null");
         requireNonNull(classPathResourceBase, "classPathResourceBase is null");
@@ -81,6 +84,7 @@ public ClassPathResourceHandler(String baseUri, String classPathResourceBase, Li
             files.add(welcomeFile);
         }
         this.welcomeFiles = files.build();
+        this.extraHeaders = ImmutableMap.copyOf(extraHeaders);
     }
 
     @Override
@@ -129,6 +133,7 @@ public void handle(String target, Request baseRequest, HttpServletRequest reques
 
             String contentType = MIME_TYPES.getMimeByExtension(resource.toString());
             response.setContentType(contentType);
+            extraHeaders.forEach((name, value) -> response.setHeader(name, value));
 
             if (skipContent) {
                 return;

http-server/src/main/java/com/facebook/airlift/http/server/HttpServer.java

@@ -366,7 +366,11 @@ public HttpServer(HttpServerInfo httpServerInfo,
 
         for (HttpResourceBinding resource : resources) {
             GzipHandler gzipHandler = new GzipHandler();
-            gzipHandler.setHandler(new ClassPathResourceHandler(resource.getBaseUri(), resource.getClassPathResourceBase(), resource.getWelcomeFiles()));
+            gzipHandler.setHandler(new ClassPathResourceHandler(
+                    resource.getBaseUri(),
+                    resource.getClassPathResourceBase(),
+                    resource.getWelcomeFiles(),
+                    resource.getExtraHeaders()));
             handlers.addHandler(gzipHandler);
         }
 

http-server/src/main/java/com/facebook/airlift/http/server/HttpServerBinder.java

@@ -1,11 +1,14 @@
 package com.facebook.airlift.http.server;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.inject.Binder;
 
 import java.lang.annotation.Annotation;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 import static com.google.inject.multibindings.Multibinder.newSetBinder;
 import static java.util.Objects.requireNonNull;
@@ -43,6 +46,7 @@ public static class HttpResourceBinding
         private final String baseUri;
         private final String classPathResourceBase;
         private final List<String> welcomeFiles = new ArrayList<>();
+        private final Map<String, String> extraHeaders = new HashMap<>();
 
         public HttpResourceBinding(String baseUri, String classPathResourceBase)
         {
@@ -70,5 +74,16 @@ public HttpResourceBinding withWelcomeFile(String welcomeFile)
             welcomeFiles.add(welcomeFile);
             return this;
         }
+
+        public HttpResourceBinding extraHeaders(String name, String value)
+        {
+            extraHeaders.put(name, value);
+            return this;
+        }
+
+        public Map<String, String> getExtraHeaders()
+        {
+            return ImmutableMap.copyOf(extraHeaders);
+        }
     }
 }

http-server/src/test/java/com/facebook/airlift/http/server/TestHttpServerModule.java

@@ -181,7 +181,8 @@ public void testServer()
                     binder.bind(Servlet.class).annotatedWith(TheServlet.class).to(DummyServlet.class);
                     newSetBinder(binder, Filter.class, TheServlet.class).addBinding().to(DummyFilter.class).in(Scopes.SINGLETON);
                     httpServerBinder(binder).bindResource("/", "webapp/user").withWelcomeFile("user-welcome.txt");
-                    httpServerBinder(binder).bindResource("/", "webapp/user2");
+                    httpServerBinder(binder).bindResource("/", "webapp/user2")
+                            .extraHeaders("Content-Security-Policy", "frame-ancestors 'self'");
                     httpServerBinder(binder).bindResource("path", "webapp/user").withWelcomeFile("user-welcome.txt");
                     httpServerBinder(binder).bindResource("path", "webapp/user2");
                 });
@@ -202,12 +203,14 @@ public void testServer()
             response = client.execute(prepareGet().setUri(httpUri.resolve("/filter")).build(), createStatusResponseHandler());
 
             assertEquals(response.getStatusCode(), HttpServletResponse.SC_PAYMENT_REQUIRED);
-
+            Map<String, String> extraHeaders = new ImmutableMap.Builder<String, String>()
+                    .put("Content-Security-Policy", "frame-ancestors 'self'")
+                    .build();
             // test http resources
             assertResource(httpUri, client, "", "welcome user!");
             assertResource(httpUri, client, "user-welcome.txt", "welcome user!");
             assertResource(httpUri, client, "user.txt", "user");
-            assertResource(httpUri, client, "user2.txt", "user2");
+            assertResource(httpUri, client, "user2.txt", "user2", extraHeaders);
             assertRedirect(httpUri, client, "path", "/path/");
             assertResource(httpUri, client, "path/", "welcome user!");
             assertResource(httpUri, client, "path/user-welcome.txt", "welcome user!");
@@ -220,6 +223,11 @@ public void testServer()
     }
 
     private void assertResource(URI baseUri, HttpClient client, String path, String contents)
+    {
+        assertResource(baseUri, client, path, contents, ImmutableMap.of());
+    }
+
+    private void assertResource(URI baseUri, HttpClient client, String path, String contents, Map<String, String> hasHeaders)
     {
         HttpUriBuilder uriBuilder = uriBuilderFrom(baseUri);
         StringResponse response = client.execute(prepareGet().setUri(uriBuilder.appendPath(path).build()).build(), createStringResponseHandler());
@@ -229,6 +237,10 @@ private void assertResource(URI baseUri, HttpClient client, String path, String
         MediaType mediaType = MediaType.parse(contentType);
         assertTrue(PLAIN_TEXT_UTF_8.is(mediaType), "Expected text/plain but got " + mediaType);
         assertEquals(response.getBody().trim(), contents);
+        hasHeaders.forEach((name, value) -> {
+            assertNotNull(response.getHeader(name), name + " header is absent");
+            assertEquals(value, response.getHeader(name));
+        });
     }
 
     private void assertRedirect(URI baseUri, HttpClient client, String path, String redirect)