Support all sub-classes of 'Response' in security response filters

Previously security response filters were just skipped if a request
handler or (security) request filter returned a response that was not a
sub-class of HttpResponse.

4 files changed, 74 insertions, 36 deletions

diff --git a/container-core/src/main/java/com/yahoo/jdisc/http/CookieHelper.java b/container-core/src/main/java/com/yahoo/jdisc/http/CookieHelper.java
new file mode 100644
index 00000000000..897c18d1129
--- /dev/null
+++ b/container-core/src/main/java/com/yahoo/jdisc/http/CookieHelper.java
@@ -0,0 +1,38 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.jdisc.http;
+
+import com.yahoo.jdisc.HeaderFields;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+
+/**
+ * Helper for encoding/decoding cookies on request/response.
+ *
+ * @author bjorncs
+ */
+public class CookieHelper {
+
+    private CookieHelper() {}
+
+    public static List<Cookie> decodeSetCookieHeader(HeaderFields headers) {
+        List<String> cookies = headers.get(HttpHeaders.Names.SET_COOKIE);
+        if (cookies == null) {
+            return Collections.emptyList();
+        }
+        List<Cookie> ret = new LinkedList<>();
+        for (String cookie : cookies) {
+            ret.add(Cookie.fromSetCookieHeader(cookie));
+        }
+        return ret;
+    }
+
+    public static void encodeSetCookieHeader(HeaderFields headers, List<Cookie> cookies) {
+        headers.remove(HttpHeaders.Names.SET_COOKIE);
+        for (Cookie cookie : cookies) {
+            headers.add(HttpHeaders.Names.SET_COOKIE, Cookie.toSetCookieHeaders(Arrays.asList(cookie)));
+        }
+    }
+}
diff --git a/container-core/src/main/java/com/yahoo/jdisc/http/HttpResponse.java b/container-core/src/main/java/com/yahoo/jdisc/http/HttpResponse.java
index f7138ba0e2b..2e2553e421a 100644
--- a/container-core/src/main/java/com/yahoo/jdisc/http/HttpResponse.java
+++ b/container-core/src/main/java/com/yahoo/jdisc/http/HttpResponse.java
@@ -9,9 +9,6 @@ import com.yahoo.jdisc.handler.ContentChannel;
 import com.yahoo.jdisc.handler.ResponseHandler;
 import com.yahoo.jdisc.http.servlet.ServletOrJdiscHttpResponse;
 
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.LinkedList;
 import java.util.List;
 
 /**
@@ -62,23 +59,13 @@ public class HttpResponse extends Response implements ServletOrJdiscHttpResponse
         target.addAll(headers());
     }
 
+    @Override
     public List<Cookie> decodeSetCookieHeader() {
-        List<String> cookies = headers().get(HttpHeaders.Names.SET_COOKIE);
-        if (cookies == null) {
-            return Collections.emptyList();
-        }
-        List<Cookie> ret = new LinkedList<>();
-        for (String cookie : cookies) {
-            ret.add(Cookie.fromSetCookieHeader(cookie));
-        }
-        return ret;
+        return CookieHelper.decodeSetCookieHeader(headers());
     }
 
     public void encodeSetCookieHeader(List<Cookie> cookies) {
-        headers().remove(HttpHeaders.Names.SET_COOKIE);
-        for (Cookie cookie : cookies) {
-            headers().add(HttpHeaders.Names.SET_COOKIE, Cookie.toSetCookieHeaders(Arrays.asList(cookie)));
-        }
+        CookieHelper.encodeSetCookieHeader(headers(), cookies);
     }
 
     /**
diff --git a/container-core/src/main/java/com/yahoo/jdisc/http/filter/JdiscFilterResponse.java b/container-core/src/main/java/com/yahoo/jdisc/http/filter/JdiscFilterResponse.java
index ff81359f93c..32a0c7e2ed6 100644
--- a/container-core/src/main/java/com/yahoo/jdisc/http/filter/JdiscFilterResponse.java
+++ b/container-core/src/main/java/com/yahoo/jdisc/http/filter/JdiscFilterResponse.java
@@ -1,22 +1,26 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.jdisc.http.filter;
 
+import com.yahoo.jdisc.HeaderFields;
+import com.yahoo.jdisc.Response;
 import com.yahoo.jdisc.http.Cookie;
+import com.yahoo.jdisc.http.CookieHelper;
 import com.yahoo.jdisc.http.HttpResponse;
+import com.yahoo.jdisc.http.servlet.ServletOrJdiscHttpResponse;
 
 import java.util.List;
+import java.util.Map;
 
 /**
  * JDisc implementation of a filter request.
- *
- * @since 5.27
  */
