Fixed session recovery for Azure SQL DB in redirect mode connected with AAD auth

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

@@ -7252,27 +7252,33 @@ final boolean complete(LogonCommand logonCommand, TDSReader tdsReader) throws SQ
 
         LogonProcessor logonProcessor = new LogonProcessor(authentication);
         TDSReader tdsReader;
+        sessionRecovery.setConnectionRecoveryPossible(false);
         do {
             tdsReader = logonCommand.startResponse();
-            sessionRecovery.setConnectionRecoveryPossible(false);
             TDSParser.parse(tdsReader, logonProcessor);
         } while (!logonProcessor.complete(logonCommand, tdsReader));
 
-        if (sessionRecovery.isReconnectRunning() && !sessionRecovery.isConnectionRecoveryPossible()) {
+        connectionReconveryCheck(sessionRecovery.isReconnectRunning(), sessionRecovery.isConnectionRecoveryPossible(),
+                routingInfo);
+
+        if (connectRetryCount > 0 && !sessionRecovery.isReconnectRunning()) {
+            sessionRecovery.getSessionStateTable().setOriginalCatalog(sCatalog);
+            sessionRecovery.getSessionStateTable().setOriginalCollation(databaseCollation);
+            sessionRecovery.getSessionStateTable().setOriginalLanguage(sLanguage);
+        }
+    }
+
+    private void connectionReconveryCheck(boolean isReconnectRunning, boolean isConnectionRecoveryPossible,
+                    ServerPortPlaceHolder routingDetails) throws SQLServerException {
+        if (isReconnectRunning && !isConnectionRecoveryPossible && routingDetails == null) {
             if (connectionlogger.isLoggable(Level.WARNING)) {
                 connectionlogger.warning(this.toString()
                         + "SessionRecovery feature extension ack was not sent by the server during reconnection.");
             }
             terminate(SQLServerException.DRIVER_ERROR_INVALID_TDS,
                     SQLServerException.getErrString("R_crClientNoRecoveryAckFromLogin"));
         }
-        if (connectRetryCount > 0 && !sessionRecovery.isReconnectRunning()) {
-            sessionRecovery.getSessionStateTable().setOriginalCatalog(sCatalog);
-            sessionRecovery.getSessionStateTable().setOriginalCollation(databaseCollation);
-            sessionRecovery.getSessionStateTable().setOriginalLanguage(sLanguage);
-        }
     }
-
     /* --------------- JDBC 3.0 ------------- */
 
     /**

src/test/java/com/microsoft/sqlserver/jdbc/SQLServerConnectionTest.java

@@ -9,10 +9,21 @@
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 import java.io.IOException;
 import java.io.Reader;
 import java.lang.management.ManagementFactory;
+import java.lang.reflect.Method;
 import java.sql.Clob;
 import java.sql.Connection;
 import java.sql.DriverManager;
@@ -29,6 +40,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
+import java.util.logging.Level;
 import java.util.logging.Logger;
 
 import javax.sql.ConnectionEvent;
@@ -60,6 +72,9 @@ public class SQLServerConnectionTest extends AbstractTest {
 
     String randomServer = RandomUtil.getIdentifier("Server");
 
+    SQLServerConnection mockConnection;
+    Logger mockLogger;
+
     @BeforeAll
     public static void setupTests() throws Exception {
         setConnection();
@@ -1528,6 +1543,49 @@ public void testManagedIdentityWithEncryptStrict() {
         } catch (SQLException e) {
             fail("Connection failed: " + e.getMessage());
         }
-    } 
+    }
+
+    public Method mockedConnectionRecoveryCheck() throws Exception {
+        mockConnection = spy(new SQLServerConnection("test"));
+        mockLogger = mock(Logger.class);
+        doReturn(true).when(mockLogger).isLoggable(Level.WARNING);
+        doNothing().when(mockConnection).terminate(anyInt(), anyString());
+
+        Method method = SQLServerConnection.class.getDeclaredMethod("connectionReconveryCheck", boolean.class,
+                boolean.class, ServerPortPlaceHolder.class);
+        method.setAccessible(true);
+        return method;
+    }
+
+    @Test
+    void testConnectionRecoveryCheckThrowsWhenAllConditionsMet() throws Exception {
+        Method method = mockedConnectionRecoveryCheck();
+        method.invoke(mockConnection, true, false, null);
+        verify(mockConnection, times(1)).terminate(eq(SQLServerException.DRIVER_ERROR_INVALID_TDS),
+                eq(SQLServerException.getErrString("R_crClientNoRecoveryAckFromLogin")));
+    }
+
+    @Test
+    void testConnectionRecoveryCheckDoesNotThrowWhenNotReconnectRunning() throws Exception {
+        Method method = mockedConnectionRecoveryCheck();
+        method.invoke(mockConnection, false, false, null);
+        verify(mockConnection, never()).terminate(anyInt(), anyString());
+    }
+
+    @Test
+    void testConnectionRecoveryCheckDoesNotThrowWhenRecoveryPossible() throws Exception {
+        Method method = mockedConnectionRecoveryCheck();
+        method.invoke(mockConnection, true, true, null);
+        verify(mockConnection, never()).terminate(anyInt(), anyString());
+    }
+
+    @Test
+    void testConnectionRecoveryCheckDoesNotThrowWhenRoutingDetailsNotNull() throws Exception {
+        Method method = mockedConnectionRecoveryCheck();
+        ServerPortPlaceHolder routingDetails = mock(ServerPortPlaceHolder.class);
+        method.setAccessible(true);
+        method.invoke(mockConnection, true, false, routingDetails);
+        verify(mockConnection, never()).terminate(anyInt(), anyString());
+    }
 
 }