Introduce support for "Data Classification Specifications" on fetched resultsets

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

@@ -73,6 +73,9 @@
 import javax.net.ssl.TrustManager;
 import javax.net.ssl.TrustManagerFactory;
 import javax.net.ssl.X509TrustManager;
+
+import com.microsoft.sqlserver.jdbc.dataclassification.SensitivityClassification;
+
 import java.nio.Buffer;
 
 final class TDS {
@@ -100,9 +103,11 @@ 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;
+    static final byte TDS_FEATURE_EXT_FEDAUTH = 0x02;
     static final int TDS_FEDAUTH_LIBRARY_SECURITYTOKEN = 0x01;
     static final int TDS_FEDAUTH_LIBRARY_ADAL = 0x02;
     static final int TDS_FEDAUTH_LIBRARY_RESERVED = 0x7F;
@@ -112,9 +117,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
-    static final int TDS_FEATURE_EXT_AE = 0x04;
-    static final int MAX_SUPPORTED_TCE_VERSION = 0x01; // max version
+    // 0x03 is for x_eFeatureExtensionId_Rcs
+    static final byte TDS_FEATURE_EXT_AE = 0x04;
+    static final byte 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 byte TDS_FEATURE_EXT_DATACLASSIFICATION = 0x09;
+    static final byte DATA_CLASSIFICATION_NOT_ENABLED = 0x00;
+    static final byte 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 +190,8 @@ static final String getTokenName(int tdsTokenType) {
                 return "TDS_DONEINPROC (0xFF)";
             case TDS_FEDAUTHINFO:
                 return "TDS_FEDAUTHINFO (0xEE)";
+            case TDS_FEATURE_EXT_DATACLASSIFICATION:
+                return "TDS_FEATURE_EXT_DATACLASSIFICATION (0x09)";
             default:
                 return "unknown token (0x" + Integer.toHexString(tdsTokenType).toUpperCase() + ")";
         }
@@ -6367,7 +6382,7 @@ final TDSCommand getCommand() {
     final SQLServerConnection getConnection() {
         return con;
     }
-
+    
     private TDSPacket currentPacket = new TDSPacket(0);
     private TDSPacket lastPacket = currentPacket;
     private int payloadOffset = 0;
@@ -6376,8 +6391,12 @@ 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];
+
+    protected SensitivityClassification sensitivityClassification;
+
     private static final AtomicInteger lastReaderID = new AtomicInteger(0);
 
     private static int nextReaderID() {
@@ -6404,6 +6423,7 @@ private static int nextReaderID() {
             useColumnEncryption = true;
         }
         serverSupportsColumnEncryption = con.getServerSupportsColumnEncryption();
+        serverSupportsDataClassification = con.getServerSupportsDataClassification();
     }
 
     final boolean isColumnEncryptionSettingEnabled() {
@@ -6414,6 +6434,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);
@@ -7105,7 +7129,7 @@ final void skip(int bytesToSkip) throws SQLServerException {
         }
     }
 
