Allow an AuthenticationResult to return metadata (#34382)

PR #34290 made it impossible to use thread-context values to pass
authentication metadata out of a realm. The SAML realm used this
technique to allow the SamlAuthenticateAction to process the parsed
SAML token, and apply them to the access token that was generated.

This new method adds metadata to the AuthenticationResult itself, and
then the authentication service makes this result available on the
thread context.

Closes: #34332

Author: tvernum

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/authc/AuthenticationResult.java

@@ -8,6 +8,8 @@
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.xpack.core.security.user.User;
 
+import java.util.Collections;
+import java.util.Map;
 import java.util.Objects;
 
 /**
@@ -21,7 +23,9 @@
  * </ol>
  */
 public final class AuthenticationResult {
-    private static final AuthenticationResult NOT_HANDLED = new AuthenticationResult(Status.CONTINUE, null, null, null);
+    private static final AuthenticationResult NOT_HANDLED = new AuthenticationResult(Status.CONTINUE, null, null, null, null);
+
+    public static String THREAD_CONTEXT_KEY = "_xpack_security_auth_result";
 
     public enum Status {
         SUCCESS,
@@ -33,12 +37,15 @@ public enum Status {
     private final User user;
     private final String message;
     private final Exception exception;
+    private final Map<String, Object> metadata;
 
-    private AuthenticationResult(Status status, @Nullable User user, @Nullable String message, @Nullable Exception exception) {
+    private AuthenticationResult(Status status, @Nullable User user, @Nullable String message, @Nullable Exception exception,
+                                 @Nullable Map<String, Object> metadata) {
         this.status = status;
         this.user = user;
         this.message = message;
         this.exception = exception;
+        this.metadata = metadata == null ? Collections.emptyMap() : Collections.unmodifiableMap(metadata);
     }
 
     public Status getStatus() {
@@ -57,6 +64,10 @@ public Exception getException() {
         return exception;
     }
 
+    public Map<String, Object> getMetadata() {
+        return metadata;
+    }
+
     /**
      * Creates an {@code AuthenticationResult} that indicates that the supplied {@link User}
      * has been successfully authenticated.
@@ -69,7 +80,16 @@ public Exception getException() {
      */
     public static AuthenticationResult success(User user) {
         Objects.requireNonNull(user);
-        return new AuthenticationResult(Status.SUCCESS, user, null, null);
+        return success(user, null);
+    }
+
+    /**
+     * Creates a successful result, with optional metadata
+     *
+     * @see #success(User)
+     */
+    public static AuthenticationResult success(User user, @Nullable Map<String, Object> metadata) {
+        return new AuthenticationResult(Status.SUCCESS, user, null, null, metadata);
     }
 
     /**
@@ -96,7 +116,7 @@ public static AuthenticationResult notHandled() {
      */
     public static AuthenticationResult unsuccessful(String message, @Nullable Exception cause) {
         Objects.requireNonNull(message);
-        return new AuthenticationResult(Status.CONTINUE, null, message, cause);
+        return new AuthenticationResult(Status.CONTINUE, null, message, cause, null);
     }
 
     /**
@@ -110,7 +130,7 @@ public static AuthenticationResult unsuccessful(String message, @Nullable Except
      * </p>
      */
     public static AuthenticationResult terminate(String message, @Nullable Exception cause) {
-        return new AuthenticationResult(Status.TERMINATE, null, message, cause);
+        return new AuthenticationResult(Status.TERMINATE, null, message, cause, null);
     }
 
     public boolean isAuthenticated() {

x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/action/saml/TransportSamlAuthenticateAction.java

@@ -22,6 +22,7 @@
 import org.elasticsearch.xpack.core.security.action.saml.SamlAuthenticateRequest;
 import org.elasticsearch.xpack.core.security.action.saml.SamlAuthenticateResponse;
 import org.elasticsearch.xpack.core.security.authc.Authentication;
+import org.elasticsearch.xpack.core.security.authc.AuthenticationResult;
 import org.elasticsearch.xpack.security.authc.AuthenticationService;
 import org.elasticsearch.xpack.security.authc.TokenService;
 import org.elasticsearch.xpack.security.authc.saml.SamlRealm;
@@ -54,7 +55,12 @@ protected void doExecute(Task task, SamlAuthenticateRequest request, ActionListe
         Authentication originatingAuthentication = Authentication.getAuthentication(threadContext);
         try (ThreadContext.StoredContext ignore = threadContext.stashContext()) {
             authenticationService.authenticate(SamlAuthenticateAction.NAME, request, saml, ActionListener.wrap(authentication -> {
-                final Map<String, Object> tokenMeta = threadContext.getTransient(SamlRealm.CONTEXT_TOKEN_DATA);
+                AuthenticationResult result = threadContext.getTransient(AuthenticationResult.THREAD_CONTEXT_KEY);
+                if (result == null) {
+                    listener.onFailure(new IllegalStateException("Cannot find AuthenticationResult on thread context"));
+                    return;
+                }
+                final Map<String, Object> tokenMeta = (Map<String, Object>) result.getMetadata().get(SamlRealm.CONTEXT_TOKEN_DATA);
                 tokenService.createUserToken(authentication, originatingAuthentication,
                         ActionListener.wrap(tuple -> {
                             final String tokenString = tokenService.getUserTokenString(tuple.v1());

x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/authc/AuthenticationService.java

@@ -140,6 +140,7 @@ class Authenticator {
         private RealmRef authenticatedBy = null;
         private RealmRef lookedupBy = null;
         private AuthenticationToken authenticationToken = null;
+        private AuthenticationResult authenticationResult = null;
 
         Authenticator(RestRequest request, ActionListener<Authentication> listener) {
             this(new AuditableRestRequest(auditTrail, failureHandler, threadContext, request), null, listener);
@@ -267,6 +268,7 @@ private void consumeToken(AuthenticationToken token) {
                             if (result.getStatus() == AuthenticationResult.Status.SUCCESS) {
                                 // user was authenticated, populate the authenticated by information
                                 authenticatedBy = new RealmRef(realm.name(), realm.type(), nodeName);
+                                authenticationResult = result;
                                 userListener.onResponse(result.getUser());
                             } else {
                                 // the user was not authenticated, call this so we can audit the correct event
@@ -360,6 +362,7 @@ private void consumeUser(User user, Map<Realm, Tuple<String, Exception>> message
                 });
                 listener.onFailure(request.authenticationFailed(authenticationToken));
             } else {
+                threadContext.putTransient(AuthenticationResult.THREAD_CONTEXT_KEY, authenticationResult);
                 if (runAsEnabled) {
                     final String runAsUsername = threadContext.getHeader(AuthenticationServiceField.RUN_AS_USER_HEADER);
                     if (runAsUsername != null && runAsUsername.isEmpty() == false) {

x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/authc/saml/SamlRealm.java

@@ -426,7 +426,10 @@ private void buildUser(SamlAttributes attributes, ActionListener<AuthenticationR
         final Map<String, Object> tokenMetadata = createTokenMetadata(attributes.name(), attributes.session());
         ActionListener<AuthenticationResult> wrappedListener = ActionListener.wrap(auth -> {
             if (auth.isAuthenticated()) {
-                config.threadContext().putTransient(CONTEXT_TOKEN_DATA, tokenMetadata);
+                // Add the SAML token details as metadata on the authentication
+                Map<String, Object> metadata = new HashMap<>(auth.getMetadata());
+                metadata.put(CONTEXT_TOKEN_DATA, tokenMetadata);
+                auth = AuthenticationResult.success(auth.getUser(), metadata);
             }
             baseListener.onResponse(auth);
         }, baseListener::onFailure);