Get list of headersToCopy from core and use getHeader(String headerName) instead of getHeaders()

CHANGELOG.md

@@ -16,6 +16,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Resource Sharing] Allow multiple sharable resource types in single resource index ([#5713](https://github.com/opensearch-project/security/pull/5713))
 - Adding Alerting V2 roles to roles.yml ([#5747](https://github.com/opensearch-project/security/pull/5747))
 - add suggest api to ad read access role ([#5754](https://github.com/opensearch-project/security/pull/5754))
+- Get list of headersToCopy from core and use getHeader(String headerName) instead of getHeaders() ([#5769](https://github.com/opensearch-project/security/pull/5769))
 
 ### Bug Fixes
 - Create a WildcardMatcher.NONE when creating a WildcardMatcher with an empty string ([#5694](https://github.com/opensearch-project/security/pull/5694))

src/main/java/org/opensearch/security/OpenSearchSecurityPlugin.java

@@ -123,6 +123,7 @@
 import org.opensearch.repositories.RepositoriesService;
 import org.opensearch.rest.RestController;
 import org.opensearch.rest.RestHandler;
+import org.opensearch.rest.RestHeaderDefinition;
 import org.opensearch.script.ScriptService;
 import org.opensearch.search.internal.InternalScrollSearchRequest;
 import org.opensearch.search.internal.ReaderContext;
@@ -708,13 +709,13 @@ public List<RestHandler> getRestHandlers(
     }
 
     @Override
-    public UnaryOperator<RestHandler> getRestHandlerWrapper(final ThreadContext threadContext) {
+    public UnaryOperator<RestHandler> getRestHandlerWrapper(final ThreadContext threadContext, Set<RestHeaderDefinition> headersToCopy) {
 
         if (client || disabled || SSLConfig.isSslOnlyMode()) {
             return (rh) -> rh;
         }
 
-        return (rh) -> securityRestHandler.wrap(rh, adminDns);
+        return (rh) -> securityRestHandler.wrap(rh, adminDns, headersToCopy);
     }
 
     @Override

src/main/java/org/opensearch/security/filter/SecurityRestFilter.java

@@ -29,6 +29,7 @@
 import java.io.IOException;
 import java.nio.file.Path;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -50,6 +51,7 @@
 import org.opensearch.rest.NamedRoute;
 import org.opensearch.rest.RestChannel;
 import org.opensearch.rest.RestHandler;
+import org.opensearch.rest.RestHeaderDefinition;
 import org.opensearch.rest.RestRequest;
 import org.opensearch.rest.RestRequestFilter;
 import org.opensearch.security.auditlog.AuditLog;
@@ -124,10 +126,12 @@ public SecurityRestFilter(
 
     class AuthczRestHandler extends DelegatingRestHandler {
         private final AdminDNs adminDNs;
+        private final Set<RestHeaderDefinition> headersToCopy;
 
-        public AuthczRestHandler(RestHandler original, AdminDNs adminDNs) {
+        public AuthczRestHandler(RestHandler original, AdminDNs adminDNs, Set<RestHeaderDefinition> headersToCopy) {
             super(original);
             this.adminDNs = adminDNs;
+            this.headersToCopy = headersToCopy;
         }
 
         @Override
@@ -144,13 +148,23 @@ public void handleRequest(RestRequest request, RestChannel channel, NodeClient c
             }
 
             NettyAttribute.popFrom(request, Netty4HttpRequestHeaderVerifier.CONTEXT_TO_RESTORE).ifPresent(storedContext -> {
-                // X_OPAQUE_ID will be overritten on restore - save to apply after restoring the saved context
-                final Map<String, String> tmpHeaders = threadContext.getHeaders();
-                storedContext.restore();
-                for (Map.Entry<String, String> header : tmpHeaders.entrySet()) {
-                    threadContext.putHeader(header.getKey(), header.getValue());
+                // X_OPAQUE_ID will be overwritten on restore - save to apply after restoring the saved context
+                Map<String, String> tmpHeaders = null;
+                for (RestHeaderDefinition header : headersToCopy) {
+                    final String value = threadContext.getHeader(header.getName());
+                    if (value != null) {
+                        if (tmpHeaders == null) {
+                            tmpHeaders = new HashMap<>();
+                        }
+                        tmpHeaders.put(header.getName(), value);
+                    }
                 }
-                if (!tmpHeaders.isEmpty()) {
+                storedContext.restore();
+
+                if (tmpHeaders != null) {
+                    for (Map.Entry<String, String> header : tmpHeaders.entrySet()) {
+                        threadContext.putHeader(header.getKey(), header.getValue());
+                    }
                     threadContext.putHeader(OPENSEARCH_SECURITY_REQUEST_HEADERS, String.join(",", tmpHeaders.keySet()));
                 }
             });
@@ -249,8 +263,8 @@ RestRequest maybeFilterRestRequest(RestRequest request) throws IOException {
      * See {@link AllowlistApiAction} for the implementation of this API.
      * SuperAdmin is identified by credentials, which can be passed in the curl request.
      */
-    public RestHandler wrap(RestHandler original, AdminDNs adminDNs) {
-        return new AuthczRestHandler(original, adminDNs);
+    public RestHandler wrap(RestHandler original, AdminDNs adminDNs, Set<RestHeaderDefinition> headersToCopy) {
+        return new AuthczRestHandler(original, adminDNs, headersToCopy);
     }
 
     /**

src/test/java/org/opensearch/security/filter/SecurityRestFilterUnitTests.java

@@ -12,6 +12,7 @@
 package org.opensearch.security.filter;
 
 import java.nio.file.Path;
+import java.util.HashSet;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -80,7 +81,7 @@ public void setUp() throws NoSuchMethodException {
     public void testSecurityRestFilterWrap() throws Exception {
         AdminDNs adminDNs = mock(AdminDNs.class);
 
-        RestHandler wrappedRestHandler = sf.wrap(testRestHandler, adminDNs);
+        RestHandler wrappedRestHandler = sf.wrap(testRestHandler, adminDNs, new HashSet<>());
 
         assertTrue(wrappedRestHandler instanceof SecurityRestFilter.AuthczRestHandler);
         assertFalse(wrappedRestHandler instanceof TestRestHandler);
@@ -92,7 +93,7 @@ public void testDoesCallDelegateOnSuccessfulAuthorization() throws Exception {
         AdminDNs adminDNs = mock(AdminDNs.class);
 
         RestHandler testRestHandlerSpy = spy(testRestHandler);
-        RestHandler wrappedRestHandler = filterSpy.wrap(testRestHandlerSpy, adminDNs);
+        RestHandler wrappedRestHandler = filterSpy.wrap(testRestHandlerSpy, adminDNs, new HashSet<>());
 
         doReturn(false).when(filterSpy).userIsSuperAdmin(any(), any());