Add configurable response headers for blocked requests

3 files changed, 72 insertions, 7 deletions

diff --git a/jdisc-security-filters/src/main/java/com/yahoo/jdisc/http/filter/security/rule/RuleBasedRequestFilter.java b/jdisc-security-filters/src/main/java/com/yahoo/jdisc/http/filter/security/rule/RuleBasedRequestFilter.java
index 71f1965c764..7bdc386e4b4 100644
--- a/jdisc-security-filters/src/main/java/com/yahoo/jdisc/http/filter/security/rule/RuleBasedRequestFilter.java
+++ b/jdisc-security-filters/src/main/java/com/yahoo/jdisc/http/filter/security/rule/RuleBasedRequestFilter.java
@@ -3,6 +3,7 @@ package com.yahoo.jdisc.http.filter.security.rule;
 
 import com.google.inject.Inject;
 import com.yahoo.jdisc.Metric;
+import com.yahoo.jdisc.Response;
 import com.yahoo.jdisc.http.filter.DiscFilterRequest;
 import com.yahoo.jdisc.http.filter.security.base.JsonSecurityRequestFilterBase;
 import com.yahoo.jdisc.http.filter.security.rule.RuleBasedFilterConfig.Rule.Action;
@@ -56,7 +57,11 @@ public class RuleBasedRequestFilter extends JsonSecurityRequestFilterBase {
     private static ErrorResponse createDefaultResponse(RuleBasedFilterConfig.DefaultRule defaultRule) {
         switch (defaultRule.action()) {
             case ALLOW: return null;
-            case BLOCK: return new ErrorResponse(defaultRule.blockResponseCode(), defaultRule.blockResponseMessage());
+            case BLOCK: {
+                Response response = new Response(defaultRule.blockResponseCode());
+                defaultRule.blockResponseHeaders().forEach(h -> response.headers().add(h.name(), h.value()));
+                return new ErrorResponse(response, defaultRule.blockResponseMessage());
+            }
             default: throw new IllegalArgumentException(defaultRule.action().name());
         }
     }
@@ -100,9 +105,13 @@ public class RuleBasedRequestFilter extends JsonSecurityRequestFilterBase {
                     .map(m -> m.name().toUpperCase())
                     .collect(Collectors.toSet());
             this.pathGlobExpressions = Set.copyOf(config.pathExpressions());
-            this.response = config.action() == Action.Enum.BLOCK
-                    ? new ErrorResponse(config.blockResponseCode(), config.blockResponseMessage())
-                    : null;
+            this.response = config.action() == Action.Enum.BLOCK ? createResponse(config) : null;
+        }
+
+        private static ErrorResponse createResponse(RuleBasedFilterConfig.Rule config) {
+            Response response = new Response(config.blockResponseCode());
+            config.blockResponseHeaders().forEach(h -> response.headers().add(h.name(), h.value()));
+            return new ErrorResponse(response, config.blockResponseMessage());
         }
 
         boolean matches(String method, URI uri) {
diff --git a/jdisc-security-filters/src/main/resources/configdefinitions/jdisc.http.filter.security.rule.rule-based-filter.def b/jdisc-security-filters/src/main/resources/configdefinitions/jdisc.http.filter.security.rule.rule-based-filter.def
index 374c3ca69c6..6abf3d43a7d 100644
--- a/jdisc-security-filters/src/main/resources/configdefinitions/jdisc.http.filter.security.rule.rule-based-filter.def
+++ b/jdisc-security-filters/src/main/resources/configdefinitions/jdisc.http.filter.security.rule.rule-based-filter.def
@@ -5,6 +5,8 @@ dryrun bool default=false
 defaultRule.action enum { ALLOW, BLOCK }
 defaultRule.blockResponseCode int default=403
 defaultRule.blockResponseMessage string default=""
+defaultRule.blockResponseHeaders[].name string
+defaultRule.blockResponseHeaders[].value string
 rule[].name string
 rule[].action enum { ALLOW, BLOCK }
 rule[].hostNames[] string
@@ -12,5 +14,5 @@ rule[].methods[] enum { GET, POST, PUT, PATCH, DELETE }
 rule[].pathExpressions[] string
 rule[].blockResponseCode int default=403
 rule[].blockResponseMessage string default=""
-
-
+rule[].blockResponseHeaders[].name string
+rule[].blockResponseHeaders[].value string
diff --git a/jdisc-security-filters/src/test/java/com/yahoo/jdisc/http/filter/security/rule/RuleBasedRequestFilterTest.java b/jdisc-security-filters/src/test/java/com/yahoo/jdisc/http/filter/security/rule/RuleBasedRequestFilterTest.java
index c67d3b430c8..3bd606d7edd 100644
--- a/jdisc-security-filters/src/test/java/com/yahoo/jdisc/http/filter/security/rule/RuleBasedRequestFilterTest.java
+++ b/jdisc-security-filters/src/test/java/com/yahoo/jdisc/http/filter/security/rule/RuleBasedRequestFilterTest.java
@@ -11,11 +11,11 @@ import com.yahoo.jdisc.http.filter.DiscFilterRequest;
 import com.yahoo.jdisc.http.filter.security.rule.RuleBasedFilterConfig.DefaultRule;
 import com.yahoo.jdisc.http.filter.security.rule.RuleBasedFilterConfig.Rule;
 import com.yahoo.test.json.JsonTestHelper;
-import com.yahoo.vespa.jdk8compat.List;
 import org.junit.jupiter.api.Test;
 
 import java.io.IOException;
 import java.net.URI;
+import java.util.List;
 import java.util.Set;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -148,6 +148,60 @@ class RuleBasedRequestFilterTest {
         assertAllowed(responseHandler, metric);
     }
 
+    @Test
+    void includes_default_rule_response_headers_in_response_for_blocked_request() throws IOException {
+        RuleBasedFilterConfig config = new RuleBasedFilterConfig.Builder()
+                .dryrun(false)
+                .defaultRule(new DefaultRule.Builder()
+                        .action(DefaultRule.Action.Enum.BLOCK)
+                        .blockResponseHeaders(new DefaultRule.BlockResponseHeaders.Builder()
+                                .name("Response-Header-1").value("first-header"))
+                        .blockResponseHeaders(new DefaultRule.BlockResponseHeaders.Builder()
+                                .name("Response-Header-2").value("second-header")))
+                .build();
+
+        Metric metric = mock(Metric.class);
+        RuleBasedRequestFilter filter = new RuleBasedRequestFilter(metric, config);
+        MockResponseHandler responseHandler = new MockResponseHandler();
+        filter.filter(request("GET", "http://myserver:80/"), responseHandler);
+
+        assertBlocked(responseHandler, metric, 403, "");
+        Response response = responseHandler.getResponse();
+        assertResponseHeader(response, "Response-Header-1", "first-header");
+        assertResponseHeader(response, "Response-Header-2", "second-header");
+    }
+
+    @Test
+    void includes_rule_response_headers_in_response_for_blocked_request() throws IOException {
+        RuleBasedFilterConfig config = new RuleBasedFilterConfig.Builder()
+                .dryrun(false)
+                .defaultRule(new DefaultRule.Builder()
+                        .action(DefaultRule.Action.Enum.ALLOW))
+                .rule(new Rule.Builder()
+                        .name("rule")
+                        .pathExpressions("/path-to-resource")
+                        .action(Rule.Action.Enum.BLOCK)
+                        .blockResponseHeaders(new Rule.BlockResponseHeaders.Builder()
+                                .name("Response-Header-1").value("first-header")))
+                .build();
+
+        Metric metric = mock(Metric.class);
+        RuleBasedRequestFilter filter = new RuleBasedRequestFilter(metric, config);
+        MockResponseHandler responseHandler = new MockResponseHandler();
+        filter.filter(request("GET", "http://myserver/path-to-resource"), responseHandler);
+
+        assertBlocked(responseHandler, metric, 403, "");
+        Response response = responseHandler.getResponse();
+        assertResponseHeader(response, "Response-Header-1", "first-header");
+    }
+
+    private void assertResponseHeader(Response response, String name, String expectedValue) {
+        List<String> actualValues = response.headers().get(name);
+        assertNotNull(actualValues);
+        assertEquals(1, actualValues.size());
+        assertEquals(expectedValue, actualValues.get(0));
+    }
+
     private static DiscFilterRequest request(String method, String uri) {
         DiscFilterRequest request = mock(DiscFilterRequest.class);
         when(request.getMethod()).thenReturn(method);