From 1bb2a1502d705ea966b2a0d3d05025907f0ff28b Mon Sep 17 00:00:00 2001
From: Tim Vernum <tim@adjective.org>
Date: Fri, 5 Oct 2018 12:08:21 +1000
Subject: [PATCH] Preserve thread context during authentication. (#34290)

There may be values in the thread context that ought to be preseved
for later use, even if one or more realms perform asynchronous
authentication.

This commit changes the AuthenticationService to wrap the potentially
asynchronous calls in a ContextPreservingActionListener that retains
the original thread context for the authentication.
---
 .../security/authc/AuthenticationService.java |  7 ++++---
 .../authc/AuthenticationServiceTests.java     | 20 ++++++++++++++-----
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/authc/AuthenticationService.java b/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/authc/AuthenticationService.java
index c3888ba9453..037ed11ac1d 100644
--- a/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/authc/AuthenticationService.java
+++ b/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/authc/AuthenticationService.java
@@ -9,6 +9,7 @@ import org.apache.logging.log4j.message.ParameterizedMessage;
 import org.apache.logging.log4j.util.Supplier;
 import org.elasticsearch.ElasticsearchSecurityException;
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.ContextPreservingActionListener;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.common.component.AbstractComponent;
@@ -294,9 +295,9 @@ public class AuthenticationService extends AbstractComponent {
                     }
                 };
                 final IteratingActionListener<User, Realm> authenticatingListener =
-                        new IteratingActionListener<>(ActionListener.wrap(
-                                (user) -> consumeUser(user, messages),
-                                (e) -> listener.onFailure(request.exceptionProcessingRequest(e, token))),
+                    new IteratingActionListener<>(ContextPreservingActionListener.wrapPreservingContext(ActionListener.wrap(
+                        (user) -> consumeUser(user, messages),
+                        (e) -> listener.onFailure(request.exceptionProcessingRequest(e, token))), threadContext),
                         realmAuthenticatingConsumer, realmsList, threadContext);
                 try {
                     authenticatingListener.run();
diff --git a/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/authc/AuthenticationServiceTests.java b/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/authc/AuthenticationServiceTests.java
index 65f69b397ba..ef5b0386bc2 100644
--- a/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/authc/AuthenticationServiceTests.java
+++ b/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/authc/AuthenticationServiceTests.java
@@ -716,7 +716,7 @@ public class AuthenticationServiceTests extends ESTestCase {
         when(secondRealm.supports(token)).thenReturn(true);
         mockAuthenticate(secondRealm, token, new User("lookup user", new String[]{"user"}));
         mockRealmLookupReturnsNull(firstRealm, "run_as");
-        doThrow(authenticationError("realm doesn't want to " + "lookup"))
+        doThrow(authenticationError("realm doesn't want to lookup"))
                 .when(secondRealm).lookupUser(eq("run_as"), any(ActionListener.class));
 
         try {
@@ -1029,12 +1029,22 @@ public class AuthenticationServiceTests extends ESTestCase {
 
     @SuppressWarnings("unchecked")
     private void mockAuthenticate(Realm realm, AuthenticationToken token, User user) {
-        doAnswer((i) -> {
+        final boolean separateThread = randomBoolean();
+        doAnswer(i -> {
             ActionListener<AuthenticationResult> listener = (ActionListener<AuthenticationResult>) i.getArguments()[1];
-            if (user == null) {
-                listener.onResponse(AuthenticationResult.notHandled());
+            Runnable run = () -> {
+                if (user == null) {
+                    listener.onResponse(AuthenticationResult.notHandled());
+                } else {
+                    listener.onResponse(AuthenticationResult.success(user));
+                }
+            };
+            if (separateThread) {
+                final Thread thread = new Thread(run);
+                thread.start();
+                thread.join();
             } else {
-                listener.onResponse(AuthenticationResult.success(user));
+                run.run();
             }
             return null;
         }).when(realm).authenticate(eq(token), any(ActionListener.class));