content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements. See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License. You may obtain a copy of the License at
   *
   *    http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.kafka.common.security.oauthbearer;
  
  import org.apache.kafka.common.config.SaslConfigs;
  import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
  import org.apache.kafka.common.security.auth.Login;
  import org.apache.kafka.common.security.auth.SaslExtensions;
  import org.apache.kafka.common.security.auth.SaslExtensionsCallback;
  import org.apache.kafka.common.security.oauthbearer.internals.OAuthBearerSaslClientProvider;
  import org.apache.kafka.common.security.oauthbearer.internals.OAuthBearerSaslServerProvider;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.io.IOException;
  import java.util.Collections;
  import java.util.Iterator;
  import java.util.Map;
  import java.util.Objects;
  
  import javax.security.auth.Subject;
  import javax.security.auth.callback.Callback;
  import javax.security.auth.callback.CallbackHandler;
  import javax.security.auth.callback.UnsupportedCallbackException;
  import javax.security.auth.login.LoginException;
  import javax.security.auth.spi.LoginModule;
  
  /**
   * The {@code LoginModule} for the SASL/OAUTHBEARER mechanism. When a client
   * (whether a non-broker client or a broker when SASL/OAUTHBEARER is the
   * inter-broker protocol) connects to Kafka the {@code OAuthBearerLoginModule}
   * instance asks its configured {@link AuthenticateCallbackHandler}
   * implementation to handle an instance of {@link OAuthBearerTokenCallback} and
   * return an instance of {@link OAuthBearerToken}. A default, builtin
   * {@link AuthenticateCallbackHandler} implementation creates an unsecured token
   * as defined by these JAAS module options:
   * <p>
   * <table>
   * <tr>
   * <th>JAAS Module Option for Unsecured Token Retrieval</th>
   * <th>Documentation</th>
   * </tr>
   * <tr>
   * <td>{@code unsecuredLoginStringClaim_<claimname>="value"}</td>
   * <td>Creates a {@code String} claim with the given name and value. Any valid
   * claim name can be specified except '{@code iat}' and '{@code exp}' (these are
   * automatically generated).</td>
   * </tr>
   * <tr>
   * <td>{@code unsecuredLoginNumberClaim_<claimname>="value"}</td>
   * <td>Creates a {@code Number} claim with the given name and value. Any valid
   * claim name can be specified except '{@code iat}' and '{@code exp}' (these are
   * automatically generated).</td>
   * </tr>
   * <tr>
   * <td>{@code unsecuredLoginListClaim_<claimname>="value"}</td>
   * <td>Creates a {@code String List} claim with the given name and values parsed
   * from the given value where the first character is taken as the delimiter. For
   * example: {@code unsecuredLoginListClaim_fubar="|value1|value2"}. Any valid
   * claim name can be specified except '{@code iat}' and '{@code exp}' (these are
   * automatically generated).</td>
   * </tr>
   * <tr>
   * <td>{@code unsecuredLoginPrincipalClaimName}</td>
   * <td>Set to a custom claim name if you wish the name of the {@code String}
   * claim holding the principal name to be something other than
   * '{@code sub}'.</td>
   * </tr>
   * <tr>
   * <td>{@code unsecuredLoginLifetimeSeconds}</td>
   * <td>Set to an integer value if the token expiration is to be set to something
   * other than the default value of 3600 seconds (which is 1 hour). The
   * '{@code exp}' claim will be set to reflect the expiration time.</td>
   * </tr>
   * <tr>
   * <td>{@code unsecuredLoginScopeClaimName}</td>
   * <td>Set to a custom claim name if you wish the name of the {@code String} or
   * {@code String List} claim holding any token scope to be something other than
   * '{@code scope}'.</td>
   * </tr>
   * </table>
   * <p>
   * <p>
   * You can also add custom unsecured SASL extensions when using the default, builtin {@link AuthenticateCallbackHandler}
   * implementation through using the configurable option {@code unsecuredLoginExtension_<extensionname>}. Note that there
   * are validations for the key/values in order to conform to the SASL/OAUTHBEARER standard
   * (https://tools.ietf.org/html/rfc7628#section-3.1), including the reserved key at
   * {@link org.apache.kafka.common.security.oauthbearer.internals.OAuthBearerClientInitialResponse#AUTH_KEY}.
   * The {@code OAuthBearerLoginModule} instance also asks its configured {@link AuthenticateCallbackHandler}
   * implementation to handle an instance of {@link SaslExtensionsCallback} and return an instance of {@link SaslExtensions}.
   * The configured callback handler does not need to handle this callback, though -- any {@code UnsupportedCallbackException}
   * that is thrown is ignored, and no SASL extensions will be associated with the login.
   * <p>
   * Production use cases will require writing an implementation of
   * {@link AuthenticateCallbackHandler} that can handle an instance of
   * {@link OAuthBearerTokenCallback} and declaring it via either the
   * {@code sasl.login.callback.handler.class} configuration option for a
   * non-broker client or via the
   * {@code listener.name.sasl_ssl.oauthbearer.sasl.login.callback.handler.class}
   * configuration option for brokers (when SASL/OAUTHBEARER is the inter-broker
   * protocol).
   * <p>
   * This class stores the retrieved {@link OAuthBearerToken} in the
   * {@code Subject}'s private credentials where the {@code SaslClient} can
   * retrieve it. An appropriate, builtin {@code SaslClient} implementation is
   * automatically used and configured such that it can perform that retrieval.
   * <p>
   * Here is a typical, basic JAAS configuration for a client leveraging unsecured
   * SASL/OAUTHBEARER authentication:
   *
   * <pre>
   * KafkaClient {
   *      org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule Required
   *      unsecuredLoginStringClaim_sub="thePrincipalName";
   * };
   * </pre>
   *
   * An implementation of the {@link Login} interface specific to the
   * {@code OAUTHBEARER} mechanism is automatically applied; it periodically
   * refreshes any token before it expires so that the client can continue to make
   * connections to brokers. The parameters that impact how the refresh algorithm
   * operates are specified as part of the producer/consumer/broker configuration
   * and are as follows. See the documentation for these properties elsewhere for
   * details.
   * <p>
   * <table>
   * <tr>
   * <th>Producer/Consumer/Broker Configuration Property</th>
   * </tr>
   * <tr>
   * <td>{@code sasl.login.refresh.window.factor}</td>
   * </tr>
   * <tr>
   * <td>{@code sasl.login.refresh.window.jitter}</td>
   * </tr>
   * <tr>
   * <td>{@code sasl.login.refresh.min.period.seconds}</td>
   * </tr>
   * <tr>
   * <td>{@code sasl.login.refresh.min.buffer.seconds}</td>
   * </tr>
   * </table>
   * <p>
   * When a broker accepts a SASL/OAUTHBEARER connection the instance of the
   * builtin {@code SaslServer} implementation asks its configured
   * {@link AuthenticateCallbackHandler} implementation to handle an instance of
   * {@link OAuthBearerValidatorCallback} constructed with the OAuth 2 Bearer
   * Token's compact serialization and return an instance of
   * {@link OAuthBearerToken} if the value validates. A default, builtin
   * {@link AuthenticateCallbackHandler} implementation validates an unsecured
   * token as defined by these JAAS module options:
   * <p>
   * <table>
   * <tr>
   * <th>JAAS Module Option for Unsecured Token Validation</th>
   * <th>Documentation</th>
   * </tr>
   * <tr>
   * <td>{@code unsecuredValidatorPrincipalClaimName="value"}</td>
   * <td>Set to a non-empty value if you wish a particular {@code String} claim
   * holding a principal name to be checked for existence; the default is to check
   * for the existence of the '{@code sub}' claim.</td>
   * </tr>
   * <tr>
   * <td>{@code unsecuredValidatorScopeClaimName="value"}</td>
   * <td>Set to a custom claim name if you wish the name of the {@code String} or
   * {@code String List} claim holding any token scope to be something other than
   * '{@code scope}'.</td>
   * </tr>
   * <tr>
   * <td>{@code unsecuredValidatorRequiredScope="value"}</td>
   * <td>Set to a space-delimited list of scope values if you wish the
   * {@code String/String List} claim holding the token scope to be checked to
   * make sure it contains certain values.</td>
   * </tr>
   * <tr>
   * <td>{@code unsecuredValidatorAllowableClockSkewMs="value"}</td>
   * <td>Set to a positive integer value if you wish to allow up to some number of
   * positive milliseconds of clock skew (the default is 0).</td>
   * </tr>
   * </table>
   * <p>
   * Here is a typical, basic JAAS configuration for a broker leveraging unsecured
   * SASL/OAUTHBEARER validation:
   *
   * <pre>
   * KafkaServer {
   *      org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule Required
   *      unsecuredLoginStringClaim_sub="thePrincipalName";
   * };
   * </pre>
   *
   * Production use cases will require writing an implementation of
   * {@link AuthenticateCallbackHandler} that can handle an instance of
   * {@link OAuthBearerValidatorCallback} and declaring it via the
   * {@code listener.name.sasl_ssl.oauthbearer.sasl.server.callback.handler.class}
   * broker configuration option.
   * <p>
   * The builtin {@code SaslServer} implementation for SASL/OAUTHBEARER in Kafka
   * makes the instance of {@link OAuthBearerToken} available upon successful
   * authentication via the negotiated property "{@code OAUTHBEARER.token}"; the
   * token could be used in a custom authorizer (to authorize based on JWT claims
   * rather than ACLs, for example).
   * <p>
   * This implementation's {@code logout()} method will logout the specific token
   * that this instance logged in if it's {@code Subject} instance is shared
   * across multiple {@code LoginContext}s and there happen to be multiple tokens
   * on the {@code Subject}. This functionality is useful because it means a new
   * token with a longer lifetime can be created before a soon-to-expire token is
   * actually logged out. Otherwise, if multiple simultaneous tokens were not
   * supported like this, the soon-to-be expired token would have to be logged out
   * first, and then if the new token could not be retrieved (maybe the
   * authorization server is temporarily unavailable, for example) the client
   * would be left without a token and would be unable to create new connections.
   * Better to mitigate this possibility by leaving the existing token (which
   * still has some lifetime left) in place until a new replacement token is
   * actually retrieved. This implementation supports this.
   *
   * @see SaslConfigs#SASL_LOGIN_REFRESH_WINDOW_FACTOR_DOC
   * @see SaslConfigs#SASL_LOGIN_REFRESH_WINDOW_JITTER_DOC
   * @see SaslConfigs#SASL_LOGIN_REFRESH_MIN_PERIOD_SECONDS_DOC
   * @see SaslConfigs#SASL_LOGIN_REFRESH_BUFFER_SECONDS_DOC
   */
  public class OAuthBearerLoginModule implements LoginModule {
  
      /**
       * Login state transitions:
       *   Initial state: NOT_LOGGED_IN
       *   login()      : NOT_LOGGED_IN => LOGGED_IN_NOT_COMMITTED
       *   commit()     : LOGGED_IN_NOT_COMMITTED => COMMITTED
       *   abort()      : LOGGED_IN_NOT_COMMITTED => NOT_LOGGED_IN
       *   logout()     : Any state => NOT_LOGGED_IN
       */
      private enum LoginState {
          NOT_LOGGED_IN,
          LOGGED_IN_NOT_COMMITTED,
          COMMITTED
      }
  
      /**
       * The SASL Mechanism name for OAuth 2: {@code OAUTHBEARER}
       */
      public static final String OAUTHBEARER_MECHANISM = "OAUTHBEARER";
      private static final Logger log = LoggerFactory.getLogger(OAuthBearerLoginModule.class);
      private static final SaslExtensions EMPTY_EXTENSIONS = new SaslExtensions(Collections.emptyMap());
      private Subject subject = null;
      private AuthenticateCallbackHandler callbackHandler = null;
      private OAuthBearerToken tokenRequiringCommit = null;
      private OAuthBearerToken myCommittedToken = null;
      private SaslExtensions extensionsRequiringCommit = null;
      private SaslExtensions myCommittedExtensions = null;
      private LoginState loginState;
  
      static {
          OAuthBearerSaslClientProvider.initialize(); // not part of public API
          OAuthBearerSaslServerProvider.initialize(); // not part of public API
      }
  
      @Override
      public void initialize(Subject subject, CallbackHandler callbackHandler, Map<String, ?> sharedState,
              Map<String, ?> options) {
          this.subject = Objects.requireNonNull(subject);
          if (!(Objects.requireNonNull(callbackHandler) instanceof AuthenticateCallbackHandler))
              throw new IllegalArgumentException(String.format("Callback handler must be castable to %s: %s",
                      AuthenticateCallbackHandler.class.getName(), callbackHandler.getClass().getName()));
          this.callbackHandler = (AuthenticateCallbackHandler) callbackHandler;
      }
  
      @Override
      public boolean login() throws LoginException {
          if (loginState == LoginState.LOGGED_IN_NOT_COMMITTED) {
              if (tokenRequiringCommit != null)
                  throw new IllegalStateException(String.format(
                      "Already have an uncommitted token with private credential token count=%d", committedTokenCount()));
              else
                  throw new IllegalStateException("Already logged in without a token");
          }
          if (loginState == LoginState.COMMITTED) {
              if (myCommittedToken != null)
                  throw new IllegalStateException(String.format(
                      "Already have a committed token with private credential token count=%d; must login on another login context or logout here first before reusing the same login context",
                      committedTokenCount()));
              else
                  throw new IllegalStateException("Login has already been committed without a token");
          }
  
          identifyToken();
          if (tokenRequiringCommit != null)
              identifyExtensions();
          else
              log.debug("Logged in without a token, this login cannot be used to establish client connections");
  
          loginState = LoginState.LOGGED_IN_NOT_COMMITTED;
          log.debug("Login succeeded; invoke commit() to commit it; current committed token count={}",
                  committedTokenCount());
          return true;
      }
  
      private void identifyToken() throws LoginException {
          OAuthBearerTokenCallback tokenCallback = new OAuthBearerTokenCallback();
          try {
              callbackHandler.handle(new Callback[] {tokenCallback});
          } catch (IOException | UnsupportedCallbackException e) {
              log.error(e.getMessage(), e);
              throw new LoginException("An internal error occurred while retrieving token from callback handler");
          }
  
          tokenRequiringCommit = tokenCallback.token();
          if (tokenCallback.errorCode() != null) {
              log.info("Login failed: {} : {} (URI={})", tokenCallback.errorCode(), tokenCallback.errorDescription(),
                      tokenCallback.errorUri());
              throw new LoginException(tokenCallback.errorDescription());
          }
      }
  
      /**
       * Attaches SASL extensions to the Subject
       */
      private void identifyExtensions() throws LoginException {
          SaslExtensionsCallback extensionsCallback = new SaslExtensionsCallback();
          try {
              callbackHandler.handle(new Callback[] {extensionsCallback});
              extensionsRequiringCommit = extensionsCallback.extensions();
          } catch (IOException e) {
              log.error(e.getMessage(), e);
              throw new LoginException("An internal error occurred while retrieving SASL extensions from callback handler");
          } catch (UnsupportedCallbackException e) {
              extensionsRequiringCommit = EMPTY_EXTENSIONS;
              log.debug("CallbackHandler {} does not support SASL extensions. No extensions will be added", callbackHandler.getClass().getName());
          }
          if (extensionsRequiringCommit ==  null) {
              log.error("SASL Extensions cannot be null. Check whether your callback handler is explicitly setting them as null.");
              throw new LoginException("Extensions cannot be null.");
          }
      }
  
      @Override
      public boolean logout() {
          if (loginState == LoginState.LOGGED_IN_NOT_COMMITTED)
              throw new IllegalStateException(
                      "Cannot call logout() immediately after login(); need to first invoke commit() or abort()");
          if (loginState != LoginState.COMMITTED) {
              log.debug("Nothing here to log out");
              return false;
          }
          if (myCommittedToken != null) {
              log.trace("Logging out my token; current committed token count = {}", committedTokenCount());
              for (Iterator<Object> iterator = subject.getPrivateCredentials().iterator(); iterator.hasNext(); ) {
                  Object privateCredential = iterator.next();
                  if (privateCredential == myCommittedToken) {
                      iterator.remove();
                      myCommittedToken = null;
                      break;
                  }
              }
              log.debug("Done logging out my token; committed token count is now {}", committedTokenCount());
          } else
              log.debug("No tokens to logout for this login");
  
          if (myCommittedExtensions != null) {
              log.trace("Logging out my extensions");
              if (subject.getPublicCredentials().removeIf(e -> myCommittedExtensions == e))
                  myCommittedExtensions = null;
              log.debug("Done logging out my extensions");
          } else
              log.debug("No extensions to logout for this login");
  
          loginState = LoginState.NOT_LOGGED_IN;
          return true;
      }
  
      @Override
      public boolean commit() {
          if (loginState != LoginState.LOGGED_IN_NOT_COMMITTED) {
              log.debug("Nothing here to commit");
              return false;
          }
  
          if (tokenRequiringCommit != null) {
              log.trace("Committing my token; current committed token count = {}", committedTokenCount());
              subject.getPrivateCredentials().add(tokenRequiringCommit);
              myCommittedToken = tokenRequiringCommit;
              tokenRequiringCommit = null;
              log.debug("Done committing my token; committed token count is now {}", committedTokenCount());
          } else
              log.debug("No tokens to commit, this login cannot be used to establish client connections");
  
          if (extensionsRequiringCommit != null) {
              subject.getPublicCredentials().add(extensionsRequiringCommit);
              myCommittedExtensions = extensionsRequiringCommit;
              extensionsRequiringCommit = null;
          }
  
          loginState = LoginState.COMMITTED;
          return true;
      }
  
      @Override
      public boolean abort() {
          if (loginState == LoginState.LOGGED_IN_NOT_COMMITTED) {
              log.debug("Login aborted");
              tokenRequiringCommit = null;
              extensionsRequiringCommit = null;
              loginState = LoginState.NOT_LOGGED_IN;
              return true;
          }
          log.debug("Nothing here to abort");
          return false;
      }
  
      private int committedTokenCount() {
          return subject.getPrivateCredentials(OAuthBearerToken.class).size();
      }
  }
  