Introduce support for "Data Classification Specifications" on fetched resultsets

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

@@ -0,0 +1,16 @@
+package com.microsoft.sqlserver.jdbc;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class ColumnSensitivity {
+    public List<SensitivityProperty> sensitivityProperties;
+
+	public ColumnSensitivity(List<SensitivityProperty> sensitivityProperties) {
+        this.sensitivityProperties = new ArrayList<SensitivityProperty>(sensitivityProperties);
+    }
+
+    public List<SensitivityProperty> getSensitivityProperties() {
+		return sensitivityProperties;
+	}
+}

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

@@ -73,6 +73,7 @@
 import javax.net.ssl.TrustManager;
 import javax.net.ssl.TrustManagerFactory;
 import javax.net.ssl.X509TrustManager;
+
 import java.nio.Buffer;
 
 final class TDS {
@@ -100,6 +101,8 @@ final class TDS {
     static final int TDS_DONEPROC = 0xFE;
     static final int TDS_DONEINPROC = 0xFF;
     static final int TDS_FEDAUTHINFO = 0xEE;
+    static final int TDS_SQLRESCOLSRCS = 0xa2;
+    static final int TDS_SQLDATACLASSIFICATION = 0xa3;
 
     // FedAuth
     static final int TDS_FEATURE_EXT_FEDAUTH = 0x02;
@@ -112,9 +115,17 @@ final class TDS {
     static final byte FEDAUTH_INFO_ID_SPN = 0x02; // FedAuthInfoData is the SPN to use for acquiring fed auth token
 
     // AE constants
+    // 0x03 is for x_eFeatureExtensionId_Rcs
     static final int TDS_FEATURE_EXT_AE = 0x04;
     static final int MAX_SUPPORTED_TCE_VERSION = 0x01; // max version
     static final int CUSTOM_CIPHER_ALGORITHM_ID = 0; // max version
+    // 0x06 is for x_eFeatureExtensionId_LoginToken 
+    // 0x07 is for x_eFeatureExtensionId_ClientSideTelemetry 
+    // Data Classification constants
+    static final int TDS_FEATUREEXT_DATACLASSIFICATION = 0x09;
+    static final int DATA_CLASSIFICATION_NOT_ENABLED = 0x00;
+    static final int MAX_SUPPORTED_DATA_CLASSIFICATION_VERSION = 0x01;
+
     static final int AES_256_CBC = 1;
     static final int AEAD_AES_256_CBC_HMAC_SHA256 = 2;
     static final int AE_METADATA = 0x08;
@@ -177,6 +188,8 @@ static final String getTokenName(int tdsTokenType) {
                 return "TDS_DONEINPROC (0xFF)";
             case TDS_FEDAUTHINFO:
                 return "TDS_FEDAUTHINFO (0xEE)";
+            case TDS_FEATUREEXT_DATACLASSIFICATION:
+                return "TDS_FEATUREEXT_DATACLASSIFICATION (0x09)";
             default:
                 return "unknown token (0x" + Integer.toHexString(tdsTokenType).toUpperCase() + ")";
         }
@@ -6367,7 +6380,7 @@ final TDSCommand getCommand() {
     final SQLServerConnection getConnection() {
         return con;
     }
-
+    
     private TDSPacket currentPacket = new TDSPacket(0);
     private TDSPacket lastPacket = currentPacket;
     private int payloadOffset = 0;
@@ -6376,9 +6389,13 @@ final SQLServerConnection getConnection() {
     private boolean isStreaming = true;
     private boolean useColumnEncryption = false;
     private boolean serverSupportsColumnEncryption = false;
+    private boolean serverSupportsDataClassification = false;
 
     private final byte valueBytes[] = new byte[256];
-    private static final AtomicInteger lastReaderID = new AtomicInteger(0);
+
+	protected SensitivityClassification sensitivityClassification;
+
+	private static final AtomicInteger lastReaderID = new AtomicInteger(0);
 
     private static int nextReaderID() {
         return lastReaderID.incrementAndGet();
@@ -6404,6 +6421,7 @@ private static int nextReaderID() {
             useColumnEncryption = true;
         }
         serverSupportsColumnEncryption = con.getServerSupportsColumnEncryption();
+        serverSupportsDataClassification = con.getServerSupportsDataClassification();
     }
 
     final boolean isColumnEncryptionSettingEnabled() {
@@ -6414,6 +6432,10 @@ final boolean getServerSupportsColumnEncryption() {
         return serverSupportsColumnEncryption;
     }
 
+    final boolean getServerSupportsDataClassification() {
+        return serverSupportsDataClassification;
+    }
+
     final void throwInvalidTDS() throws SQLServerException {
         if (logger.isLoggable(Level.SEVERE))
             logger.severe(toString() + " got unexpected value in TDS response at offset:" + payloadOffset);
@@ -7114,6 +7136,13 @@ final void TryProcessFeatureExtAck(boolean featureExtAckReceived) throws SQLServ
         if (isColumnEncryptionSettingEnabled() && !featureExtAckReceived)
             throw new SQLServerException(this, SQLServerException.getErrString("R_AE_NotSupportedByServer"), null, 0, false);
     }
+
+    boolean TrySetSensitivityClassification(SensitivityClassification sensitivityClassification) {
+        this.sensitivityClassification = sensitivityClassification;
+        return true;
+    }
+
+
 }
 
 /**

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

@@ -0,0 +1,20 @@
+package com.microsoft.sqlserver.jdbc;
+
+public class InformationType {
+    public String name;
+    public String id;
+
+    public InformationType(String name, String id)
+    {
+        this.name = name;
+        this.id = id;
+    }
+
+    public String getName() {
+		return name;
+	}
+
+	public String getId() {
+		return id;
+	}
+}

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

@@ -0,0 +1,19 @@
+package com.microsoft.sqlserver.jdbc;
+
+public class Label {
+    public String name;
+    public String id;
+
+	public Label(String name, String id) {
+        this.name = name;
+        this.id = id;
+    }
+
+    public String getName() {
+		return name;
+	}
+
+	public String getId() {
+		return id;
+	}
+}

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

@@ -559,6 +559,12 @@ boolean getServerSupportsColumnEncryption() {
         return serverSupportsColumnEncryption;
     }
 
+    private boolean serverSupportsDataClassification = false;
+
+    boolean getServerSupportsDataClassification() {
+        return serverSupportsDataClassification;
+    }
+
     static boolean isWindows;
     static Map<String, SQLServerColumnEncryptionKeyStoreProvider> globalSystemColumnEncryptionKeyStoreProviders = new HashMap<>();
     static {
@@ -3337,7 +3343,6 @@ int writeFedAuthFeatureRequest(boolean write,
                                                                                                                   * if false just calculates the
                                                                                                                   * length
                                                                                                                   */
-
         assert (fedAuthFeatureExtensionData.libraryType == TDS.TDS_FEDAUTH_LIBRARY_ADAL
                 || fedAuthFeatureExtensionData.libraryType == TDS.TDS_FEDAUTH_LIBRARY_SECURITYTOKEN);
 
@@ -3421,7 +3426,20 @@ int writeFedAuthFeatureRequest(boolean write,
         }
         return totalLen;
     }
+
+    int writeDataClassificationFeatureRequest (boolean write /* if false just calculates the length */,
+    		TDSWriter tdsWriter) throws SQLServerException {
+        int len = 6; // 1byte = featureID, 4bytes = featureData length, 1 bytes = Version
 
+        if (write) {
+            // Write Feature ID, length of the version# field and Sensitivity Classification Version#
+        	tdsWriter.writeByte((byte) TDS.TDS_FEATUREEXT_DATACLASSIFICATION);
+        	tdsWriter.writeInt(1);
+        	tdsWriter.writeByte((byte) TDS.MAX_SUPPORTED_DATA_CLASSIFICATION_VERSION);
+        }
+        return len; // size of data written
+    }
+
     private final class LogonCommand extends UninterruptableTDSCommand {
         LogonCommand() {
             super("logon");
@@ -4116,11 +4134,40 @@ private void onFeatureExtAck(int featureId,
                     throw new SQLServerException(SQLServerException.getErrString("R_InvalidAEVersionNumber"), null);
                 }
 
-                assert supportedTceVersion == TDS.MAX_SUPPORTED_TCE_VERSION; // Client support TCE version 1
                 serverSupportsColumnEncryption = true;
                 break;
             }
+            case TDS.TDS_FEATUREEXT_DATACLASSIFICATION: {
+            	if (connectionlogger.isLoggable(Level.FINER)) {
+                    connectionlogger.fine(toString() + " Received feature extension acknowledgement for Data Classification.");
+                }
+
+                if (1 > data.length) {
+                	if (connectionlogger.isLoggable(Level.SEVERE)) {
+                        connectionlogger.severe(toString() + " Unknown token for Data Classification.");
+                    }
+                    throw new SQLServerException(SQLServerException.getErrString("R_UnknownDataClsTokenNumber"), null);
+                }
+
+                byte supportedDataClassificationVersion = data[0];
+                if ((0 == supportedDataClassificationVersion) || (supportedDataClassificationVersion > TDS.MAX_SUPPORTED_DATA_CLASSIFICATION_VERSION)) {
+                	if (connectionlogger.isLoggable(Level.SEVERE)) {
+                        connectionlogger.severe(toString() + " Invalid version number for Data Classification");
+                    }
+                    throw new SQLServerException(SQLServerException.getErrString("R_InvalidDataClsTokenNumber"), null);
+                }
 
+                if (data.length != 2) {
+                	if (connectionlogger.isLoggable(Level.SEVERE)) {
+                        connectionlogger.severe(toString() + " Unknown token for Data Classification");
+                    }
+                    throw new SQLServerException(SQLServerException.getErrString("R_UnknownDataClsTokenNumber"), null);
+                }
+
+                byte enabled = data[1];
+                serverSupportsDataClassification = (enabled == 0) ? false : true;
+                break;
+            }
             default: {
                 // Unknown feature ack
                 if (connectionlogger.isLoggable(Level.SEVERE)) {
@@ -4406,9 +4453,12 @@ else if (serverMajorVersion >= 9) // Yukon (9.0) --> TDS 7.2 // Prelogin disconn
         if (federatedAuthenticationInfoRequested || federatedAuthenticationRequested) {
             len2 = len2 + writeFedAuthFeatureRequest(false, tdsWriter, fedAuthFeatureExtensionData);
         }
-
-        len2 = len2 + 1; // add 1 to length becaue of FeatureEx terminator
-
+        //Data Classification is always enabled (by default)
+
+        len2 += writeDataClassificationFeatureRequest(false, tdsWriter);
+
+        len2 = len2 + 1; // add 1 to length because of FeatureEx terminator
+
         // Length of entire Login 7 packet
         tdsWriter.writeInt(len2);
         tdsWriter.writeInt(tdsVersion);
@@ -4587,7 +4637,9 @@ else if (serverMajorVersion >= 9) // Yukon (9.0) --> TDS 7.2 // Prelogin disconn
         if (federatedAuthenticationInfoRequested || federatedAuthenticationRequested) {
             writeFedAuthFeatureRequest(true, tdsWriter, fedAuthFeatureExtensionData);
         }
-
+
+        writeDataClassificationFeatureRequest(true, tdsWriter);
+
         tdsWriter.writeByte((byte) TDS.FEATURE_EXT_TERMINATOR);
         tdsWriter.setDataLoggable(true);
 

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

@@ -72,6 +72,11 @@ public final class SQLServerException extends java.sql.SQLException {
     static final int DRIVER_ERROR_INTERMITTENT_TLS_FAILED = 7;
     static final int ERROR_SOCKET_TIMEOUT = 8;
     static final int ERROR_QUERY_TIMEOUT = 9;
+    static final int DataClassificationInvalidVersion = 24;
+    static final int DataClassificationNotExpected = 25;
+    static final int DataClassificationInvalidLabelIndex = 26;
+    static final int DataClassificationInvalidInformationTypeIndex = 27;
+
     private int driverErrorCode = DRIVER_ERROR_NONE;
 
     final int getDriverErrorCode() {

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

@@ -393,5 +393,7 @@ protected Object[][] getContents() {
 				{"R_invalidSSLProtocol", "SSL Protocol {0} label is not valid. Only TLS, TLSv1, TLSv1.1, and TLSv1.2 are supported."},
 				{"R_cancelQueryTimeoutPropertyDescription", "The number of seconds to wait to cancel sending a query timeout."},
 				{"R_invalidCancelQueryTimeout", "The cancel timeout value {0} is not valid."},
+				{"R_UnknownDataClsTokenNumber","Unknown token for Data Classification."},	// From Server
+				{"R_InvalidDataClsVersionNumber","Invalid version number {0} for Data Classification."},	// From Server
     };
 }

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

@@ -212,12 +212,22 @@ protected TDSReader getTDSReader() {
     }
 
     private final FetchBuffer fetchBuffer;
-
+
+    /**
+     * Exposes Data Classification information for the current ResultSet
+     * For SQL Servers that do not support Data Classification or results that do not fetch any classified columns, this data can be null
+     * 
+     * @return SensitivityClassification
+     */
+    public SensitivityClassification getSensitivityClassification() {
+    	return tdsReader.sensitivityClassification;
+    }
+
     /**
      * Make a new result set
      * 
      * @param stmtIn
-     *            the generating statement
+     * the generating statement
      */
     SQLServerResultSet(SQLServerStatement stmtIn) throws SQLServerException {
         int resultSetID = nextResultSetID();

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

@@ -0,0 +1,30 @@
+package com.microsoft.sqlserver.jdbc;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class SensitivityClassification {
+    public List<Label> labels;
+    public List<InformationType> informationTypes;
+    public List<ColumnSensitivity> columnSensitivities;
+
+    // Creating new ArrayList here assures that 'informationTypes' and 'labels' properties will not be null.
+    // The Count of the ColumnSensitivities property will be equal to the number of output columns for the query result set.
+	public SensitivityClassification(List<Label> labels, List<InformationType> informationTypes, List<ColumnSensitivity> columnSensitivity) {
+        this.labels = new ArrayList<Label>(labels);
+        this.informationTypes = new ArrayList<InformationType>(informationTypes);
+        this.columnSensitivities = new ArrayList<ColumnSensitivity>(columnSensitivity);
+    }
+
+    public List<Label> getLabels() {
+		return labels;
+	}
+
+	public List<InformationType> getInformationTypes() {
+		return informationTypes;
+	}
+
+	public List<ColumnSensitivity> getColumnSensitivities() {
+		return columnSensitivities;
+	}
+}

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

@@ -0,0 +1,19 @@
+package com.microsoft.sqlserver.jdbc;
+
+public class SensitivityProperty {
+    public Label label;
+    public InformationType informationType;
+
+	public SensitivityProperty(Label label, InformationType informationType) {
+        this.label = label;
+        this.informationType = informationType;
+    }
+
+    public Label getLabel() {
+		return label;
+	}
+
+	public InformationType getInformationType() {
+		return informationType;
+	}
+}