Merge pull request #17 from IronCoreLabs/leasedDataTest

Add test that ensures that keys wrapped with a leased key can be decr…
IronCoreLabs · Jul 13, 2020 · 532d8f1 · 532d8f1
2 parents 4745b68 + 43026d2
commit 532d8f1
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 6 deletions.
diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityKMSClient.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityKMSClient.java
@@ -93,7 +93,7 @@ public TenantSecurityKMSClient(String tspDomain, String apiKey, int requestThrea
      * @param requestThreadSize Number of threads to use for fixed-size web request thread pool
      * @param aesThreadSize     Number of threads to use for fixed-size AES operations threadpool
      * @param timeout           Request to TSP read and connect timeout in ms.
-     * 
+     *
      * @throws Exception If the provided domain is invalid.
      */
     public TenantSecurityKMSClient(String tspDomain, String apiKey, int requestThreadSize,
@@ -237,13 +237,12 @@ private static int getHeaderSize(byte[] bytes) {
      * @param bytes bytes to be checked
      */
     public static boolean isCiphertext(byte[] bytes) {
-        // Header size is currently always 0 for CMK encrypted docs. Expect at least one
-        // byte following the header that would have been encrypted.
-        // whenever header size is not 0, this should include a check that
-        // bytes.length > META_LENGTH + headerSize
+        // Header size is variable for CMK encrypted docs depending on whether
+        // the header is present. Expect at least one byte following the header
+        // that would have been encrypted.
         return bytes.length > DOCUMENT_HEADER_META_LENGTH
                 && bytes[0] == CURRENT_DOCUMENT_HEADER_VERSION && containsIroncoreMagic(bytes)
-                && getHeaderSize(bytes) == 0;
+                && getHeaderSize(bytes) >= 0;
     }
 
     /**

diff --git a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/DevIntegrationTest.java b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/DevIntegrationTest.java
@@ -9,6 +9,7 @@
 import java.util.concurrent.CompletionException;
 import java.util.stream.Collectors;
 import org.testng.annotations.Test;
+import java.util.Base64;
 
 @Test(groups = {"dev-integration"})
 public class DevIntegrationTest {
@@ -596,4 +597,26 @@ public void largeBatchTestNewApi() throws Exception {
         assertEquals(0, decryptedValues.getFailures().size());
         assertEquals(batchSize, decryptedValues.getDocuments().size());
     }
+
+    public void leasedDataTest() throws Exception {
+        DocumentMetadata metadata = getRoundtripMetadata(this.GCP_TENANT_ID);
+
+        String leasedDocumentEdek = "Cr8BCjA7nnuAiXpD0Jkjc6mOBgcSyxcjFYX813WQhhYg0oKnsDJTmeyAaLs3t9pzkR6mU9cQ7AMY3gQiDCEN6aQFtglBZ0DX7yp3CnUKcAokABW+8Gfu/FSC8WQTqxw528aQXwrpvY0MjlHurZJ6yHx9S/2zEkgAs0w57oTuIHzVmauLGDi/S9zCQH20dezcc/jtw/nqCDnAtAPSB9m17YvGOVpN5xO8960C86NA4AJCoVJ291YW9OkIKto48/YQ7AM=";
+        String leasedDocumentBytes = "A0lST04AOwocjKi8E65AAxBCqUjeSqQDc7veZVQehempBfsABBobChlJTlRFR1JBVElPTi1URVNULURFVjEtR0NQbZ+1yhYOoCNdtV+VVTMTUfAQm1FdqtGyjqeE7iYxfW9TKwTc2C0=";
+
+        Map<String, byte[]> documentMap = new HashMap<>();
+        documentMap.put("doc", Base64.getDecoder().decode(leasedDocumentBytes));
+        EncryptedDocument leasedDoc = new EncryptedDocument(documentMap, leasedDocumentEdek);
+
+        CompletableFuture<PlaintextDocument> roundtrip = getClient().thenCompose(client -> {
+            try {
+                return client.decrypt(leasedDoc, metadata);
+            } catch (Exception e) {
+                throw new CompletionException(e);
+            }
+        });
+
+        Map<String, byte[]> decryptedValuesMap = roundtrip.get().getDecryptedFields();
+        assertEqualBytes(decryptedValuesMap.get("doc"), "new daters".getBytes("UTF-8"));
+    }
 }