fix: Preserve response headers across context restore in SecurityInterceptor

Signed-off-by: Kishore Kumaar Natarajan <kkumaarn@amazon.com>

Author: KishoreKicha14

src/main/java/org/opensearch/security/transport/SecurityInterceptor.java

@@ -178,8 +178,12 @@ public <T extends TransportResponse> void sendRequestDecorate(
             requestHeadersToCopy.removeAll(Task.REQUEST_HEADERS); // Special case where this header is preserved during stashContext.
         }
 
+        final Supplier<ThreadContext.StoredContext> restorableContextSupplier = getThreadContext().newRestorableContext(true);
         try (ThreadContext.StoredContext stashedContext = getThreadContext().stashContext()) {
-            final TransportResponseHandler<T> restoringHandler = new RestoringTransportResponseHandler<T>(handler, stashedContext);
+            final TransportResponseHandler<T> restoringHandler = new RestoringTransportResponseHandler<T>(
+                handler,
+                restorableContextSupplier
+            );
             getThreadContext().putHeader("_opendistro_security_remotecn", cs.getClusterName().value());
 
             final Map<String, String> headerMap = new HashMap<>(
@@ -377,10 +381,13 @@ private ThreadContext getThreadContext() {
     // which is private scoped
     private class RestoringTransportResponseHandler<T extends TransportResponse> implements TransportResponseHandler<T> {
 
-        private final ThreadContext.StoredContext contextToRestore;
+        private final Supplier<ThreadContext.StoredContext> contextToRestore;
         private final TransportResponseHandler<T> innerHandler;
 
-        private RestoringTransportResponseHandler(TransportResponseHandler<T> innerHandler, ThreadContext.StoredContext contextToRestore) {
+        private RestoringTransportResponseHandler(
+            TransportResponseHandler<T> innerHandler,
+            Supplier<ThreadContext.StoredContext> contextToRestore
+        ) {
             this.contextToRestore = contextToRestore;
             this.innerHandler = innerHandler;
         }
@@ -407,7 +414,7 @@ public void handleResponse(T response) {
             List<String> dlsResponseHeader = responseHeaders.get(ConfigConstants.OPENDISTRO_SECURITY_DLS_QUERY_HEADER);
             List<String> maskedFieldsResponseHeader = responseHeaders.get(ConfigConstants.OPENDISTRO_SECURITY_MASKED_FIELD_HEADER);
 
-            contextToRestore.restore();
+            contextToRestore.get();
 
             final boolean isDebugEnabled = log.isDebugEnabled();
             if (response instanceof ClusterSearchShardsResponse) {
@@ -438,7 +445,7 @@ public void handleResponse(T response) {
 
         @Override
         public void handleException(TransportException e) {
-            contextToRestore.restore();
+            contextToRestore.get();
             innerHandler.handleException(e);
         }