From 574b1e1aefecd1162bccf064a26c8e76a9a2c5bf Mon Sep 17 00:00:00 2001
From: David Engel <v-daenge@microsoft.com>
Date: Fri, 11 Jun 2021 15:01:35 -0700
Subject: [PATCH] Send client process ID to server on Java 9+ (#1595)

---
 .../sqlserver/jdbc/SQLServerConnection.java   | 17 +++++-
 .../jdbc/connection/ClientProcessIdTest.java  | 55 +++++++++++++++++++
 2 files changed, 71 insertions(+), 1 deletion(-)
 create mode 100644 src/test/java/com/microsoft/sqlserver/jdbc/connection/ClientProcessIdTest.java

diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
index e58ac4a2f0..31744fe1f1 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
@@ -1216,6 +1216,21 @@ final SQLCollation getDatabaseCollation() {
             .getLogger("com.microsoft.sqlserver.jdbc.Connection");
     private static String loggingClassName = "com.microsoft.sqlserver.jdbc.SQLServerConnection:";
 
+    /** Client process ID sent during login */
+    private static int pid = 0;
+
+    static {
+        long pidLong = 0;
+        try {
+            pidLong = ProcessHandle.current().pid();
+        } catch (NoClassDefFoundError e) { // ProcessHandle is Java 9+
+            if (loggerExternal.isLoggable(Level.FINER) && Util.isActivityTraceOn()) {
+                loggerExternal.finer(loggingClassName + " NoClassDefFoundError for ProcessHandle. ProcessId will be 0.");
+            }
+        }
+        pid = (pidLong > Integer.MAX_VALUE) ? 0 : (int)pidLong;
+    }
+
     /**
      * There are three ways to get a failover partner connection string, from the failover map, the connecting server
      * returned the following variable only stores the serverReturned failver information.
@@ -5370,7 +5385,7 @@ final boolean complete(LogonCommand logonCommand, TDSReader tdsReader) throws SQ
         tdsWriter.writeInt(tdsVersion);
         tdsWriter.writeInt(requestedPacketSize);
         tdsWriter.writeBytes(interfaceLibVersionBytes); // writeBytes() is little endian
-        tdsWriter.writeInt(0); // Client process ID (0 = ??)
+        tdsWriter.writeInt(pid); // Client process ID
         tdsWriter.writeInt(0); // Primary server connection ID
 
         tdsWriter.writeByte((byte) (// OptionFlags1:
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/connection/ClientProcessIdTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/connection/ClientProcessIdTest.java
new file mode 100644
index 0000000000..8872c9b7c6
--- /dev/null
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/connection/ClientProcessIdTest.java
@@ -0,0 +1,55 @@
+/*
+ * Microsoft JDBC Driver for SQL Server Copyright(c) Microsoft Corporation All rights reserved. This program is made
+ * available under the terms of the MIT License. See the LICENSE file in the project root for more information.
+ */
+package com.microsoft.sqlserver.jdbc.connection;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+
+import org.junit.jupiter.api.Test;
+import org.junit.platform.runner.JUnitPlatform;
+import org.junit.runner.RunWith;
+
+import com.microsoft.sqlserver.jdbc.SQLServerDataSource;
+import com.microsoft.sqlserver.testframework.AbstractTest;
+
+/*
+ * This test is for validating that client process ID gets registered with the server when available to the driver.
+ */
+@RunWith(JUnitPlatform.class)
+public class ClientProcessIdTest extends AbstractTest {
+
+    private static int pid = 0;
+
+    static {
+        long pidLong = 0;
+        try {
+            pidLong = ProcessHandle.current().pid();
+        } catch (NoClassDefFoundError e) { // ProcessHandle is Java 9+
+        }
+        pid = (pidLong > Integer.MAX_VALUE) ? 0 : (int)pidLong;
+    }
+
+    @Test
+    public void testClientProcessId() throws SQLException {
+        SQLServerDataSource ds = new SQLServerDataSource();
+        ds.setURL(connectionString);
+        String sqlSelect = "select host_process_id from sys.dm_exec_sessions where session_id = @@SPID";
+
+        try (Connection con = ds.getConnection(); Statement stmt = con.createStatement()) {
+            try (ResultSet rs = stmt.executeQuery(sqlSelect)) {
+                if (rs.next()) {
+                    assertEquals(rs.getInt(1), pid);
+                } else {
+                    assertTrue(false, "Expected row of data was not found.");
+                }
+            }
+        }
+    }
+}
\ No newline at end of file