Custom authenticator followup + tests

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

@@ -7,17 +7,26 @@
 
 package org.elasticsearch.xpack.security.authc;
 
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.xpack.core.common.IteratingActionListener;
 import org.elasticsearch.xpack.core.security.authc.Authentication;
 import org.elasticsearch.xpack.core.security.authc.AuthenticationResult;
 import org.elasticsearch.xpack.core.security.authc.AuthenticationToken;
 import org.elasticsearch.xpack.core.security.authc.CustomAuthenticator;
 
 import java.util.Collections;
 import java.util.List;
+import java.util.function.BiConsumer;
+import java.util.function.Function;
+
+import static org.elasticsearch.common.Strings.format;
 
 public class PluggableAuthenticatorChain implements Authenticator {
 
+    private static final Logger logger = LogManager.getLogger(PluggableAuthenticatorChain.class);
+
     private final List<CustomAuthenticator> customAuthenticators;
 
     public PluggableAuthenticatorChain(List<CustomAuthenticator> customAuthenticators) {
@@ -59,28 +68,47 @@ public void authenticate(Context context, ActionListener<AuthenticationResult<Au
         }
         AuthenticationToken token = context.getMostRecentAuthenticationToken();
         if (token != null) {
-            // TODO switch to IteratingActionListener
-            for (CustomAuthenticator customAuthenticator : customAuthenticators) {
-                if (customAuthenticator.supports(token)) {
-                    customAuthenticator.authenticate(token, ActionListener.wrap(response -> {
-                        if (response.isAuthenticated()) {
-                            listener.onResponse(response);
-                        } else if (response.getStatus() == AuthenticationResult.Status.TERMINATE) {
-                            final Exception ex = response.getException();
-                            if (ex == null) {
-                                listener.onFailure(context.getRequest().authenticationFailed(token));
-                            } else {
-                                listener.onFailure(context.getRequest().exceptionProcessingRequest(ex, token));
-                            }
-                        } else if (response.getStatus() == AuthenticationResult.Status.CONTINUE) {
-                            listener.onResponse(AuthenticationResult.notHandled());
-                        }
-                    }, ex -> listener.onFailure(context.getRequest().exceptionProcessingRequest(ex, token))));
-                    return;
-                }
+            var iteratingListener = new IteratingActionListener<>(
+                listener,
+                getAuthConsumer(context),
+                customAuthenticators,
+                context.getThreadContext(),
+                Function.identity(),
+                result -> result.getStatus() == AuthenticationResult.Status.CONTINUE
+            );
+            try {
+                iteratingListener.run();
+            } catch (Exception e) {
+                logger.debug(() -> format("Authentication of token [%s] failed", token.getClass().getName()), e);
+                listener.onFailure(context.getRequest().exceptionProcessingRequest(e, token));
             }
+            return;
         }
         listener.onResponse(AuthenticationResult.notHandled());
     }
 
+    private BiConsumer<CustomAuthenticator, ActionListener<AuthenticationResult<Authentication>>> getAuthConsumer(Context context) {
+        AuthenticationToken token = context.getMostRecentAuthenticationToken();
+        return (authenticator, iteratingListener) -> {
+            if (authenticator.supports(token)) {
+                authenticator.authenticate(token, ActionListener.wrap(response -> {
+                    if (response.isAuthenticated()) {
+                        iteratingListener.onResponse(response);
+                    } else if (response.getStatus() == AuthenticationResult.Status.TERMINATE) {
+                        final Exception ex = response.getException();
+                        if (ex == null) {
+                            iteratingListener.onFailure(context.getRequest().authenticationFailed(token));
+                        } else {
+                            iteratingListener.onFailure(context.getRequest().exceptionProcessingRequest(ex, token));
+                        }
+                    } else if (response.getStatus() == AuthenticationResult.Status.CONTINUE) {
+                        iteratingListener.onResponse(AuthenticationResult.notHandled());
+                    }
+                }, ex -> iteratingListener.onFailure(context.getRequest().exceptionProcessingRequest(ex, token))));
+            } else {
+                iteratingListener.onResponse(AuthenticationResult.notHandled()); // try the next custom authenticator
+            }
+        };
+    }
+
 }