-public class JdiscFilterResponse extends DiscFilterResponse {
+class JdiscFilterResponse extends DiscFilterResponse {
 
-    private final HttpResponse parent;
+    private final Response parent;
 
-    public JdiscFilterResponse(HttpResponse parent) {
-        super(parent);
+    JdiscFilterResponse(Response parent) {
+        // A separate adapter is required as super-constructor will invoke methods from ServletOrJdiscHttpResponse in its constructor
+        super(parent instanceof HttpResponse ? (HttpResponse)parent : new Adapter(parent));
         this.parent = parent;
     }
 
@@ -61,7 +65,20 @@ public class JdiscFilterResponse extends DiscFilterResponse {
 
     @Override
     public void setCookies(List<Cookie> cookies) {
-        parent.encodeSetCookieHeader(cookies);
+        CookieHelper.encodeSetCookieHeader(parent.headers(), cookies);
+    }
+
+    private static class Adapter implements ServletOrJdiscHttpResponse {
+        private final Response response;
+
+        Adapter(Response response) {
+            this.response = response;
+        }
+
+        @Override public void copyHeaders(HeaderFields target) { target.addAll(response.headers()); }
+        @Override public int getStatus() { return response.getStatus(); }
+        @Override public Map<String, Object> context() { return response.context(); }
+        @Override public List<Cookie> decodeSetCookieHeader() { return CookieHelper.decodeSetCookieHeader(response.headers()); }
     }
 
 }
diff --git a/container-core/src/main/java/com/yahoo/jdisc/http/filter/SecurityResponseFilterChain.java b/container-core/src/main/java/com/yahoo/jdisc/http/filter/SecurityResponseFilterChain.java
index d45b406a375..ad0fb75ebff 100644
--- a/container-core/src/main/java/com/yahoo/jdisc/http/filter/SecurityResponseFilterChain.java
+++ b/container-core/src/main/java/com/yahoo/jdisc/http/filter/SecurityResponseFilterChain.java
@@ -1,6 +1,11 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.jdisc.http.filter;
 
+import com.yahoo.jdisc.AbstractResource;
+import com.yahoo.jdisc.Request;
+import com.yahoo.jdisc.Response;
+import com.yahoo.jdisc.http.HttpRequest;
+
 import java.net.URI;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -8,16 +13,10 @@ import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
 
-import com.yahoo.jdisc.AbstractResource;
-import com.yahoo.jdisc.Request;
-import com.yahoo.jdisc.Response;
-import com.yahoo.jdisc.http.HttpRequest;
-import com.yahoo.jdisc.http.HttpResponse;
-
 /**
  * Implementation of TypedFilterChain for DiscFilterResponse
- * @author tejalk
  *
+ * @author Tejal Knot
  */
 public class SecurityResponseFilterChain extends AbstractResource implements ResponseFilter {
 
@@ -31,12 +30,9 @@ public class SecurityResponseFilterChain extends AbstractResource implements Res
 
     @Override
     public void filter(Response response, Request request) {
-        if(response instanceof HttpResponse) {
-            DiscFilterResponse discFilterResponse = new JdiscFilterResponse((HttpResponse)response);
-            RequestView requestView = new RequestViewImpl(request);
-            filter(requestView, discFilterResponse);
-        }
-
+        DiscFilterResponse discFilterResponse = new JdiscFilterResponse(response);
+        RequestView requestView = new RequestViewImpl(request);
+        filter(requestView, discFilterResponse);
     }
 
     public void filter(RequestView requestView, DiscFilterResponse response) {