OAuth2 authenticator for support this kind of auth in JavaMail IMAP folder requests

OAuth2Authenticator.scala

import com.typesafe.scalalogging.slf4j.Logging

import com.sun.mail.imap.IMAPSSLStore
import javax.mail.{Store, Session}
import java.security.{Provider, Security}
import java.util.Properties

import OAuth2SaslClientFactory

/**
 * Performs OAuth2 authentication.
 */
object OAuth2Authenticator extends Logging {

  val SOCKET_CONNECTION_TIMEOUT_MS = 10000

  val SOCKET_IO_TIMEOUT_MS = 10000

  /**
   * Connects and authenticates to an IMAP server with OAuth2.
   *
   * @param host        Hostname of the imap server, eg. imap.gmail.com.
   * @param port        Port of the imap server, for example 993
   * @param userEmail   Email address of the user to authenticate, eg. [email protected].
   * @param oauthToken  The user's OAuth token.
   * @param debug       Whether to enable debug logging on the IMAP connection.
   *
   * @return An authenticated IMAPStore that can be used for IMAP operations.
   */
  def getIMAPStore(host: String,
                    port: Int,
                    userEmail: String,
                    oauthToken: String,
                    debug: Boolean): Store = {

    /**
     * Installs the OAuth2 SASL provider. This is needed for functioning of SASL authentication
     * for requesting messages in imap folders.
     */
    Security.addProvider(new OAuth2Provider())

    val props = new Properties()
    props.setProperty("mail.imaps.connectiontimeout", SOCKET_CONNECTION_TIMEOUT_MS.toString)
    props.setProperty("mail.imaps.timeout", SOCKET_IO_TIMEOUT_MS.toString)
    props.setProperty("mail.imaps.sasl.enable", "true")
    props.setProperty("mail.imaps.sasl.mechanisms", "XOAUTH2")
    props.setProperty(OAuth2SaslClientFactory.OAUTH_TOKEN_PROP, oauthToken)
    val session = Session.getInstance(props)
    session.setDebug(debug)

    new IMAPSSLStore(session, null)
  }

}

class OAuth2Provider extends Provider("Tocobox OAuth2 Provider", 1.0, "Provides the XOAUTH2 SASL Mechanism") {
  val serialVersionUID: Long = 1L
  put("SaslClientFactory.XOAUTH2", "com.tocobox.util.oauth2.OAuth2SaslClientFactory")
}

OAuth2SaslClient.java

import java.io.IOException;
import java.util.logging.Logger;

import javax.security.auth.callback.Callback;
import javax.security.auth.callback.CallbackHandler;
import javax.security.auth.callback.NameCallback;
import javax.security.auth.callback.UnsupportedCallbackException;
import javax.security.sasl.SaslClient;
import javax.security.sasl.SaslException;


/**
 * An OAuth2 implementation of SaslClient.
 */
class OAuth2SaslClient implements SaslClient {
    private static final Logger logger =
            Logger.getLogger(OAuth2SaslClient.class.getName());

    private final String oauthToken;
    private final CallbackHandler callbackHandler;

    private boolean isComplete = false;

    /**
     * Creates a new instance of the OAuth2SaslClient. This will ordinarily only
     * be called from OAuth2SaslClientFactory.
     */
    public OAuth2SaslClient(String oauthToken,
                            CallbackHandler callbackHandler) {
        this.oauthToken = oauthToken;
        this.callbackHandler = callbackHandler;
    }

    public String getMechanismName() {
        return "XOAUTH2";
    }

    public boolean hasInitialResponse() {
        return true;
    }

    public byte[] evaluateChallenge(byte[] challenge) throws SaslException {
        if (isComplete) {
            // Empty final response from server, just ignore it.
            return new byte[] { };
        }

        NameCallback nameCallback = new NameCallback("Enter name");
        Callback[] callbacks = new Callback[] { nameCallback };
        try {
            callbackHandler.handle(callbacks);
        } catch (UnsupportedCallbackException e) {
            throw new SaslException("Unsupported callback: " + e);
        } catch (IOException e) {
            throw new SaslException("Failed to execute callback: " + e);
        }
        String email = nameCallback.getName();

        byte[] response = String.format("user=%s\1auth=Bearer %s\1\1", email,
                oauthToken).getBytes();
        isComplete = true;
        return response;
    }

    public boolean isComplete() {
        return isComplete;
    }

    public byte[] unwrap(byte[] incoming, int offset, int len)
            throws SaslException {
        throw new IllegalStateException();
    }

    public byte[] wrap(byte[] outgoing, int offset, int len)
            throws SaslException {
        throw new IllegalStateException();
    }

    public Object getNegotiatedProperty(String propName) {
        if (!isComplete()) {
            throw new IllegalStateException();
        }
        return null;
    }

    public void dispose() throws SaslException {
    }
}

OAuth2SaslClientFactory.java

import java.util.Map;

import java.util.logging.Logger;

import javax.security.auth.callback.CallbackHandler;
import javax.security.sasl.SaslClient;
import javax.security.sasl.SaslClientFactory;

/**
 * A SaslClientFactory that returns instances of OAuth2SaslClient.
 *
 * <p>Only the "XOAUTH2" mechanism is supported. The {@code callbackHandler} is
 * passed to the OAuth2SaslClient. Other parameters are ignored.
 */
public class OAuth2SaslClientFactory implements SaslClientFactory {
    private static final Logger logger =
            Logger.getLogger(OAuth2SaslClientFactory.class.getName());

    public static final String OAUTH_TOKEN_PROP =
            "mail.imaps.sasl.mechanisms.oauth2.oauthToken";

    public SaslClient createSaslClient(String[] mechanisms,
                                       String authorizationId,
                                       String protocol,
                                       String serverName,
                                       Map<String, ?> props,
                                       CallbackHandler callbackHandler) {
        boolean matchedMechanism = false;
        for (int i = 0; i < mechanisms.length; ++i) {
            if ("XOAUTH2".equalsIgnoreCase(mechanisms[i])) {
                matchedMechanism = true;
                break;
            }
        }
        if (!matchedMechanism) {
            logger.info("Failed to match any mechanisms");
            return null;
        }
        return new OAuth2SaslClient((String) props.get(OAUTH_TOKEN_PROP),
                callbackHandler);
    }

    public String[] getMechanismNames(Map<String, ?> props) {
        return new String[] {"XOAUTH2"};
    }
}