Feature | Added support for Data Classification v2 (#1338)

microsoft · Jun 24, 2020 · c984628 · c984628
1 parent 8e8a765
commit c984628
Show file tree

Hide file tree

Showing 7 changed files with 189 additions and 73 deletions.
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java b/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java
@@ -125,7 +125,7 @@ final class TDS {
     // 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 byte MAX_SUPPORTED_DATA_CLASSIFICATION_VERSION = 0x02;
 
     static final int AES_256_CBC = 1;
     static final int AEAD_AES_256_CBC_HMAC_SHA256 = 2;
@@ -6434,6 +6434,8 @@ final SQLServerConnection getConnection() {
     private boolean useColumnEncryption = false;
     private boolean serverSupportsColumnEncryption = false;
     private boolean serverSupportsDataClassification = false;
+    private byte serverSupportedDataClassificationVersion = TDS.DATA_CLASSIFICATION_NOT_ENABLED;
+
     private ColumnEncryptionVersion columnEncryptionVersion;
 
     private final byte valueBytes[] = new byte[256];
@@ -6461,6 +6463,7 @@ private static int nextReaderID() {
         serverSupportsColumnEncryption = con.getServerSupportsColumnEncryption();
         columnEncryptionVersion = con.getServerColumnEncryptionVersion();
         serverSupportsDataClassification = con.getServerSupportsDataClassification();
+        serverSupportedDataClassificationVersion = con.getServerSupportedDataClassificationVersion();
     }
 
     final boolean isColumnEncryptionSettingEnabled() {
@@ -6475,6 +6478,10 @@ final boolean getServerSupportsDataClassification() {
         return serverSupportsDataClassification;
     }
 
+    final byte getServerSupportedDataClassificationVersion() {
+        return serverSupportedDataClassificationVersion;
+    }
+
     final void throwInvalidTDS() throws SQLServerException {
         if (logger.isLoggable(Level.SEVERE))
             logger.severe(toString() + " got unexpected value in TDS response at offset:" + payloadOffset);

diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
@@ -680,11 +680,16 @@ ColumnEncryptionVersion getServerColumnEncryptionVersion() {
     }
 
     private boolean serverSupportsDataClassification = false;
+    private byte serverSupportedDataClassificationVersion = TDS.DATA_CLASSIFICATION_NOT_ENABLED;
 
     boolean getServerSupportsDataClassification() {
         return serverSupportsDataClassification;
     }
 
+    byte getServerSupportedDataClassificationVersion() {
+        return serverSupportedDataClassificationVersion;
+    }
+
     // Boolean that indicates whether LOB objects created by this connection should be loaded into memory
     private boolean delayLoadingLobs = SQLServerDriverBooleanProperty.DELAY_LOADING_LOBS.getDefaultValue();
 
@@ -4669,9 +4674,9 @@ private void onFeatureExtAck(byte featureId, byte[] data) throws SQLServerExcept
                     throw new SQLServerException(SQLServerException.getErrString("R_UnknownDataClsTokenNumber"), null);
                 }
 
-                byte supportedDataClassificationVersion = data[0];
-                if ((0 == supportedDataClassificationVersion)
-                        || (supportedDataClassificationVersion > TDS.MAX_SUPPORTED_DATA_CLASSIFICATION_VERSION)) {
+                serverSupportedDataClassificationVersion = data[0];
+                if ((0 == serverSupportedDataClassificationVersion)
+                        || (serverSupportedDataClassificationVersion > TDS.MAX_SUPPORTED_DATA_CLASSIFICATION_VERSION)) {
                     throw new SQLServerException(SQLServerException.getErrString("R_InvalidDataClsVersionNumber"),
                             null);
                 }

diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDriver.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerDriver.java
@@ -455,6 +455,36 @@ public String toString() {
 }
 
 
+enum SensitivityRank {
+    NOT_DEFINED(-1),
+    NONE(0),
+    LOW(10),
+    MEDIUM(20),
+    HIGH(30),
+    CRITICAL(40);
+
+    private static final SensitivityRank[] VALUES = values();
+    private int rank;
+
+    private SensitivityRank(int rank) {
+        this.rank = rank;
+    }
+
+    public int getValue() {
+        return rank;
+    }
+
+    static boolean isValid(int rank) throws SQLServerException {
+        for (SensitivityRank r : VALUES) {
+            if (r.getValue() == rank) {
+                return true;
+            }
+        }
+        return false;
+    }
+}
+
+
 /**
  * Provides methods to connect to a SQL Server database and to obtain information about the JDBC driver.
  */

diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/StreamColumns.java b/src/main/java/com/microsoft/sqlserver/jdbc/StreamColumns.java
@@ -259,55 +259,73 @@ private SensitivityClassification processDataClassification(TDSReader tdsReader)
         SensitivityClassification sensitivityClassification = null;
 
         // get the label count
-        int numLabels = tdsReader.readUnsignedShort();
-        List<Label> labels = new ArrayList<Label>(numLabels);
+        int sensitivityLabelCount = tdsReader.readUnsignedShort();
+        List<Label> sensitivityLabels = new ArrayList<Label>(sensitivityLabelCount);
 
-        for (int i = 0; i < numLabels; i++) {
-            labels.add(readSensitivityLabel(tdsReader));
+        for (int i = 0; i < sensitivityLabelCount; i++) {
+            sensitivityLabels.add(readSensitivityLabel(tdsReader));
         }
 
         // get the information type count
-        int numInformationTypes = tdsReader.readUnsignedShort();
+        int informationTypeCount = tdsReader.readUnsignedShort();
 
-        List<InformationType> informationTypes = new ArrayList<InformationType>(numInformationTypes);
-        for (int i = 0; i < numInformationTypes; i++) {
+        List<InformationType> informationTypes = new ArrayList<InformationType>(informationTypeCount);
+        for (int i = 0; i < informationTypeCount; i++) {
             informationTypes.add(readSensitivityInformationType(tdsReader));
         }
 
+        int sensitivityRank = SensitivityRank.NOT_DEFINED.getValue();
+        if (TDS.MAX_SUPPORTED_DATA_CLASSIFICATION_VERSION <= tdsReader.getServerSupportedDataClassificationVersion()) {
+            sensitivityRank = tdsReader.readInt();
+            if (!SensitivityRank.isValid(sensitivityRank)) {
+                tdsReader.throwInvalidTDS();
+            }
+        }
+
         // get the per column classification data (corresponds to order of output columns for query)
-        int numResultColumns = tdsReader.readUnsignedShort();
+        int numResultSetColumns = tdsReader.readUnsignedShort();
 
-        List<ColumnSensitivity> columnSensitivities = new ArrayList<ColumnSensitivity>(numResultColumns);
-        for (int columnNum = 0; columnNum < numResultColumns; columnNum++) {
+        List<ColumnSensitivity> columnSensitivities = new ArrayList<ColumnSensitivity>(numResultSetColumns);
+        for (int columnNum = 0; columnNum < numResultSetColumns; 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++) {
+            int numSensitivityProperties = tdsReader.readUnsignedShort();
+            List<SensitivityProperty> sensitivityProperties = new ArrayList<SensitivityProperty>(
+                    numSensitivityProperties);
+            for (int sourceNum = 0; sourceNum < numSensitivityProperties; sourceNum++) {
 
                 // get the label index and then lookup label to use for source
-                int labelIndex = tdsReader.readUnsignedShort();
+                int sensitivityLabelIndex = tdsReader.readUnsignedShort();
                 Label label = null;
-                if (labelIndex != Integer.MAX_VALUE) {
-                    if (labelIndex >= labels.size()) {
+                if (sensitivityLabelIndex != Integer.MAX_VALUE) {
+                    if (sensitivityLabelIndex >= sensitivityLabels.size()) {
                         tdsReader.throwInvalidTDS();
                     }
-                    label = labels.get(labelIndex);
+                    label = sensitivityLabels.get(sensitivityLabelIndex);
                 }
-
                 // 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);
                 }
+
+                if (TDS.MAX_SUPPORTED_DATA_CLASSIFICATION_VERSION <= tdsReader
+                        .getServerSupportedDataClassificationVersion()) {
+                    sensitivityRank = tdsReader.readInt();
+                    if (!SensitivityRank.isValid(sensitivityRank)) {
+                        tdsReader.throwInvalidTDS();
+                    }
+                }
+
                 // add sensitivity properties for the source
-                sensitivityProperties.add(new SensitivityProperty(label, informationType));
+                sensitivityProperties.add(new SensitivityProperty(label, informationType, sensitivityRank));
             }
             columnSensitivities.add(new ColumnSensitivity(sensitivityProperties));
         }
-        sensitivityClassification = new SensitivityClassification(labels, informationTypes, columnSensitivities);
+        sensitivityClassification = new SensitivityClassification(sensitivityLabels, informationTypes,
+                columnSensitivities);
         return sensitivityClassification;
     }
 

diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/dataclassification/SensitivityProperty.java b/src/main/java/com/microsoft/sqlserver/jdbc/dataclassification/SensitivityProperty.java
@@ -11,6 +11,7 @@
 public class SensitivityProperty {
     private Label label;
     private InformationType informationType;
+    private int sensitivityRank;
 
     /**
      * Constructs a SensitivityProperty
@@ -25,6 +26,22 @@ public SensitivityProperty(Label label, InformationType informationType) {
         this.informationType = informationType;
     }
 
+    /**
+     * Constructs a SensitivityProperty
+     * 
+     * @param label
+     *        Label as received from SQL Server for this SensitivityProperty
+     * @param informationType
+     *        InformationType as received from SQL Server for this SensitivityProperty
+     * @param sensitivityRank
+     *        sensitivity rank as received from SQL Server for this SensitivityProperty
+     */
+    public SensitivityProperty(Label label, InformationType informationType, int sensitivityRank) {
+        this.label = label;
+        this.informationType = informationType;
+        this.sensitivityRank = sensitivityRank;
+    }
+
     /**
      * Returns the label data for this <code>SensitivityProperty</code> Object
      * 
@@ -42,4 +59,8 @@ public Label getLabel() {
     public InformationType getInformationType() {
         return informationType;
     }
+
+    public int getSensitivityRank() {
+        return sensitivityRank;
+    }
 }
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java b/src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java
@@ -187,5 +187,6 @@ protected Object[][] getContents() {
             {"R_invalidEnclaveSessionFailed", "invalidate enclave session failed."},
             {"R_invalidEnclaveType", "Invalid enclave type {0}."},
             {"R_keystorePassword", "keystore password was incorrect"},
-            {"R_enclaveNotEnabled", "The statement triggers enclave computations"}};
+            {"R_enclaveNotEnabled", "The statement triggers enclave computations"},
+            {"R_dataClassificationNotSupported", "Data Classification is not supported on this server."}};
 }