shield: also restore original context to transport handlers

See elastic/elasticsearch#1380

Original commit: elastic/x-pack-elasticsearch@afbd964f18

elasticsearch/x-pack/shield/src/main/java/org/elasticsearch/shield/transport/ShieldServerTransportService.java

@@ -78,7 +78,7 @@ public class ShieldServerTransportService extends TransportService {
 
     @Override
     public <T extends TransportResponse> void sendRequest(DiscoveryNode node, String action, TransportRequest request, TransportRequestOptions options, TransportResponseHandler<T> handler) {
-        try (ThreadContext.StoredContext original = threadPool.getThreadContext().newStoredContext()) {
+        final ThreadContext.StoredContext original = threadPool.getThreadContext().newStoredContext();
         // FIXME this is really just a hack. What happens is that we send a request and we always copy headers over
         // Sometimes a system action gets executed like a internal create index request or update mappings request
         // which means that the user is copied over to system actions and these really fail for internal things...
@@ -86,7 +86,7 @@ public class ShieldServerTransportService extends TransportService {
             try (ThreadContext.StoredContext ctx = threadPool.getThreadContext().stashContext()) {
                 try {
                     clientFilter.outbound(action, request);
-                        super.sendRequest(node, action, request, options, handler);
+                    super.sendRequest(node, action, request, options, new ContextRestoreResponseHandler<>(original, handler));
                 } catch (Throwable t) {
                     handler.handleException(new TransportException("failed sending request", t));
                 }
@@ -94,13 +94,12 @@ public class ShieldServerTransportService extends TransportService {
         } else {
             try {
                 clientFilter.outbound(action, request);
-                    super.sendRequest(node, action, request, options, handler);
+                super.sendRequest(node, action, request, options, new ContextRestoreResponseHandler<>(original, handler));
             } catch (Throwable t) {
                 handler.handleException(new TransportException("failed sending request", t));
             }
         }
     }
-    }
 
     @Override
     public <Request extends TransportRequest> void registerRequestHandler(String action, Supplier<Request> requestFactory, String executor, TransportRequestHandler<Request> handler) {
@@ -204,4 +203,38 @@ public class ShieldServerTransportService extends TransportService {
         }
     }
 
+    /**
+     * This handler wrapper ensures that the response thread executes with the correct thread context. Before any of the4 handle methods
+     * are invoked we restore the context.
+     */
+    private final static class ContextRestoreResponseHandler<T extends TransportResponse> implements TransportResponseHandler<T> {
+        private final TransportResponseHandler<T> delegate;
+        private final ThreadContext.StoredContext threadContext;
+        private ContextRestoreResponseHandler(ThreadContext.StoredContext threadContext, TransportResponseHandler<T> delegate) {
+            this.delegate = delegate;
+            this.threadContext = threadContext;
+        }
+
+        @Override
+        public T newInstance() {
+            return delegate.newInstance();
+        }
+
+        @Override
+        public void handleResponse(T response) {
+            threadContext.restore();
+            delegate.handleResponse(response);
+        }
+
+        @Override
+        public void handleException(TransportException exp) {
+            threadContext.restore();
+            delegate.handleException(exp);
+        }
+
+        @Override
+        public String executor() {
+            return delegate.executor();
+        }
+    }
 }