-    final void TryProcessFeatureExtAck(boolean featureExtAckReceived) throws SQLServerException {
+    final void tryProcessFeatureExtAck(boolean featureExtAckReceived) throws SQLServerException {
         // in case of redirection, do not check if TDS_FEATURE_EXTENSION_ACK is received or not.
         if (null != this.con.getRoutingInfo()) {
             return;
@@ -7114,6 +7138,10 @@ final void TryProcessFeatureExtAck(boolean featureExtAckReceived) throws SQLServ
         if (isColumnEncryptionSettingEnabled() && !featureExtAckReceived)
             throw new SQLServerException(this, SQLServerException.getErrString("R_AE_NotSupportedByServer"), null, 0, false);
     }
+
+    final void trySetSensitivityClassification(SensitivityClassification sensitivityClassification) {
+        this.sensitivityClassification = sensitivityClassification;
+    }
 }
 
 /**

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

@@ -560,6 +560,12 @@ boolean getServerSupportsColumnEncryption() {
         return serverSupportsColumnEncryption;
     }
 
+    private boolean serverSupportsDataClassification = false;
+
+    boolean getServerSupportsDataClassification() {
+        return serverSupportsDataClassification;
+    }
+
     static boolean isWindows;
     static Map<String, SQLServerColumnEncryptionKeyStoreProvider> globalSystemColumnEncryptionKeyStoreProviders = new HashMap<>();
     static {
@@ -3334,9 +3340,9 @@ int writeAEFeatureRequest(boolean write,
         int len = 6; // (1byte = featureID, 4bytes = featureData length, 1 bytes = Version)
 
         if (write) {
-            tdsWriter.writeByte((byte) TDS.TDS_FEATURE_EXT_AE); // FEATUREEXT_TCE
+            tdsWriter.writeByte(TDS.TDS_FEATURE_EXT_AE); // FEATUREEXT_TCE
             tdsWriter.writeInt(1);
-            tdsWriter.writeByte((byte) TDS.MAX_SUPPORTED_TCE_VERSION);
+            tdsWriter.writeByte(TDS.MAX_SUPPORTED_TCE_VERSION);
         }
         return len;
     }
@@ -3347,7 +3353,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);
 
@@ -3431,7 +3436,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(TDS.TDS_FEATURE_EXT_DATACLASSIFICATION);
+            tdsWriter.writeInt(1);
+            tdsWriter.writeByte(TDS.MAX_SUPPORTED_DATA_CLASSIFICATION_VERSION);
+        }
+        return len; // size of data written
+    }
+
     private final class LogonCommand extends UninterruptableTDSCommand {
         LogonCommand() {
             super("logon");
@@ -4056,7 +4074,7 @@ final void processFeatureExtAck(TDSReader tdsReader) throws SQLServerException {
         while (featureId != TDS.FEATURE_EXT_TERMINATOR);
     }
 
-    private void onFeatureExtAck(int featureId,
+    private void onFeatureExtAck(byte featureId,
             byte[] data) throws SQLServerException {
         if (null != routingInfo) {
             return;
@@ -4126,11 +4144,34 @@ 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_FEATURE_EXT_DATACLASSIFICATION: {
+                if (connectionlogger.isLoggable(Level.FINER)) {
+                    connectionlogger.fine(toString() + " Received feature extension acknowledgement for Data Classification.");
+                }
+
+                if (2 != 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_InvalidDataClsVersionNumber"), null);
+                }
 
+                byte enabled = data[1];
+                serverSupportsDataClassification = (enabled == 0) ? false : true;
+                break;
+            }
             default: {
                 // Unknown feature ack
                 if (connectionlogger.isLoggable(Level.SEVERE)) {
@@ -4417,8 +4458,11 @@ else if (serverMajorVersion >= 9) // Yukon (9.0) --> TDS 7.2 // Prelogin disconn
             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);
@@ -4598,6 +4642,8 @@ else if (serverMajorVersion >= 9) // Yukon (9.0) --> TDS 7.2 // Prelogin disconn
             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 DATA_CLASSIFICATION_INVALID_VERSION = 24;
+    static final int DATA_CLASSIFICATION_NOT_EXPECTED = 25;
+    static final int DATA_CLASSIFICATION_INVALID_LABEL_INDEX = 26;
+    static final int DATA_CLASSIFICATION_INVALID_INFORMATION_TYPE_INDEX = 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

@@ -31,6 +31,8 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.logging.Level;
 
+import com.microsoft.sqlserver.jdbc.dataclassification.SensitivityClassification;
+
 /**
  * Indicates the type of the row received from the server
  */
@@ -212,7 +214,17 @@ 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
      * 

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

@@ -8,6 +8,16 @@
 
 package com.microsoft.sqlserver.jdbc;
 
+import java.util.List;
+
+import com.microsoft.sqlserver.jdbc.dataclassification.ColumnSensitivity;
+import com.microsoft.sqlserver.jdbc.dataclassification.InformationType;
+import com.microsoft.sqlserver.jdbc.dataclassification.Label;
+import com.microsoft.sqlserver.jdbc.dataclassification.SensitivityClassification;
+import com.microsoft.sqlserver.jdbc.dataclassification.SensitivityProperty;
+
+import java.util.ArrayList;
+
 /**
  * StreamColumns stores the column meta data for a result set. StreamColumns parses the inbound TDS packet stream to determine column meta data.
  */
@@ -210,6 +220,95 @@ void setFromTDS(TDSReader tdsReader) throws SQLServerException {
                 this.columns[numColumns] = new Column(typeInfo, columnName, tableName, null);
             }
         }
