From a5389cec65e2b0e0b53a8ac930d9832cb7a347fc Mon Sep 17 00:00:00 2001
From: Martin Ndegwa <mndegwa@ona.io>
Date: Wed, 2 Oct 2024 09:13:34 +0300
Subject: [PATCH] Add support for removing query params in request mutations

---
 .../BearerAuthorizationInterceptor.java       |  6 ++-
 .../gateway/interfaces/RequestMutation.java   |  5 +++
 .../BearerAuthorizationInterceptorTest.java   | 39 +++++++++++++++++++
 3 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/server/src/main/java/com/google/fhir/gateway/BearerAuthorizationInterceptor.java b/server/src/main/java/com/google/fhir/gateway/BearerAuthorizationInterceptor.java
index 2fe9a467..ae3b692e 100755
--- a/server/src/main/java/com/google/fhir/gateway/BearerAuthorizationInterceptor.java
+++ b/server/src/main/java/com/google/fhir/gateway/BearerAuthorizationInterceptor.java
@@ -283,10 +283,14 @@ private void serveWellKnown(ServletRequestDetails request) {
   void mutateRequest(RequestDetails requestDetails, AccessDecision accessDecision) {
     RequestMutation mutation =
         accessDecision.getRequestMutation(new RequestDetailsToReader(requestDetails));
-    if (mutation == null || CollectionUtils.isEmpty(mutation.getQueryParams())) {
+    if (mutation == null
+        || (CollectionUtils.isEmpty(mutation.getQueryParams())
+            && CollectionUtils.isEmpty(mutation.getDiscardQueryParams()))) {
       return;
     }
 
+    mutation.getDiscardQueryParams().forEach(requestDetails::removeParameter);
+
     mutation
         .getQueryParams()
         .forEach((key, value) -> requestDetails.addParameter(key, value.toArray(new String[0])));
diff --git a/server/src/main/java/com/google/fhir/gateway/interfaces/RequestMutation.java b/server/src/main/java/com/google/fhir/gateway/interfaces/RequestMutation.java
index fdcac3dd..c0e62bf1 100644
--- a/server/src/main/java/com/google/fhir/gateway/interfaces/RequestMutation.java
+++ b/server/src/main/java/com/google/fhir/gateway/interfaces/RequestMutation.java
@@ -15,6 +15,7 @@
  */
 package com.google.fhir.gateway.interfaces;
 
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -32,4 +33,8 @@ public class RequestMutation {
   // parameter in RequestMutation is already present in the original request).
   // Old parameter values should be explicitly retained while mutating values for that parameter.
   @Builder.Default Map<String, List<String>> queryParams = new HashMap<>();
+
+  // Query parameters that are no longer needed when forwarding the request to the upstream server
+  // Parameters with the keys in this list will be removed
+  @Builder.Default List<String> discardQueryParams = new ArrayList<>();
 }
diff --git a/server/src/test/java/com/google/fhir/gateway/BearerAuthorizationInterceptorTest.java b/server/src/test/java/com/google/fhir/gateway/BearerAuthorizationInterceptorTest.java
index 87e6b02a..a4441056 100644
--- a/server/src/test/java/com/google/fhir/gateway/BearerAuthorizationInterceptorTest.java
+++ b/server/src/test/java/com/google/fhir/gateway/BearerAuthorizationInterceptorTest.java
@@ -46,6 +46,7 @@
 import java.io.Writer;
 import java.net.URL;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -53,6 +54,7 @@
 import javax.servlet.http.HttpServletResponse;
 import org.apache.http.HttpResponse;
 import org.apache.http.HttpStatus;
+import org.hamcrest.Matchers;
 import org.hl7.fhir.instance.model.api.IBaseResource;
 import org.hl7.fhir.r4.model.CapabilityStatement;
 import org.junit.Before;
@@ -307,6 +309,43 @@ public String postProcess(
         arrayContainingInAnyOrder("param3-value2", "param3-value1"));
   }
 
+  @Test
+  public void mutateRequestRemoveQueryParams() {
+
+    List<String> queryParamsToRemove = new ArrayList<>();
+    queryParamsToRemove.add("param1");
+
+    ServletRequestDetails requestDetails = new ServletRequestDetails();
+    requestDetails.addParameter("param1", new String[] {"param1-value1"});
+    requestDetails.addParameter("param2", new String[] {"param2-value1"});
+
+    AccessDecision mutableAccessDecision =
+        new AccessDecision() {
+          public boolean canAccess() {
+            return true;
+          }
+
+          public RequestMutation getRequestMutation(RequestDetailsReader requestDetailsReader) {
+            RequestMutation requestMutation = RequestMutation.builder().build();
+            requestMutation.getDiscardQueryParams().addAll(queryParamsToRemove);
+            return requestMutation;
+          }
+
+          public String postProcess(
+              RequestDetailsReader requestDetailsReader, HttpResponse response) throws IOException {
+            return null;
+          }
+        };
+    assertThat(requestDetails.getParameters().size(), Matchers.equalTo(2));
+
+    testInstance.mutateRequest(requestDetails, mutableAccessDecision);
+
+    assertThat(requestDetails.getParameters().size(), Matchers.equalTo(1));
+
+    assertThat(
+        requestDetails.getParameters().get("param2"), arrayContainingInAnyOrder("param2-value1"));
+  }
+
   @Test
   public void shouldSendGzippedResponseWhenRequested() throws IOException {
     testInstance = createTestInstance(true, null);