jwt header fix

Signed-off-by: Stephen Crawford <[email protected]>
stephen-crawford · Dec 15, 2023 · fcccd8e · fcccd8e
1 parent 4234823
commit fcccd8e
Showing 1 changed file with 14 additions and 2 deletions.
diff --git a/src/main/java/org/opensearch/security/auditlog/impl/AuditMessage.java b/src/main/java/org/opensearch/security/auditlog/impl/AuditMessage.java
@@ -28,6 +28,7 @@
 import org.apache.commons.codec.digest.DigestUtils;
 import org.apache.hc.core5.net.URIBuilder;
 
+import org.apache.http.HttpHeaders;
 import org.opensearch.ExceptionsHelper;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.collect.Tuple;
@@ -127,10 +128,12 @@ public final class AuditMessage {
     private static final DateTimeFormatter DEFAULT_FORMAT = DateTimeFormat.forPattern("yyyy-MM-dd'T'HH:mm:ss.SSSZZ");
     private final Map<String, Object> auditInfo = new HashMap<String, Object>(50);
     private final AuditCategory msgCategory;
+    private final String customHeader;
 
     public AuditMessage(final AuditCategory msgCategory, final ClusterService clusterService, final Origin origin, final Origin layer) {
         this.msgCategory = Objects.requireNonNull(msgCategory);
         final String currentTime = currentTime();
+        this.customHeader = clusterService.getSettings().get("jwt_header", HttpHeaders.AUTHORIZATION);
         auditInfo.put(FORMAT_VERSION, 4);
         auditInfo.put(CATEGORY, Objects.requireNonNull(msgCategory));
         auditInfo.put(UTC_TIMESTAMP, currentTime);
@@ -360,7 +363,11 @@ public void addRestHeaders(Map<String, List<String>> headers, boolean excludeSen
         if (headers != null && !headers.isEmpty()) {
             final Map<String, List<String>> headersClone = new HashMap<>(headers);
             if (excludeSensitiveHeaders) {
-                headersClone.keySet().removeIf(AUTHORIZATION_HEADER);
+                if (headersClone.containsKey(AUTHORIZATION_HEADER)) { //Look for default "Authorization header
+                    headersClone.keySet().removeIf(AUTHORIZATION_HEADER);
+                } else { // This means it was replaced by a custom header
+                    headersClone.keySet().remove(this.customHeader);
+                }
             }
             auditInfo.put(REST_REQUEST_HEADERS, headersClone);
         }
@@ -417,7 +424,12 @@ public void addTransportHeaders(Map<String, String> headers, boolean excludeSens
         if (headers != null && !headers.isEmpty()) {
             final Map<String, String> headersClone = new HashMap<>(headers);
             if (excludeSensitiveHeaders) {
-                headersClone.keySet().removeIf(AUTHORIZATION_HEADER);
+                if (headersClone.containsKey(AUTHORIZATION_HEADER)) { //Look for default "Authorization header
+                    headersClone.keySet().removeIf(AUTHORIZATION_HEADER);
+                } else { // This means it was replaced by a custom header
+                    headersClone.keySet().remove(customHeader);
+
+                }
             }
             auditInfo.put(TRANSPORT_REQUEST_HEADERS, headersClone);
         }