From 17caf22234e632407499ad4e5a234c9da2de5aef Mon Sep 17 00:00:00 2001
From: Terry Chow <32403408+tkyc@users.noreply.github.com>
Date: Thu, 23 May 2024 11:44:39 -0700
Subject: [PATCH] [Backport] SocketTimeout should be unbounded by loginTimeout
 after a successful connection open (2355) (#2431)

* socketTimeout should reset to original value after a successful connection open (#2355)

* SocketTimeout should be unbounded by loginTimeout after a successful connection open
---
 CHANGELOG.md                                         |  1 +
 .../java/com/microsoft/sqlserver/jdbc/IOBuffer.java  |  4 ++++
 .../sqlserver/jdbc/SQLServerConnection.java          |  6 ++++++
 .../sqlserver/jdbc/connection/TimeoutTest.java       | 12 ++++++++++++
 4 files changed, 23 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6402d8942..16d66b8e4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/)
 - Clear prepared statement handle before reconnect [#2422](https://github.com/microsoft/mssql-jdbc/pull/2422)
 - RPC calls for CallableStatements will be executed directly [#2427](https://github.com/microsoft/mssql-jdbc/pull/2427)
 - Corrected authentication token object to accept expiration in milliseconds [#2428](https://github.com/microsoft/mssql-jdbc/pull/2428)
+- SocketTimeout should be unbounded by loginTimeout after a successful connection open [#2431](https://github.com/microsoft/mssql-jdbc/pull/2431)
 
 ## [12.6.1] Hotfix & Stable Release
 ### Fixed issues
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java b/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java
index 6e09bd756..2173b928f 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java
@@ -2366,6 +2366,10 @@ final int getNetworkTimeout() throws IOException {
     final void setNetworkTimeout(int timeout) throws IOException {
         tcpSocket.setSoTimeout(timeout);
     }
+
+    void resetTcpSocketTimeout() throws SocketException {
+        this.tcpSocket.setSoTimeout(con.getSocketTimeoutMilliseconds());
+    }
 }
 
 
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
index 96f216b37..1434cada3 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
@@ -3225,9 +3225,15 @@ else if (0 == requestedPacketSize)
 
             state = State.OPENED;
 
+            // Socket timeout is bounded by loginTimeout during the login phase.
+            // Reset socket timeout back to the original value.
+            tdsChannel.resetTcpSocketTimeout();
+
             if (connectionlogger.isLoggable(Level.FINER)) {
                 connectionlogger.finer(toString() + " End of connect");
             }
+        } catch (SocketException e) {
+            throw new SQLServerException(e.getMessage(), null);
         } finally {
             // once we exit the connect function, the connection can be only in one of two
             // states, Opened or Closed(if an exception occurred)
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/connection/TimeoutTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/connection/TimeoutTest.java
index 13064bf7f..129d13a47 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/connection/TimeoutTest.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/connection/TimeoutTest.java
@@ -281,6 +281,18 @@ public void testConnectRetryTimeout() {
                 "total time: " + totalTime + " interval: " + TimeUnit.SECONDS.toMillis(interval));
     }
 
+    @Test
+    public void testSocketTimeoutBoundedByLoginTimeoutReset() throws Exception {
+        try (Connection con = PrepUtil.getConnection(connectionString + ";socketTimeout=90000;loginTimeout=10;");
+                Statement stmt = con.createStatement()) {
+            // Login timeout (10s) is less than the 15s sec WAITFOR DELAY. Upon a login attempt, socketTimeout should be bounded
+            // by loginTimeout. After a successful login, when executing a query, socketTimeout should be reset to the
+            // original 90000ms timeout. The statement below should successfully execute as socketTimeout should not be bounded
+            // by loginTimeout, otherwise the test fails with a socket read timeout error.
+            stmt.execute("WAITFOR DELAY '00:00:15';");
+        }
+    }
+
     // Test for detecting Azure server for connection retries
     @Test
     public void testAzureEndpointRetry() {