fix: add possibility to set header values for StreamResource (#11860)

fixes #5471

(cherry picked from commit 508c0c4)

flow-server/src/main/java/com/vaadin/flow/server/StreamResource.java

@@ -20,6 +20,9 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Optional;
 
 import com.vaadin.flow.dom.Element;
@@ -50,6 +53,8 @@ public class StreamResource extends AbstractStreamResource {
 
     private ContentTypeResolver resolver = DEFAULT_RESOLVER;
 
+    private Map<String, String> headers;
+
     private static class DefaultResolver implements ContentTypeResolver {
 
         @Override
@@ -217,6 +222,54 @@ public ContentTypeResolver getContentTypeResolver() {
         return resolver;
     }
 
+    /**
+     * Sets the value of a generic response header. If the header had already
+     * been set, the new value overwrites the previous one.
+     * 
+     * @param name
+     *            a header name
+     * @param value
+     *            value of the header
+     * @return this resource
+     */
+    public StreamResource setHeader(String name, String value) {
+        if (headers == null) {
+            headers = new HashMap<>();
+        }
+        headers.put(name, value);
+        return this;
+    }
+
+    /**
+     * Gets the value for header {@code name} set for the resource.
+     * 
+     * @param name
+     *            name of header to get value for
+     * @return an optional with header value, or an empty optional if it has not
+     *         been set
+     */
+    public Optional<String> getHeader(String name) {
+        if (headers != null) {
+            return Optional.ofNullable(headers.get(name));
+        }
+        return Optional.empty();
+    }
+
+    /**
+     * Gets the additionally configured headers for the resource.
+     * <p>
+     * This method doesn't return headers which are set via explicit setters
+     * like {@link #setContentType(String)} and {@link #setCacheTime(long)}.
+     * 
+     * @return a map of headers and their values
+     */
+    public Map<String, String> getHeaders() {
+        if (headers == null) {
+            return Collections.emptyMap();
+        }
+        return Collections.unmodifiableMap(headers);
+    }
+
     @Override
     public String getName() {
         return fileName;

flow-server/src/main/java/com/vaadin/flow/server/communication/StreamResourceHandler.java

@@ -67,6 +67,8 @@ public void handleRequest(VaadinSession session, VaadinRequest request,
             response.setContentType(streamResource.getContentTypeResolver()
                     .apply(streamResource, context));
             response.setCacheTime(streamResource.getCacheTime());
+            streamResource.getHeaders()
+                    .forEach((name, value) -> response.setHeader(name, value));
             writer = streamResource.getWriter();
             if (writer == null) {
                 throw new IOException(

flow-server/src/test/java/com/vaadin/flow/server/StreamResourceTest.java

@@ -15,17 +15,17 @@
  */
 package com.vaadin.flow.server;
 
+import javax.servlet.ServletContext;
+
 import java.io.ByteArrayInputStream;
 import java.io.InputStream;
-
-import javax.servlet.ServletContext;
+import java.util.Map;
 
 import org.junit.Assert;
 import org.junit.Test;
 import org.mockito.Mockito;
 
 import com.vaadin.flow.function.ContentTypeResolver;
-import com.vaadin.flow.server.StreamResource;
 
 public class StreamResourceTest {
 
@@ -62,6 +62,20 @@ public void setContentType() {
         assertContentType(resource, resource.getContentTypeResolver());
     }
 
+    @Test
+    public void setHeader_headerIsInHeadersListAndGetterReturnsTheValue() {
+        StreamResource resource = new StreamResource("foo",
+                () -> makeEmptyStream());
+
+        resource.setHeader("foo", "bar");
+
+        Assert.assertEquals("bar", resource.getHeader("foo").get());
+
+        Map<String, String> headers = resource.getHeaders();
+        Assert.assertEquals(1, headers.size());
+        Assert.assertEquals("bar", headers.get("foo"));
+    }
+
     private void assertContentType(StreamResource resource,
             ContentTypeResolver resolver) {
         ServletContext context = Mockito.mock(ServletContext.class);

flow-server/src/test/java/com/vaadin/flow/server/communication/StreamResourceHandlerTest.java

@@ -107,4 +107,17 @@ public StreamResourceWriter getWriter() {
         Mockito.verify(response)
                 .setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
     }
+
+    @Test
+    public void inputStreamResourceHasHeader_headerIsWritten()
+            throws IOException {
+        StreamResource res = new StreamResource("readme.md",
+                () -> new ByteArrayInputStream(new byte[0]));
+
+        res.setHeader("foo", "bar");
+
+        handler.handleRequest(session, request, response, res);
+
+        Mockito.verify(response).setHeader("foo", "bar");
+    }
 }