Access token callback

src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java

@@ -204,6 +204,7 @@ final class TDS {
     static final byte ADALWORKFLOW_ACTIVEDIRECTORYMANAGEDIDENTITY = 0x03;
     static final byte ADALWORKFLOW_ACTIVEDIRECTORYINTERACTIVE = 0x03;
     static final byte ADALWORKFLOW_DEFAULTAZURECREDENTIAL = 0x03;
+    static final byte ADALWORKFLOW_ACCESSTOKENCALLBACK = 0x03;
     static final byte ADALWORKFLOW_ACTIVEDIRECTORYSERVICEPRINCIPAL = 0x01; // Using the Password byte as that is the
                                                                            // closest we have.
     static final byte FEDAUTH_INFO_ID_STSURL = 0x01; // FedAuthInfoData is token endpoint URL from which to acquire fed

src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerConnection.java

@@ -411,13 +411,18 @@ CallableStatement prepareCall(String sql, int nType, int nConcur, int nHold,
     /**
      * Deprecated. Time-to-live is no longer supported for the cached Managed Identity tokens.
      * This method will always return 0 and is for backwards compatibility only.
+     *
+     * @return Method will always return 0.
      */
     @Deprecated
     int getMsiTokenCacheTtl();
 
     /**
      * Deprecated. Time-to-live is no longer supported for the cached Managed Identity tokens.
      * This method is a no-op for backwards compatibility only.
+     *
+     * @param timeToLive
+     *        Time-to-live is no longer supported.
      */
     @Deprecated
     void setMsiTokenCacheTtl(int timeToLive);

src/main/java/com/microsoft/sqlserver/jdbc/ISQLServerDataSource.java

@@ -1216,14 +1216,34 @@ public interface ISQLServerDataSource extends javax.sql.CommonDataSource {
     /**
      * Deprecated. Time-to-live is no longer supported for the cached Managed Identity tokens.
      * This method is a no-op for backwards compatibility only.
+     *
+     * @param timeToLive
+     *        Time-to-live is no longer supported.
      */
     @Deprecated
     void setMsiTokenCacheTtl(int timeToLive);
 
     /**
      * Deprecated. Time-to-live is no longer supported for the cached Managed Identity tokens.
      * This method will always return 0 and is for backwards compatibility only.
+     *
+     * @return Method will always return 0.
      */
     @Deprecated
     int getMsiTokenCacheTtl();
+
+    /**
+     * Sets the {@link SQLServerAccessTokenCallback} delegate.
+     *
+     * @param accessTokenCallback
+     *        Access token callback delegate.
+     */
+    void setAccessTokenCallback(SQLServerAccessTokenCallback accessTokenCallback);
+
+    /**
+     * Returns a {@link SQLServerAccessTokenCallback}, the access token callback delegate.
+     *
+     * @return Access token callback delegate.
+     */
+    SQLServerAccessTokenCallback getAccessTokenCallback();
 }

src/main/java/com/microsoft/sqlserver/jdbc/SQLServerAccessTokenCallback.java

@@ -0,0 +1,21 @@
+package com.microsoft.sqlserver.jdbc;
+
+/**
+ * Provides SqlAuthenticationToken callback to be implemented by client code.
+ */
+public interface SQLServerAccessTokenCallback {
+
+    /**
+     * For an example of callback usage, look under the project's code samples.
+     *
+     * Returns the access token for the authentication request
+     *
+     * @param stsurl
+     *        - Security token service URL.
+     * @param spn
+     *        - Service principal name.
+     *
+     * @return Returns a {@link SqlAuthenticationToken}.
+     */
+    SqlAuthenticationToken getAccessToken(String stsurl, String spn);
+}

src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java

@@ -162,7 +162,7 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
     private byte[] accessTokenInByte = null;
 
     /** fedAuth token */
-    private SqlFedAuthToken fedAuthToken = null;
+    private SqlAuthenticationToken fedAuthToken = null;
 
     /** original hostNameInCertificate */
     private String originalHostNameInCertificate = null;
@@ -537,6 +537,13 @@ class FederatedAuthenticationFeatureExtensionData implements Serializable {
                     this.authentication = SqlAuthentication.ActiveDirectoryInteractive;
                     break;
                 default:
+                    // If authenticationString not specified, check if access token callback was set.
+                    // If access token callback is set, break.
+                    if (null != activeConnectionProperties
+                            .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString())) {
+                        this.authentication = SqlAuthentication.NotSpecified;
+                        break;
+                    }
                     assert (false);
                     MessageFormat form = new MessageFormat(
                             SQLServerException.getErrString("R_InvalidConnectionSetting"));
@@ -1669,7 +1676,7 @@ void checkClosed() throws SQLServerException {
      * @return true/false
      */
     protected boolean needsReconnect() {
-        return (null != fedAuthToken && Util.checkIfNeedNewAccessToken(this, fedAuthToken.expiresOn));
+        return (null != fedAuthToken && Util.checkIfNeedNewAccessToken(this, fedAuthToken.getExpiresOn()));
     }
 
     /**
@@ -2423,6 +2430,17 @@ Connection connectInternal(Properties propsIn,
                     ntlmAuthentication = true;
                 }
 
+                SQLServerAccessTokenCallback callback = (SQLServerAccessTokenCallback) activeConnectionProperties
+                        .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString());
+
+                if (null != callback && (!activeConnectionProperties
+                        .getProperty(SQLServerDriverStringProperty.USER.toString()).isEmpty()
+                        || !activeConnectionProperties.getProperty(SQLServerDriverStringProperty.PASSWORD.toString())
+                                .isEmpty())) {
+                    throw new SQLServerException(
+                            SQLServerException.getErrString("R_AccessTokenCallbackWithUserPassword"), null);
+                }
+
                 sPropKey = SQLServerDriverStringProperty.AUTHENTICATION.toString();
                 sPropValue = activeConnectionProperties.getProperty(sPropKey);
                 if (null == sPropValue) {
@@ -2434,7 +2452,7 @@ Connection connectInternal(Properties propsIn,
                         && (!activeConnectionProperties.getProperty(SQLServerDriverStringProperty.PASSWORD.toString())
                                 .isEmpty())) {
                     MessageFormat form = new MessageFormat(
-                            SQLServerException.getErrString("R_MSIAuthenticationWithPassword"));
+                            SQLServerException.getErrString("R_ManagedIdentityAuthenticationWithPassword"));
                     throw new SQLServerException(form.format(new Object[] {authenticationString}), null);
                 }
 
@@ -2466,7 +2484,7 @@ Connection connectInternal(Properties propsIn,
                         && (!activeConnectionProperties.getProperty(SQLServerDriverStringProperty.PASSWORD.toString())
                                 .isEmpty())) {
                     MessageFormat form = new MessageFormat(
-                            SQLServerException.getErrString("R_MSIAuthenticationWithPassword"));
+                            SQLServerException.getErrString("R_ManagedIdentityAuthenticationWithPassword"));
                     throw new SQLServerException(form.format(new Object[] {authenticationString}), null);
                 }
 
@@ -3464,7 +3482,8 @@ private void executeReconnect(LogonCommand logonCommand) throws SQLServerExcepti
     void prelogin(String serverName, int portNumber) throws SQLServerException {
         // Build a TDS Pre-Login packet to send to the server.
         if ((!authenticationString.equalsIgnoreCase(SqlAuthentication.NotSpecified.toString()))
-                || (null != accessTokenInByte)) {
+                || (null != accessTokenInByte) || null != activeConnectionProperties
+                        .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString())) {
             fedAuthRequiredByUser = true;
         }
 
@@ -3846,7 +3865,8 @@ void prelogin(String serverName, int portNumber) throws SQLServerException {
                     // Or AccessToken is not null, mean token based authentication is used.
                     if (((null != authenticationString)
                             && (!authenticationString.equalsIgnoreCase(SqlAuthentication.NotSpecified.toString())))
-                            || (null != accessTokenInByte)) {
+                            || (null != accessTokenInByte) || null != activeConnectionProperties
+                                    .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString())) {
                         fedAuthRequiredPreLoginResponse = (preloginResponse[optionOffset] == 1);
                     }
                     break;
@@ -4790,6 +4810,12 @@ int writeFedAuthFeatureRequest(boolean write, /* if false just calculates the le
                             workflow = TDS.ADALWORKFLOW_ACTIVEDIRECTORYSERVICEPRINCIPAL;
                             break;
                         default:
+                            // If not specified, check if access token callback was set. If it is set, break.
+                            if (null != activeConnectionProperties
+                                    .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString())) {
+                                workflow = TDS.ADALWORKFLOW_ACCESSTOKENCALLBACK;
+                                break;
+                            }
                             assert (false); // Unrecognized Authentication type for fedauth ADAL request
                             break;
                     }
@@ -5005,7 +5031,9 @@ private void logon(LogonCommand command) throws SQLServerException {
                                 .equalsIgnoreCase(SqlAuthentication.ActiveDirectoryServicePrincipal.toString())
                         || authenticationString
                                 .equalsIgnoreCase(SqlAuthentication.ActiveDirectoryInteractive.toString()))
-                        && fedAuthRequiredPreLoginResponse)) {
+                        && fedAuthRequiredPreLoginResponse)
+                || null != activeConnectionProperties
+                        .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString())) {
             federatedAuthenticationInfoRequested = true;
             fedAuthFeatureExtensionData = new FederatedAuthenticationFeatureExtensionData(TDS.TDS_FEDAUTH_LIBRARY_ADAL,
                     authenticationString, fedAuthRequiredPreLoginResponse);
@@ -5499,9 +5527,9 @@ final class FedAuthTokenCommand extends UninterruptableTDSCommand {
         // Always update serialVersionUID when prompted.
         private static final long serialVersionUID = 1L;
         TDSTokenHandler tdsTokenHandler = null;
-        SqlFedAuthToken sqlFedAuthToken = null;
+        SqlAuthenticationToken sqlFedAuthToken = null;
 
-        FedAuthTokenCommand(SqlFedAuthToken sqlFedAuthToken, TDSTokenHandler tdsTokenHandler) {
+        FedAuthTokenCommand(SqlAuthenticationToken sqlFedAuthToken, TDSTokenHandler tdsTokenHandler) {
             super("FedAuth");
             this.tdsTokenHandler = tdsTokenHandler;
             this.sqlFedAuthToken = sqlFedAuthToken;
@@ -5530,7 +5558,16 @@ void onFedAuthInfo(SqlFedAuthInfo fedAuthInfo, TDSTokenHandler tdsTokenHandler)
         assert null != fedAuthInfo;
 
         attemptRefreshTokenLocked = true;
-        fedAuthToken = getFedAuthToken(fedAuthInfo);
+
+        SQLServerAccessTokenCallback callback = (SQLServerAccessTokenCallback) activeConnectionProperties
+                .get(SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString());
+
+        if (authenticationString.equals(SqlAuthentication.NotSpecified.toString()) && null != callback) {
+            fedAuthToken = callback.getAccessToken(fedAuthInfo.spn, fedAuthInfo.stsurl);
+        } else {
+            fedAuthToken = getFedAuthToken(fedAuthInfo);
+        }
+
         attemptRefreshTokenLocked = false;
 
         // fedAuthToken cannot be null.
@@ -5540,8 +5577,8 @@ void onFedAuthInfo(SqlFedAuthInfo fedAuthInfo, TDSTokenHandler tdsTokenHandler)
         fedAuthCommand.execute(tdsChannel.getWriter(), tdsChannel.getReader(fedAuthCommand));
     }
 
-    private SqlFedAuthToken getFedAuthToken(SqlFedAuthInfo fedAuthInfo) throws SQLServerException {
-        SqlFedAuthToken fedAuthToken = null;
+    private SqlAuthenticationToken getFedAuthToken(SqlFedAuthInfo fedAuthInfo) throws SQLServerException {
+        SqlAuthenticationToken fedAuthToken = null;
 
         // fedAuthInfo should not be null.
         assert null != fedAuthInfo;
@@ -5552,7 +5589,7 @@ private SqlFedAuthToken getFedAuthToken(SqlFedAuthInfo fedAuthInfo) throws SQLSe
         int sleepInterval = 100;
 
         if (!msalContextExists()
-                && !authenticationString.equalsIgnoreCase(SqlAuthentication.ActiveDirectoryInteractive.toString())) {
+                && !authenticationString.equalsIgnoreCase(SqlAuthentication.ActiveDirectoryIntegrated.toString())) {
             MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_MSALMissing"));
             throw new SQLServerException(form.format(new Object[] {authenticationString}), null, 0, null);
         }
@@ -5613,7 +5650,9 @@ private SqlFedAuthToken getFedAuthToken(SqlFedAuthInfo fedAuthInfo) throws SQLSe
                         byte[] accessTokenFromDLL = dllInfo.accessTokenBytes;
 
                         String accessToken = new String(accessTokenFromDLL, UTF_16LE);
-                        fedAuthToken = new SqlFedAuthToken(accessToken, dllInfo.expiresIn);
+                        Date now = new Date();
+                        now.setTime(now.getTime() + (dllInfo.expiresIn * 1000));
+                        fedAuthToken = new SqlAuthenticationToken(accessToken, now);
 
                         // Break out of the retry loop in successful case.
                         break;
@@ -5721,18 +5760,18 @@ private boolean msalContextExists() {
     /**
      * Send the access token to the server.
      */
-    private void sendFedAuthToken(FedAuthTokenCommand fedAuthCommand, SqlFedAuthToken fedAuthToken,
+    private void sendFedAuthToken(FedAuthTokenCommand fedAuthCommand, SqlAuthenticationToken fedAuthToken,
             TDSTokenHandler tdsTokenHandler) throws SQLServerException {
         assert null != fedAuthToken;
-        assert null != fedAuthToken.accessToken;
+        assert null != fedAuthToken.getAccessToken();
 
         if (connectionlogger.isLoggable(Level.FINER)) {
             connectionlogger.fine(toString() + " Sending federated authentication token.");
         }
 
         TDSWriter tdsWriter = fedAuthCommand.startRequest(TDS.PKT_FEDAUTH_TOKEN_MESSAGE);
 
-        byte[] accessToken = fedAuthToken.accessToken.getBytes(UTF_16LE);
+        byte[] accessToken = fedAuthToken.getAccessToken().getBytes(UTF_16LE);
 
         // Send total length (length of token plus 4 bytes for the token length field)
         // If we were sending a nonce, this would include that length as well

src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDataSource.java

@@ -1224,6 +1224,30 @@ public int getMsiTokenCacheTtl() {
         return 0;
     }
 
+    /**
+     * Sets the {@link SQLServerAccessTokenCallback} delegate.
+     *
+     * @param accessTokenCallback
+     *        Access token callback delegate.
+     */
+    @Override
+    public void setAccessTokenCallback(SQLServerAccessTokenCallback accessTokenCallback) {
+        setObjectProperty(connectionProps, SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString(),
+                accessTokenCallback);
+    }
+
+    /**
+     * Returns a {@link SQLServerAccessTokenCallback}, the access token callback delegate.
+     *
+     * @return Access token callback delegate.
+     */
+    @Override
+    public SQLServerAccessTokenCallback getAccessTokenCallback() {
+        return (SQLServerAccessTokenCallback) getObjectProperty(connectionProps,
+                SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.toString(),
+                SQLServerDriverObjectProperty.ACCESS_TOKEN_CALLBACK.getDefaultValue());
+    }
+
     /**
      * Sets a property string value.
      *

src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDriver.java

@@ -65,7 +65,6 @@ enum SqlAuthentication {
     SqlPassword,
     ActiveDirectoryPassword,
     ActiveDirectoryIntegrated,
-    ActiveDirectoryMSI,
     ActiveDirectoryManagedIdentity,
     ActiveDirectoryServicePrincipal,
     ActiveDirectoryInteractive,
@@ -395,7 +394,8 @@ static ApplicationIntent valueOfString(String value) throws SQLServerException {
 
 
 enum SQLServerDriverObjectProperty {
-    GSS_CREDENTIAL("gsscredential", null);
+    GSS_CREDENTIAL("gsscredential", null),
+    ACCESS_TOKEN_CALLBACK("accessTokenCallback", null);
 
     private final String name;
     private final String defaultValue;