+
+        // Data Classification
+        if (tdsReader.getServerSupportsDataClassification() && tdsReader.peekTokenType() == TDS.TDS_SQLDATACLASSIFICATION) {
+            // Read and parse
+            tdsReader.trySetSensitivityClassification(processDataClassification(tdsReader));
+        }
+    }
+
+    private String readByteString(TDSReader tdsReader) throws SQLServerException {
+        String value = "";
+        int byteLen = (int) tdsReader.readUnsignedByte();
+        value = tdsReader.readUnicodeString(byteLen);
+        return value;
+    }
+
+    private Label readSensitivityLabel(TDSReader tdsReader) throws SQLServerException {
+        String name = readByteString(tdsReader);
+        String id = readByteString(tdsReader);
+        return new Label(name, id);
+    }
+
+    private InformationType readSensitivityInformationType(TDSReader tdsReader) throws SQLServerException {
+        String name = readByteString(tdsReader);
+        String id = readByteString(tdsReader);
+        return new InformationType(name, id);
+    }
+
+    private SensitivityClassification processDataClassification(TDSReader tdsReader) throws SQLServerException {
+        if (!tdsReader.getServerSupportsDataClassification()) {
+            tdsReader.throwInvalidTDS();
+        }
+
+        int dataClassificationToken = tdsReader.readUnsignedByte();
+        assert dataClassificationToken == TDS.TDS_SQLDATACLASSIFICATION;
+
+        SensitivityClassification sensitivityClassification = null;
+
+        // get the label count
+        int numLabels = tdsReader.readUnsignedShort();
+        List<Label> labels = new ArrayList<Label>(numLabels);
+
+        for (int i = 0; i < numLabels; i++) {
+            labels.add(readSensitivityLabel(tdsReader));
+        }
+
+        // get the information type count
+        int numInformationTypes = tdsReader.readUnsignedShort();
+
+        List<InformationType> informationTypes = new ArrayList<InformationType>(numInformationTypes);
+        for (int i = 0; i < numInformationTypes; i++) {
+            informationTypes.add(readSensitivityInformationType(tdsReader));
+        }
+
+        // get the per column classification data (corresponds to order of output columns for query)
+        int numResultColumns = tdsReader.readUnsignedShort();
+
+        List<ColumnSensitivity> columnSensitivities = new ArrayList<ColumnSensitivity>(numResultColumns);
+        for (int columnNum = 0; columnNum < numResultColumns; columnNum++) {
+
+            // get sensitivity properties for all the different sources which were used in generating the column output
+            int numSources = tdsReader.readUnsignedShort();
+            List<SensitivityProperty> sensitivityProperties = new ArrayList<SensitivityProperty>(numSources);
+            for (int sourceNum = 0; sourceNum < numSources; sourceNum++) {
+
+                // get the label index and then lookup label to use for source
+                int labelIndex = tdsReader.readUnsignedShort();
+                Label label = null;
+                if (labelIndex != Integer.MAX_VALUE) {
+                    if (labelIndex >= labels.size()) {
+                        tdsReader.throwInvalidTDS();
+                    }
+                    label = labels.get(labelIndex);
+                }
+
+                // get the information type index and then lookup information type to use for source
+                int informationTypeIndex = tdsReader.readUnsignedShort();
+                InformationType informationType = null;
+                if (informationTypeIndex != Integer.MAX_VALUE) {
+                    if (informationTypeIndex >= informationTypes.size()) {
+                    }
+                    informationType = informationTypes.get(informationTypeIndex);
+                }
+                // add sensitivity properties for the source
+                sensitivityProperties.add(new SensitivityProperty(label, informationType));
+            }
+            columnSensitivities.add(new ColumnSensitivity(sensitivityProperties));
+        }
+        sensitivityClassification = new SensitivityClassification(labels, informationTypes, columnSensitivities);
+        return sensitivityClassification;
     }
 
     /**