From 006ae586127c55f207bac0f0267a88d2400dc9d7 Mon Sep 17 00:00:00 2001 From: Colt Frederickson Date: Thu, 22 Jan 2026 12:05:42 -0700 Subject: [PATCH 1/8] Initial add of cached key decryptor --- .../kms/v1/CachedKeyDecryptor.java | 244 ++++++++++++++++++ .../kms/v1/DocumentDecryptor.java | 39 +++ .../kms/v1/TenantSecurityClient.java | 88 ++++++- .../kms/v1/CachedKeyDecryptorTest.java | 210 +++++++++++++++ 4 files changed, 580 insertions(+), 1 deletion(-) create mode 100644 src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java create mode 100644 src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentDecryptor.java create mode 100644 src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptorTest.java diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java new file mode 100644 index 0000000..8b261d0 --- /dev/null +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java @@ -0,0 +1,244 @@ +package com.ironcorelabs.tenantsecurity.kms.v1; + +import java.io.Closeable; +import java.io.InputStream; +import java.io.OutputStream; +import java.time.Duration; +import java.time.Instant; +import java.util.Arrays; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.stream.Collectors; +import com.ironcorelabs.tenantsecurity.kms.v1.exception.TscException; +import com.ironcorelabs.tenantsecurity.utils.CompletableFutures; + +/** + * Holds a cached DEK (Document Encryption Key) for repeated decrypt operations without making + * additional TSP unwrap calls. The DEK is securely zeroed when close() is called. + * + *

+ * This class is thread-safe and can be used concurrently for multiple decrypt operations. Once + * closed, all decrypt operations will fail. + * + *

+ * Expiration: This decryptor automatically expires after a short time period. Caching a DEK + * for long-term use is not supported as it would undermine the security benefits of key wrapping. + * The decryptor is intended for short-lived batch operations where multiple documents sharing the + * same EDEK need to be decrypted in quick succession. Use {@link #isExpired()} to check expiration + * status. + * + *

+ * Usage with loan pattern (recommended): + * + * 

+ * client.withCachedDecryptor(edek, metadata, decryptor ->
+ *     decryptor.decrypt(encDoc1, metadata)
+ *         .thenCompose(doc1 -> decryptor.decrypt(encDoc2, metadata)))
+ *
+ * + *

+ * Usage with try-with-resources: + * + * 

+ * try (CachedKeyDecryptor decryptor = client.createCachedDecryptor(edek, metadata).get()) {
+ *   PlaintextDocument doc1 = decryptor.decrypt(encDoc1, metadata).get();
+ *   PlaintextDocument doc2 = decryptor.decrypt(encDoc2, metadata).get();
+ * } // DEK is automatically zeroed
+ *
+ * + * @author IronCore Labs + */ +public final class CachedKeyDecryptor implements DocumentDecryptor, Closeable { + + // Maximum time the decryptor can be used before it expires + private static final Duration TIMEOUT = Duration.ofMinutes(1); + + // The cached DEK bytes - zeroed on close() + private final byte[] dek; + + // The EDEK that was used to derive the DEK - used for validation + private final String edek; + + // Executor for async field decryption operations + private final ExecutorService encryptionExecutor; + + // Flag to track if close() has been called + private final AtomicBoolean closed = new AtomicBoolean(false); + + // When this decryptor was created - used for timeout enforcement + private final Instant createdAt; + + /** + * Package-private constructor. Use TenantSecurityClient.createCachedDecryptor() to create + * instances. + * + * @param dek The unwrapped document encryption key bytes (will be copied) + * @param edek The encrypted document encryption key string + * @param encryptionExecutor Executor for async decryption operations + */ + CachedKeyDecryptor(byte[] dek, String edek, ExecutorService encryptionExecutor) { + if (dek == null || dek.length != 32) { + throw new IllegalArgumentException("DEK must be exactly 32 bytes"); + } + if (edek == null || edek.isEmpty()) { + throw new IllegalArgumentException("EDEK must not be null or empty"); + } + if (encryptionExecutor == null) { + throw new IllegalArgumentException("encryptionExecutor must not be null"); + } + // Copy DEK to prevent external modification + this.dek = Arrays.copyOf(dek, dek.length); + this.edek = edek; + this.encryptionExecutor = encryptionExecutor; + this.createdAt = Instant.now(); + } + + /** + * Get the EDEK associated with this cached decryptor. Useful for verifying which documents can be + * decrypted with this instance. + * + * @return The EDEK string + */ + public String getEdek() { + return edek; + } + + /** + * Check if this decryptor has been closed. + * + * @return true if close() has been called + */ + public boolean isClosed() { + return closed.get(); + } + + /** + * Check if this decryptor has expired due to timeout. + * + * @return true if the timeout has elapsed since creation + */ + public boolean isExpired() { + return Duration.between(createdAt, Instant.now()).compareTo(TIMEOUT) > 0; + } + + /** + * Decrypt the provided EncryptedDocument using the cached DEK. + * + *

+ * The document's EDEK must match the EDEK used to create this decryptor, otherwise an error is + * returned. + * + * @param encryptedDocument Document to decrypt + * @param metadata Metadata about the document being decrypted (used for audit/logging) + * @return CompletableFuture resolving to PlaintextDocument + */ + @Override + public CompletableFuture decrypt(EncryptedDocument encryptedDocument, + DocumentMetadata metadata) { + // Check if closed or expired + if (closed.get()) { + return CompletableFuture.failedFuture(new TscException( + TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "CachedKeyDecryptor has been closed")); + } + if (isExpired()) { + return CompletableFuture.failedFuture(new TscException( + TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "CachedKeyDecryptor has expired")); + } + + // Validate EDEK matches + if (!edek.equals(encryptedDocument.getEdek())) { + return CompletableFuture + .failedFuture(new TscException(TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, + "EncryptedDocument EDEK does not match the cached EDEK. " + + "This decryptor can only decrypt documents with matching EDEKs.")); + } + + return decryptFields(encryptedDocument.getEncryptedFields(), encryptedDocument.getEdek()); + } + + /** + * Decrypt a stream using the cached DEK. + * + *

+ * The provided EDEK must match the EDEK used to create this decryptor, otherwise an error is + * returned. + * + * @param edek Encrypted document encryption key - must match this decryptor's EDEK + * @param input A stream representing the encrypted document + * @param output An output stream to write the decrypted document to + * @param metadata Metadata about the document being decrypted + * @return Future which will complete when input has been decrypted + */ + @Override + public CompletableFuture decryptStream(String edek, InputStream input, OutputStream output, + DocumentMetadata metadata) { + // Check if closed or expired + if (closed.get()) { + return CompletableFuture.failedFuture(new TscException( + TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "CachedKeyDecryptor has been closed")); + } + if (isExpired()) { + return CompletableFuture.failedFuture(new TscException( + TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "CachedKeyDecryptor has expired")); + } + + // Validate EDEK matches + if (!this.edek.equals(edek)) { + return CompletableFuture + .failedFuture(new TscException(TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, + "Provided EDEK does not match the cached EDEK. " + + "This decryptor can only decrypt documents with matching EDEKs.")); + } + + return CompletableFuture.supplyAsync( + () -> CryptoUtils.decryptStreamInternal(dek, input, output).join(), encryptionExecutor); + } + + /** + * Decrypt all fields in the document using the cached DEK. Pattern follows + * TenantSecurityClient.decryptFields(). + */ + private CompletableFuture decryptFields(Map document, + String documentEdek) { + // Check closed/expired state again before starting decryption + if (closed.get()) { + return CompletableFuture.failedFuture(new TscException( + TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "CachedKeyDecryptor has been closed")); + } + if (isExpired()) { + return CompletableFuture.failedFuture(new TscException( + TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "CachedKeyDecryptor has expired")); + } + + // Parallel decrypt each field + Map> decryptOps = document.entrySet().stream() + .collect(Collectors.toMap(Map.Entry::getKey, + entry -> CompletableFuture.supplyAsync( + () -> CryptoUtils.decryptDocument(entry.getValue(), dek).join(), + encryptionExecutor))); + + // Join all futures and build result + return CompletableFutures.tryCatchNonFatal(() -> { + Map decryptedBytes = decryptOps.entrySet().stream() + .collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().join())); + return new PlaintextDocument(decryptedBytes, documentEdek); + }); + } + + /** + * Securely zero the DEK bytes and mark this decryptor as closed. After calling close(), all + * decrypt operations will fail. + * + *

+ * This method is idempotent - calling it multiple times has no additional effect. + */ + @Override + public void close() { + if (closed.compareAndSet(false, true)) { + // Zero out the DEK bytes for security + Arrays.fill(dek, (byte) 0); + } + } +} diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentDecryptor.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentDecryptor.java new file mode 100644 index 0000000..c77c34d --- /dev/null +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentDecryptor.java @@ -0,0 +1,39 @@ +package com.ironcorelabs.tenantsecurity.kms.v1; + +import java.io.InputStream; +import java.io.OutputStream; +import java.util.concurrent.CompletableFuture; + +/** + * Interface for document decryption capabilities. Implemented by both TenantSecurityClient (for + * standard decrypt operations that unwrap the EDEK each time) and CachedKeyDecryptor (for repeated + * decrypts using a cached DEK). + * + * @author IronCore Labs + */ +public interface DocumentDecryptor { + + /** + * Decrypt the provided EncryptedDocument and return the decrypted fields. + * + * @param encryptedDocument Document to decrypt which includes encrypted bytes as well as EDEK. + * @param metadata Metadata about the document being decrypted. + * @return CompletableFuture resolving to PlaintextDocument with decrypted field bytes. + */ + CompletableFuture decrypt(EncryptedDocument encryptedDocument, + DocumentMetadata metadata); + + /** + * Decrypt a stream using the provided EDEK. + * + * @param edek Encrypted document encryption key. + * @param input A stream representing the encrypted document. + * @param output An output stream to write the decrypted document to. Note that this output should + * not be used until after the future exits successfully because the GCM tag is not fully + * verified until that time. + * @param metadata Metadata about the document being decrypted. + * @return Future which will complete when input has been decrypted. + */ + CompletableFuture decryptStream(String edek, InputStream input, OutputStream output, + DocumentMetadata metadata); +} diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java index 0074447..aa5a14c 100644 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java @@ -14,6 +14,7 @@ import java.util.concurrent.ConcurrentMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.function.Function; import java.util.stream.Collectors; import java.util.stream.Stream; import com.ironcorelabs.tenantsecurity.kms.v1.exception.TenantSecurityException; @@ -27,7 +28,7 @@ * * @author IronCore Labs */ -public final class TenantSecurityClient implements Closeable { +public final class TenantSecurityClient implements Closeable, DocumentDecryptor { private final SecureRandom secureRandom; // Use fixed size thread pool for CPU bound operations (crypto ops). Defaults to @@ -434,6 +435,7 @@ public CompletableFuture encryptStream(InputStream input, Out * @param metadata Metadata about the document being encrypted. * @return Future which will complete when input has been decrypted. */ + @Override public CompletableFuture decryptStream(String edek, InputStream input, OutputStream output, DocumentMetadata metadata) { return this.encryptionService.unwrapKey(edek, metadata).thenApplyAsync( @@ -546,6 +548,7 @@ public CompletableFuture> encryptExistingBatch( * @param metadata Metadata about the document being decrypted. * @return PlaintextDocument which contains each documents decrypted field bytes. */ + @Override public CompletableFuture decrypt(EncryptedDocument encryptedDocument, DocumentMetadata metadata) { return this.encryptionService.unwrapKey(encryptedDocument.getEdek(), metadata).thenComposeAsync( @@ -553,6 +556,89 @@ public CompletableFuture decrypt(EncryptedDocument encryptedD decryptedDocumentAESKey, encryptedDocument.getEdek())); } + /** + * Create a CachedKeyDecryptor for repeated decrypt operations using the same DEK. This unwraps + * the EDEK once and caches the resulting DEK for subsequent decrypts. + * + *

+ * Use this when you need to decrypt multiple documents that share the same EDEK, to avoid + * repeated TSP unwrap calls. + * + *

+ * The returned decryptor implements AutoCloseable and should be used with try-with-resources to + * ensure the DEK is securely zeroed when done: + * + * 

+   * try (CachedKeyDecryptor decryptor = client.createCachedDecryptor(edek, metadata).get()) {
+   *   PlaintextDocument doc1 = decryptor.decrypt(encDoc1, metadata).get();
+   *   PlaintextDocument doc2 = decryptor.decrypt(encDoc2, metadata).get();
+   * }
+   *
+ * + * @param edek The encrypted document encryption key to unwrap + * @param metadata Metadata for the unwrap operation + * @return CompletableFuture resolving to a CachedKeyDecryptor + */ + public CompletableFuture createCachedDecryptor(String edek, + DocumentMetadata metadata) { + return this.encryptionService.unwrapKey(edek, metadata) + .thenApply(dekBytes -> new CachedKeyDecryptor(dekBytes, edek, this.encryptionExecutor)); + } + + /** + * Create a CachedKeyDecryptor from an existing EncryptedDocument. Convenience method that + * extracts the EDEK from the document. + * + * @param encryptedDocument The encrypted document whose EDEK should be unwrapped + * @param metadata Metadata for the unwrap operation + * @return CompletableFuture resolving to a CachedKeyDecryptor + */ + public CompletableFuture createCachedDecryptor( + EncryptedDocument encryptedDocument, DocumentMetadata metadata) { + return createCachedDecryptor(encryptedDocument.getEdek(), metadata); + } + + /** + * Execute an operation using a CachedKeyDecryptor with automatic lifecycle management. The + * decryptor is automatically closed (and DEK zeroed) when the operation completes, whether + * successfully or with an error. + * + *

+ * This is the recommended pattern for using cached decryptors with CompletableFuture composition: + * + * 

+   * client.withCachedDecryptor(edek, metadata, decryptor ->
+   *     decryptor.decrypt(encDoc1, metadata)
+   *         .thenCompose(doc1 -> decryptor.decrypt(encDoc2, metadata)))
+   *
+ * + * @param The type returned by the operation + * @param edek The encrypted document encryption key to unwrap + * @param metadata Metadata for the unwrap operation + * @param operation Function that takes the decryptor and returns a CompletableFuture + * @return CompletableFuture resolving to the operation's result + */ + public CompletableFuture withCachedDecryptor(String edek, DocumentMetadata metadata, + Function> operation) { + return createCachedDecryptor(edek, metadata).thenCompose( + decryptor -> operation.apply(decryptor).whenComplete((result, error) -> decryptor.close())); + } + + /** + * Execute an operation using a CachedKeyDecryptor with automatic lifecycle management. + * Convenience method that extracts the EDEK from the document. + * + * @param The type returned by the operation + * @param encryptedDocument The encrypted document whose EDEK should be unwrapped + * @param metadata Metadata for the unwrap operation + * @param operation Function that takes the decryptor and returns a CompletableFuture + * @return CompletableFuture resolving to the operation's result + */ + public CompletableFuture withCachedDecryptor(EncryptedDocument encryptedDocument, + DocumentMetadata metadata, Function> operation) { + return withCachedDecryptor(encryptedDocument.getEdek(), metadata, operation); + } + /** * Re-key a document's encrypted document key (EDEK) using a new KMS config. Decrypts the EDEK * then re-encrypts it using the specified tenant's current primary KMS config. The DEK is then diff --git a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptorTest.java b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptorTest.java new file mode 100644 index 0000000..33c24b4 --- /dev/null +++ b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptorTest.java @@ -0,0 +1,210 @@ +package com.ironcorelabs.tenantsecurity.kms.v1; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.fail; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.util.Arrays; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; +import com.ironcorelabs.tenantsecurity.kms.v1.exception.TscException; + +@Test(groups = {"unit"}) +public class CachedKeyDecryptorTest { + + private ExecutorService executor; + private static final String TEST_EDEK = "test-edek-base64-string"; + private static final String DIFFERENT_EDEK = "different-edek-base64-string"; + private DocumentMetadata metadata = + new DocumentMetadata("tenantId", "requestingUserOrServiceId", "dataLabel"); + + @BeforeClass + public void setup() { + executor = Executors.newFixedThreadPool(2); + } + + @AfterClass + public void teardown() { + if (executor != null) { + executor.shutdown(); + } + } + + private byte[] createValidDek() { + byte[] dek = new byte[32]; + Arrays.fill(dek, (byte) 0x42); + return dek; + } + + // Constructor validation tests + + public void constructorRejectNullDek() { + try { + new CachedKeyDecryptor(null, TEST_EDEK, executor); + fail("Should have thrown IllegalArgumentException"); + } catch (IllegalArgumentException e) { + assertTrue(e.getMessage().contains("DEK must be exactly 32 bytes")); + } + } + + public void constructorRejectWrongSizeDek() { + byte[] shortDek = new byte[16]; + try { + new CachedKeyDecryptor(shortDek, TEST_EDEK, executor); + fail("Should have thrown IllegalArgumentException"); + } catch (IllegalArgumentException e) { + assertTrue(e.getMessage().contains("DEK must be exactly 32 bytes")); + } + } + + public void constructorRejectNullEdek() { + try { + new CachedKeyDecryptor(createValidDek(), null, executor); + fail("Should have thrown IllegalArgumentException"); + } catch (IllegalArgumentException e) { + assertTrue(e.getMessage().contains("EDEK must not be null or empty")); + } + } + + public void constructorRejectEmptyEdek() { + try { + new CachedKeyDecryptor(createValidDek(), "", executor); + fail("Should have thrown IllegalArgumentException"); + } catch (IllegalArgumentException e) { + assertTrue(e.getMessage().contains("EDEK must not be null or empty")); + } + } + + public void constructorRejectNullExecutor() { + try { + new CachedKeyDecryptor(createValidDek(), TEST_EDEK, null); + fail("Should have thrown IllegalArgumentException"); + } catch (IllegalArgumentException e) { + assertTrue(e.getMessage().contains("encryptionExecutor must not be null")); + } + } + + // Getter tests + + public void getEdekReturnsCorrectValue() { + CachedKeyDecryptor decryptor = new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor); + assertEquals(decryptor.getEdek(), TEST_EDEK); + decryptor.close(); + } + + public void isClosedReturnsFalseInitially() { + CachedKeyDecryptor decryptor = new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor); + assertFalse(decryptor.isClosed()); + decryptor.close(); + } + + public void isClosedReturnsTrueAfterClose() { + CachedKeyDecryptor decryptor = new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor); + decryptor.close(); + assertTrue(decryptor.isClosed()); + } + + // Close tests + + public void closeIsIdempotent() { + CachedKeyDecryptor decryptor = new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor); + decryptor.close(); + assertTrue(decryptor.isClosed()); + // Should not throw + decryptor.close(); + decryptor.close(); + assertTrue(decryptor.isClosed()); + } + + // Decrypt validation tests + + public void decryptFailsWhenClosed() { + CachedKeyDecryptor decryptor = new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor); + decryptor.close(); + + EncryptedDocument encDoc = new EncryptedDocument(java.util.Collections.emptyMap(), TEST_EDEK); + + try { + decryptor.decrypt(encDoc, metadata).join(); + fail("Should have thrown CompletionException"); + } catch (CompletionException e) { + assertTrue(e.getCause() instanceof TscException); + assertTrue(e.getCause().getMessage().contains("CachedKeyDecryptor has been closed")); + } + } + + public void decryptFailsWhenEdekMismatch() { + CachedKeyDecryptor decryptor = new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor); + + EncryptedDocument encDoc = + new EncryptedDocument(java.util.Collections.emptyMap(), DIFFERENT_EDEK); + + try { + decryptor.decrypt(encDoc, metadata).join(); + fail("Should have thrown CompletionException"); + } catch (CompletionException e) { + assertTrue(e.getCause() instanceof TscException); + assertTrue(e.getCause().getMessage().contains("EDEK does not match")); + } finally { + decryptor.close(); + } + } + + // DecryptStream validation tests + + public void decryptStreamFailsWhenClosed() { + CachedKeyDecryptor decryptor = new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor); + decryptor.close(); + + ByteArrayInputStream input = new ByteArrayInputStream(new byte[0]); + ByteArrayOutputStream output = new ByteArrayOutputStream(); + + try { + decryptor.decryptStream(TEST_EDEK, input, output, metadata).join(); + fail("Should have thrown CompletionException"); + } catch (CompletionException e) { + assertTrue(e.getCause() instanceof TscException); + assertTrue(e.getCause().getMessage().contains("CachedKeyDecryptor has been closed")); + } + } + + public void decryptStreamFailsWhenEdekMismatch() { + CachedKeyDecryptor decryptor = new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor); + + ByteArrayInputStream input = new ByteArrayInputStream(new byte[0]); + ByteArrayOutputStream output = new ByteArrayOutputStream(); + + try { + decryptor.decryptStream(DIFFERENT_EDEK, input, output, metadata).join(); + fail("Should have thrown CompletionException"); + } catch (CompletionException e) { + assertTrue(e.getCause() instanceof TscException); + assertTrue(e.getCause().getMessage().contains("EDEK does not match")); + } finally { + decryptor.close(); + } + } + + // DEK copying test + + public void constructorCopiesDekToPreventExternalModification() { + byte[] originalDek = createValidDek(); + CachedKeyDecryptor decryptor = new CachedKeyDecryptor(originalDek, TEST_EDEK, executor); + + // Modify the original array + Arrays.fill(originalDek, (byte) 0x00); + + // The decryptor should still have the original values (0x42) + // We can't directly test this without reflection, but we verify the constructor + // doesn't throw when we create the decryptor, showing it made a copy + assertFalse(decryptor.isClosed()); + decryptor.close(); + } +} From b7d97f8c66382b40a708d9a967ed3bb425807c54 Mon Sep 17 00:00:00 2001 From: Colt Frederickson Date: Thu, 22 Jan 2026 12:15:18 -0700 Subject: [PATCH 2/8] Self review --- .../kms/v1/CachedKeyDecryptor.java | 22 +++++-------------- .../kms/v1/DocumentDecryptor.java | 2 -- .../kms/v1/CachedKeyDecryptorTest.java | 17 +++++++++----- 3 files changed, 17 insertions(+), 24 deletions(-) diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java index 8b261d0..22c8751 100644 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java @@ -30,25 +30,13 @@ * status. * *

- * Usage with loan pattern (recommended): - * - * 

- * client.withCachedDecryptor(edek, metadata, decryptor ->
- *     decryptor.decrypt(encDoc1, metadata)
- *         .thenCompose(doc1 -> decryptor.decrypt(encDoc2, metadata)))
- *
- * - *

- * Usage with try-with-resources: - * - * 

- * try (CachedKeyDecryptor decryptor = client.createCachedDecryptor(edek, metadata).get()) {
- *   PlaintextDocument doc1 = decryptor.decrypt(encDoc1, metadata).get();
- *   PlaintextDocument doc2 = decryptor.decrypt(encDoc2, metadata).get();
- * } // DEK is automatically zeroed
- *
+ * Instances are created via {@link TenantSecurityClient#createCachedDecryptor} or + * {@link TenantSecurityClient#withCachedDecryptor}. See those methods for usage examples. * * @author IronCore Labs + * @see TenantSecurityClient#createCachedDecryptor(String, DocumentMetadata) + * @see TenantSecurityClient#withCachedDecryptor(String, DocumentMetadata, + * java.util.function.Function) */ public final class CachedKeyDecryptor implements DocumentDecryptor, Closeable { diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentDecryptor.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentDecryptor.java index c77c34d..8883a74 100644 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentDecryptor.java +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentDecryptor.java @@ -8,8 +8,6 @@ * Interface for document decryption capabilities. Implemented by both TenantSecurityClient (for * standard decrypt operations that unwrap the EDEK each time) and CachedKeyDecryptor (for repeated * decrypts using a cached DEK). - * - * @author IronCore Labs */ public interface DocumentDecryptor { diff --git a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptorTest.java b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptorTest.java index 33c24b4..4cd6fdd 100644 --- a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptorTest.java +++ b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptorTest.java @@ -7,6 +7,7 @@ import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; +import java.lang.reflect.Field; import java.util.Arrays; import java.util.concurrent.CompletionException; import java.util.concurrent.ExecutorService; @@ -194,17 +195,23 @@ public void decryptStreamFailsWhenEdekMismatch() { // DEK copying test - public void constructorCopiesDekToPreventExternalModification() { + public void constructorCopiesDekToPreventExternalModification() throws Exception { byte[] originalDek = createValidDek(); CachedKeyDecryptor decryptor = new CachedKeyDecryptor(originalDek, TEST_EDEK, executor); // Modify the original array Arrays.fill(originalDek, (byte) 0x00); - // The decryptor should still have the original values (0x42) - // We can't directly test this without reflection, but we verify the constructor - // doesn't throw when we create the decryptor, showing it made a copy - assertFalse(decryptor.isClosed()); + // Use reflection to verify internal DEK still has original values + Field dekField = CachedKeyDecryptor.class.getDeclaredField("dek"); + dekField.setAccessible(true); + byte[] internalDek = (byte[]) dekField.get(decryptor); + + // Internal DEK should still be 0x42, not 0x00 + for (byte b : internalDek) { + assertEquals(b, (byte) 0x42, "Internal DEK should not be affected by external modification"); + } + decryptor.close(); } } From 2f1924c5c70c0bda69cb737c3faafeaa1baad46b Mon Sep 17 00:00:00 2001 From: Colt Frederickson Date: Thu, 22 Jan 2026 12:20:22 -0700 Subject: [PATCH 3/8] Self review --- .../ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java | 1 - 1 file changed, 1 deletion(-) diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java index 22c8751..fbd188c 100644 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java @@ -33,7 +33,6 @@ * Instances are created via {@link TenantSecurityClient#createCachedDecryptor} or * {@link TenantSecurityClient#withCachedDecryptor}. See those methods for usage examples. * - * @author IronCore Labs * @see TenantSecurityClient#createCachedDecryptor(String, DocumentMetadata) * @see TenantSecurityClient#withCachedDecryptor(String, DocumentMetadata, * java.util.function.Function) From 174d3b508a23e49458a346fc6b9678f7324ecfc9 Mon Sep 17 00:00:00 2001 From: Craig Colegrove Date: Mon, 9 Mar 2026 15:32:28 -0700 Subject: [PATCH 4/8] Add cached encryptor --- flake.lock | 12 +- .../kms/v1/CachedKeyDecryptor.java | 131 +++--- .../kms/v1/CachedKeyEncryptor.java | 237 +++++++++++ .../kms/v1/DocumentEncryptor.java | 35 ++ .../kms/v1/TenantSecurityClient.java | 71 +++- .../kms/v1/TenantSecurityRequest.java | 27 ++ .../kms/v1/CachedKeyDecryptorTest.java | 78 +++- .../kms/v1/CachedKeyEncryptorTest.java | 247 ++++++++++++ .../kms/v1/CachedKeyOpsRoundTrip.java | 379 ++++++++++++++++++ 9 files changed, 1145 insertions(+), 72 deletions(-) create mode 100644 src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptor.java create mode 100644 src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentEncryptor.java create mode 100644 src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptorTest.java create mode 100644 src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyOpsRoundTrip.java diff --git a/flake.lock b/flake.lock index 7834c48..128d859 100644 --- a/flake.lock +++ b/flake.lock @@ -5,11 +5,11 @@ "systems": "systems" }, "locked": { - "lastModified": 1685518550, - "narHash": "sha256-o2d0KcvaXzTrPRIo0kOLV0/QXHhDQ5DTi+OxcjO8xqY=", + "lastModified": 1731533236, + "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=", "owner": "numtide", "repo": "flake-utils", - "rev": "a1720a10a6cfe8234c0e93907ffe81be440f4cef", + "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b", "type": "github" }, "original": { @@ -20,11 +20,11 @@ }, "nixpkgs": { "locked": { - "lastModified": 1686592866, - "narHash": "sha256-riGg89eWhXJcPNrQGcSwTEEm7CGxWC06oSX44hajeMw=", + "lastModified": 1772773019, + "narHash": "sha256-E1bxHxNKfDoQUuvriG71+f+s/NT0qWkImXsYZNFFfCs=", "owner": "NixOS", "repo": "nixpkgs", - "rev": "0eeebd64de89e4163f4d3cf34ffe925a5cf67a05", + "rev": "aca4d95fce4914b3892661bcb80b8087293536c6", "type": "github" }, "original": { diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java index fbd188c..b3c47d1 100644 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java @@ -10,6 +10,7 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutorService; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; import com.ironcorelabs.tenantsecurity.kms.v1.exception.TscException; import com.ironcorelabs.tenantsecurity.utils.CompletableFutures; @@ -51,12 +52,19 @@ public final class CachedKeyDecryptor implements DocumentDecryptor, Closeable { // Executor for async field decryption operations private final ExecutorService encryptionExecutor; + // For reporting operations on close + private final TenantSecurityRequest requestService; + private final DocumentMetadata metadata; + // Flag to track if close() has been called private final AtomicBoolean closed = new AtomicBoolean(false); // When this decryptor was created - used for timeout enforcement private final Instant createdAt; + // Count of successful decrypt operations performed + private final AtomicInteger operationCount = new AtomicInteger(0); + /** * Package-private constructor. Use TenantSecurityClient.createCachedDecryptor() to create * instances. @@ -64,8 +72,11 @@ public final class CachedKeyDecryptor implements DocumentDecryptor, Closeable { * @param dek The unwrapped document encryption key bytes (will be copied) * @param edek The encrypted document encryption key string * @param encryptionExecutor Executor for async decryption operations + * @param requestService TSP request service for reporting operations on close + * @param metadata Document metadata for reporting operations on close */ - CachedKeyDecryptor(byte[] dek, String edek, ExecutorService encryptionExecutor) { + CachedKeyDecryptor(byte[] dek, String edek, ExecutorService encryptionExecutor, + TenantSecurityRequest requestService, DocumentMetadata metadata) { if (dek == null || dek.length != 32) { throw new IllegalArgumentException("DEK must be exactly 32 bytes"); } @@ -75,10 +86,18 @@ public final class CachedKeyDecryptor implements DocumentDecryptor, Closeable { if (encryptionExecutor == null) { throw new IllegalArgumentException("encryptionExecutor must not be null"); } + if (requestService == null) { + throw new IllegalArgumentException("requestService must not be null"); + } + if (metadata == null) { + throw new IllegalArgumentException("metadata must not be null"); + } // Copy DEK to prevent external modification this.dek = Arrays.copyOf(dek, dek.length); this.edek = edek; this.encryptionExecutor = encryptionExecutor; + this.requestService = requestService; + this.metadata = metadata; this.createdAt = Instant.now(); } @@ -110,6 +129,31 @@ public boolean isExpired() { return Duration.between(createdAt, Instant.now()).compareTo(TIMEOUT) > 0; } + /** + * Get the number of successful decrypt operations performed with this decryptor. + * + * @return The operation count + */ + public int getOperationCount() { + return operationCount.get(); + } + + /** + * Check if this decryptor is usable (not closed and not expired). Returns a failed future if not, + * or a completed future if usable. + */ + private CompletableFuture checkUsable() { + if (closed.get()) { + return CompletableFuture.failedFuture(new TscException( + TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "CachedKeyDecryptor has been closed")); + } + if (isExpired()) { + return CompletableFuture.failedFuture(new TscException( + TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "CachedKeyDecryptor has expired")); + } + return CompletableFuture.completedFuture(null); + } + /** * Decrypt the provided EncryptedDocument using the cached DEK. * @@ -124,25 +168,21 @@ public boolean isExpired() { @Override public CompletableFuture decrypt(EncryptedDocument encryptedDocument, DocumentMetadata metadata) { - // Check if closed or expired - if (closed.get()) { - return CompletableFuture.failedFuture(new TscException( - TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "CachedKeyDecryptor has been closed")); - } - if (isExpired()) { - return CompletableFuture.failedFuture(new TscException( - TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "CachedKeyDecryptor has expired")); - } + return checkUsable().thenCompose(unused -> { + // Validate EDEK matches + if (!edek.equals(encryptedDocument.getEdek())) { + return CompletableFuture + .failedFuture(new TscException(TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, + "EncryptedDocument EDEK does not match the cached EDEK. " + + "This decryptor can only decrypt documents with matching EDEKs.")); + } - // Validate EDEK matches - if (!edek.equals(encryptedDocument.getEdek())) { - return CompletableFuture - .failedFuture(new TscException(TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, - "EncryptedDocument EDEK does not match the cached EDEK. " - + "This decryptor can only decrypt documents with matching EDEKs.")); - } - - return decryptFields(encryptedDocument.getEncryptedFields(), encryptedDocument.getEdek()); + return decryptFields(encryptedDocument.getEncryptedFields(), encryptedDocument.getEdek()) + .thenApply(result -> { + operationCount.incrementAndGet(); + return result; + }); + }); } /** @@ -161,44 +201,32 @@ public CompletableFuture decrypt(EncryptedDocument encryptedD @Override public CompletableFuture decryptStream(String edek, InputStream input, OutputStream output, DocumentMetadata metadata) { - // Check if closed or expired - if (closed.get()) { - return CompletableFuture.failedFuture(new TscException( - TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "CachedKeyDecryptor has been closed")); - } - if (isExpired()) { - return CompletableFuture.failedFuture(new TscException( - TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "CachedKeyDecryptor has expired")); - } + return checkUsable().thenCompose(unused -> { + // Validate EDEK matches + if (!this.edek.equals(edek)) { + return CompletableFuture + .failedFuture(new TscException(TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, + "Provided EDEK does not match the cached EDEK. " + + "This decryptor can only decrypt documents with matching EDEKs.")); + } - // Validate EDEK matches - if (!this.edek.equals(edek)) { return CompletableFuture - .failedFuture(new TscException(TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, - "Provided EDEK does not match the cached EDEK. " - + "This decryptor can only decrypt documents with matching EDEKs.")); - } - - return CompletableFuture.supplyAsync( - () -> CryptoUtils.decryptStreamInternal(dek, input, output).join(), encryptionExecutor); + .supplyAsync(() -> CryptoUtils.decryptStreamInternal(this.dek, input, output).join(), + encryptionExecutor) + .thenApply(result -> { + operationCount.incrementAndGet(); + return result; + }); + }); } /** * Decrypt all fields in the document using the cached DEK. Pattern follows * TenantSecurityClient.decryptFields(). */ + // Caller must call checkUsable() before invoking this method. private CompletableFuture decryptFields(Map document, String documentEdek) { - // Check closed/expired state again before starting decryption - if (closed.get()) { - return CompletableFuture.failedFuture(new TscException( - TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "CachedKeyDecryptor has been closed")); - } - if (isExpired()) { - return CompletableFuture.failedFuture(new TscException( - TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "CachedKeyDecryptor has expired")); - } - // Parallel decrypt each field Map> decryptOps = document.entrySet().stream() .collect(Collectors.toMap(Map.Entry::getKey, @@ -215,8 +243,8 @@ private CompletableFuture decryptFields(Map d } /** - * Securely zero the DEK bytes and mark this decryptor as closed. After calling close(), all - * decrypt operations will fail. + * Securely zero the DEK bytes, report operations to the TSP, and mark this decryptor as closed. + * After calling close(), all decrypt operations will fail. * *

* This method is idempotent - calling it multiple times has no additional effect. @@ -226,6 +254,11 @@ public void close() { if (closed.compareAndSet(false, true)) { // Zero out the DEK bytes for security Arrays.fill(dek, (byte) 0); + // Report operations to TSP + int count = operationCount.get(); + if (count > 0) { + requestService.reportOperations(metadata, edek, 0, count); + } } } } diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptor.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptor.java new file mode 100644 index 0000000..25f4799 --- /dev/null +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptor.java @@ -0,0 +1,237 @@ +package com.ironcorelabs.tenantsecurity.kms.v1; + +import java.io.Closeable; +import java.io.InputStream; +import java.io.OutputStream; +import java.security.SecureRandom; +import java.time.Duration; +import java.time.Instant; +import java.util.Arrays; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Collectors; +import com.ironcorelabs.tenantsecurity.kms.v1.exception.TscException; +import com.ironcorelabs.tenantsecurity.utils.CompletableFutures; + +/** + * Holds a cached DEK (Document Encryption Key) for repeated encrypt operations without making + * additional TSP wrap calls. All documents encrypted with this instance will share the same + * DEK/EDEK pair. The DEK is securely zeroed when close() is called. + * + *

+ * This class is thread-safe and can be used concurrently for multiple encrypt operations. Once + * closed, all encrypt operations will fail. + * + *

+ * Expiration: This encryptor automatically expires after a short time period. Caching a DEK + * for long-term use is not supported as it would undermine the security benefits of key wrapping. + * The encryptor is intended for short-lived batch operations where multiple documents need to be + * encrypted in quick succession with the same key. Use {@link #isExpired()} to check expiration + * status. + * + *

+ * Instances are created via {@link TenantSecurityClient#createCachedEncryptor} or + * {@link TenantSecurityClient#withCachedEncryptor}. See those methods for usage examples. + * + * @see TenantSecurityClient#createCachedEncryptor(DocumentMetadata) + * @see TenantSecurityClient#withCachedEncryptor(DocumentMetadata, java.util.function.Function) + */ +public final class CachedKeyEncryptor implements DocumentEncryptor, Closeable { + + // Maximum time the encryptor can be used before it expires + private static final Duration TIMEOUT = Duration.ofMinutes(1); + + // The cached DEK bytes - zeroed on close() + private final byte[] dek; + + // The EDEK associated with the cached DEK + private final String edek; + + // Executor for async field encryption operations + private final ExecutorService encryptionExecutor; + + // Secure random for IV generation during encryption + private final SecureRandom secureRandom; + + // For reporting operations on close + private final TenantSecurityRequest requestService; + private final DocumentMetadata metadata; + + // Flag to track if close() has been called + private final AtomicBoolean closed = new AtomicBoolean(false); + + // When this encryptor was created - used for timeout enforcement + private final Instant createdAt; + + // Count of successful encrypt operations performed + private final AtomicInteger operationCount = new AtomicInteger(0); + + /** + * Package-private constructor. Use TenantSecurityClient.createCachedEncryptor() to create + * instances. + * + * @param dek The document encryption key bytes (will be copied) + * @param edek The encrypted document encryption key string + * @param encryptionExecutor Executor for async encryption operations + * @param secureRandom Secure random for IV generation + * @param requestService TSP request service for reporting operations on close + * @param metadata Document metadata for reporting operations on close + */ + CachedKeyEncryptor(byte[] dek, String edek, ExecutorService encryptionExecutor, + SecureRandom secureRandom, TenantSecurityRequest requestService, DocumentMetadata metadata) { + if (dek == null || dek.length != 32) { + throw new IllegalArgumentException("DEK must be exactly 32 bytes"); + } + if (edek == null || edek.isEmpty()) { + throw new IllegalArgumentException("EDEK must not be null or empty"); + } + if (encryptionExecutor == null) { + throw new IllegalArgumentException("encryptionExecutor must not be null"); + } + if (secureRandom == null) { + throw new IllegalArgumentException("secureRandom must not be null"); + } + if (requestService == null) { + throw new IllegalArgumentException("requestService must not be null"); + } + if (metadata == null) { + throw new IllegalArgumentException("metadata must not be null"); + } + // Copy DEK to prevent external modification + this.dek = Arrays.copyOf(dek, dek.length); + this.edek = edek; + this.encryptionExecutor = encryptionExecutor; + this.secureRandom = secureRandom; + this.requestService = requestService; + this.metadata = metadata; + this.createdAt = Instant.now(); + } + + /** + * Get the EDEK associated with this cached encryptor. + * + * @return The EDEK string + */ + public String getEdek() { + return edek; + } + + /** + * Check if this encryptor has been closed. + * + * @return true if close() has been called + */ + public boolean isClosed() { + return closed.get(); + } + + /** + * Check if this encryptor has expired due to timeout. + * + * @return true if the timeout has elapsed since creation + */ + public boolean isExpired() { + return Duration.between(createdAt, Instant.now()).compareTo(TIMEOUT) > 0; + } + + /** + * Get the number of successful encrypt operations performed with this encryptor. + * + * @return The operation count + */ + public int getOperationCount() { + return operationCount.get(); + } + + /** + * Check if this encryptor is usable (not closed and not expired). Returns a failed future if not, + * or a completed future if usable. + */ + private CompletableFuture checkUsable() { + if (closed.get()) { + return CompletableFuture.failedFuture(new TscException( + TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED, "CachedKeyEncryptor has been closed")); + } + if (isExpired()) { + return CompletableFuture.failedFuture(new TscException( + TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED, "CachedKeyEncryptor has expired")); + } + return CompletableFuture.completedFuture(null); + } + + /** + * Encrypt the provided document fields using the cached DEK. + * + * @param document Map of field names to plaintext bytes to encrypt. + * @param metadata Metadata about the document being encrypted. + * @return CompletableFuture resolving to EncryptedDocument with encrypted field bytes and EDEK. + */ + @Override + public CompletableFuture encrypt(Map document, + DocumentMetadata metadata) { + return checkUsable() + .thenCompose(unused -> encryptFields(document, metadata).thenApply(result -> { + operationCount.incrementAndGet(); + return result; + })); + } + + /** + * Encrypt a stream of bytes using the cached DEK. + * + * @param input The input stream of plaintext bytes to encrypt. + * @param output The output stream to write encrypted bytes to. + * @param metadata Metadata about the document being encrypted. + * @return CompletableFuture resolving to StreamingResponse containing the EDEK. + */ + @Override + public CompletableFuture encryptStream(InputStream input, OutputStream output, + DocumentMetadata metadata) { + return checkUsable().thenCompose(unused -> CompletableFuture.supplyAsync( + () -> CryptoUtils.encryptStreamInternal(dek, metadata, input, output, secureRandom).join(), + encryptionExecutor).thenApply(v -> { + operationCount.incrementAndGet(); + return new StreamingResponse(edek); + })); + } + + // Caller must call checkUsable() before invoking this method. + private CompletableFuture encryptFields(Map document, + DocumentMetadata metadata) { + // Parallel encrypt each field + Map> encryptOps = document.entrySet().stream() + .collect(Collectors.toMap(Map.Entry::getKey, entry -> CompletableFuture.supplyAsync( + () -> CryptoUtils.encryptBytes(entry.getValue(), metadata, dek, secureRandom).join(), + encryptionExecutor))); + + // Join all futures and build result + return CompletableFutures.tryCatchNonFatal(() -> { + Map encryptedBytes = encryptOps.entrySet().stream() + .collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().join())); + return new EncryptedDocument(encryptedBytes, edek); + }); + } + + /** + * Securely zero the DEK bytes, report operations to the TSP, and mark this encryptor as closed. + * After calling close(), all encrypt operations will fail. + * + *

+ * This method is idempotent - calling it multiple times has no additional effect. + */ + @Override + public void close() { + if (closed.compareAndSet(false, true)) { + // Zero out the DEK bytes for security + Arrays.fill(dek, (byte) 0); + // Report operations to TSP + int count = operationCount.get(); + if (count > 0) { + requestService.reportOperations(metadata, edek, count, 0); + } + } + } +} diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentEncryptor.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentEncryptor.java new file mode 100644 index 0000000..da97ee0 --- /dev/null +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentEncryptor.java @@ -0,0 +1,35 @@ +package com.ironcorelabs.tenantsecurity.kms.v1; + +import java.io.InputStream; +import java.io.OutputStream; +import java.util.Map; +import java.util.concurrent.CompletableFuture; + +/** + * Interface for document encryption capabilities. Implemented by both TenantSecurityClient (for + * standard encrypt operations that wrap a new DEK each time) and CachedKeyEncryptor (for repeated + * encrypts using a cached DEK). + */ +public interface DocumentEncryptor { + + /** + * Encrypt the provided document fields and return the resulting encrypted document. + * + * @param document Map of field names to plaintext bytes to encrypt. + * @param metadata Metadata about the document being encrypted. + * @return CompletableFuture resolving to EncryptedDocument with encrypted field bytes and EDEK. + */ + CompletableFuture encrypt(Map document, + DocumentMetadata metadata); + + /** + * Encrypt a stream of bytes. + * + * @param input The input stream of plaintext bytes to encrypt. + * @param output The output stream to write encrypted bytes to. + * @param metadata Metadata about the document being encrypted. + * @return CompletableFuture resolving to StreamingResponse containing the EDEK. + */ + CompletableFuture encryptStream(InputStream input, OutputStream output, + DocumentMetadata metadata); +} diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java index cb606bc..ee61529 100644 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java @@ -12,6 +12,7 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; +import java.util.Arrays; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.function.Function; @@ -28,7 +29,7 @@ * * @author IronCore Labs */ -public final class TenantSecurityClient implements Closeable, DocumentDecryptor { +public final class TenantSecurityClient implements Closeable, DocumentDecryptor, DocumentEncryptor { private final SecureRandom secureRandom; // Use fixed size thread pool for CPU bound operations (crypto ops). Defaults to @@ -409,6 +410,7 @@ protected BatchResult addTspFailuresToBatchResult(BatchResult batchRes * @param metadata Metadata about the document being encrypted. * @return The edek which can be used to decrypt the resulting stream */ + @Override public CompletableFuture encryptStream(InputStream input, OutputStream output, DocumentMetadata metadata) { return this.encryptionService.wrapKey(metadata).thenApplyAsync( @@ -455,6 +457,7 @@ public CompletableFuture decryptStream(String edek, InputStream input, Out * @return Encrypted document and base64 encrypted document key (EDEK) wrapped in a * EncryptedResult class. */ + @Override public CompletableFuture encrypt(Map document, DocumentMetadata metadata) { return this.encryptionService.wrapKey(metadata) @@ -581,8 +584,12 @@ public CompletableFuture decrypt(EncryptedDocument encryptedD */ public CompletableFuture createCachedDecryptor(String edek, DocumentMetadata metadata) { - return this.encryptionService.unwrapKey(edek, metadata) - .thenApply(dekBytes -> new CachedKeyDecryptor(dekBytes, edek, this.encryptionExecutor)); + return this.encryptionService.unwrapKey(edek, metadata).thenApply(dekBytes -> { + CachedKeyDecryptor decryptor = new CachedKeyDecryptor(dekBytes, edek, + this.encryptionExecutor, this.encryptionService, metadata); + Arrays.fill(dekBytes, (byte) 0); + return decryptor; + }); } /** @@ -639,6 +646,64 @@ public CompletableFuture withCachedDecryptor(EncryptedDocument encryptedD return withCachedDecryptor(encryptedDocument.getEdek(), metadata, operation); } + /** + * Create a CachedKeyEncryptor for repeated encrypt operations using the same DEK. This wraps a + * new key once and caches the resulting DEK/EDEK pair for subsequent encrypts. All documents + * encrypted with this instance will share the same DEK/EDEK pair. + * + *

+ * Use this when you need to encrypt multiple documents for the same tenant in quick succession, + * to avoid repeated TSP wrap calls. + * + *

+ * The returned encryptor implements AutoCloseable and should be used with try-with-resources to + * ensure the DEK is securely zeroed when done: + * + * 

+   * try (CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get()) {
+   *   EncryptedDocument enc1 = encryptor.encrypt(doc1, metadata).get();
+   *   EncryptedDocument enc2 = encryptor.encrypt(doc2, metadata).get();
+   * }
+   *
+ * + * @param metadata Metadata for the wrap operation + * @return CompletableFuture resolving to a CachedKeyEncryptor + */ + public CompletableFuture createCachedEncryptor(DocumentMetadata metadata) { + return this.encryptionService.wrapKey(metadata).thenApply(wrappedKey -> { + byte[] dekBytes = wrappedKey.getDekBytes(); + CachedKeyEncryptor encryptor = new CachedKeyEncryptor(dekBytes, wrappedKey.getEdek(), + this.encryptionExecutor, this.secureRandom, this.encryptionService, metadata); + Arrays.fill(dekBytes, (byte) 0); + return encryptor; + }); + } + + /** + * Execute an operation using a CachedKeyEncryptor with automatic lifecycle management. The + * encryptor is automatically closed (and DEK zeroed) when the operation completes, whether + * successfully or with an error. + * + *

+ * This is the recommended pattern for using cached encryptors with CompletableFuture composition: + * + * 

+   * client.withCachedEncryptor(metadata, encryptor ->
+   *     encryptor.encrypt(doc1, metadata)
+   *         .thenCompose(enc1 -> encryptor.encrypt(doc2, metadata)))
+   *
+ * + * @param The type returned by the operation + * @param metadata Metadata for the wrap operation + * @param operation Function that takes the encryptor and returns a CompletableFuture + * @return CompletableFuture resolving to the operation's result + */ + public CompletableFuture withCachedEncryptor(DocumentMetadata metadata, + Function> operation) { + return createCachedEncryptor(metadata).thenCompose( + encryptor -> operation.apply(encryptor).whenComplete((result, error) -> encryptor.close())); + } + /** * Re-key a document's encrypted document key (EDEK) using a new KMS config. Decrypts the EDEK * then re-encrypts it using the specified tenant's current primary KMS config. The DEK is then diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityRequest.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityRequest.java index ec88381..9511ced 100644 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityRequest.java +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityRequest.java @@ -55,6 +55,7 @@ private static String stripTrailingSlash(String s) { private final GenericUrl rekeyEndpoint; private final GenericUrl securityEventEndpoint; private final GenericUrl deriveKeyEndpoint; + private final GenericUrl reportOperationsEndpoint; private final HttpRequestFactory requestFactory; private final int timeout; @@ -79,6 +80,7 @@ private static String stripTrailingSlash(String s) { this.rekeyEndpoint = new GenericUrl(tspApiPrefix + "document/rekey"); this.securityEventEndpoint = new GenericUrl(tspApiPrefix + "event/security-event"); this.deriveKeyEndpoint = new GenericUrl(tspApiPrefix + "key/derive-with-secret-path"); + this.reportOperationsEndpoint = new GenericUrl(tspApiPrefix + "document/report-operations"); this.webRequestExecutor = Executors.newFixedThreadPool(requestThreadSize); this.requestFactory = provideHttpRequestFactory(requestThreadSize, requestThreadSize); @@ -293,6 +295,31 @@ CompletableFuture logSecurityEvent(SecurityEvent event, EventMetadata meta return this.makeRequestAndParseFailure(this.securityEventEndpoint, postData, error); } + /** + * Report cached DEK operations to the TSP. Fire-and-forget — callers should not block on the + * result. + * + * @param metadata Metadata associated with the operations. + * @param edek The EDEK that was cached. + * @param wraps Number of encrypt operations performed with the cached key. + * @param unwraps Number of decrypt operations performed with the cached key. + * @return Void on success. Failures come back as exceptions. + */ + CompletableFuture reportOperations(DocumentMetadata metadata, String edek, int wraps, + int unwraps) { + Map postData = metadata.getAsPostData(); + Map cachedOps = new HashMap<>(); + cachedOps.put("wraps", wraps); + cachedOps.put("unwraps", unwraps); + Map operations = new HashMap<>(); + operations.put(edek, cachedOps); + postData.put("operations", operations); + String error = String.format( + "Unable to make request to Tenant Security Proxy report-operations endpoint. Endpoint requested: %s", + this.reportOperationsEndpoint); + return this.makeRequestAndParseFailure(this.reportOperationsEndpoint, postData, error); + } + /** * Request derive key endpoint. */ diff --git a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptorTest.java b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptorTest.java index 4cd6fdd..88231a4 100644 --- a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptorTest.java +++ b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptorTest.java @@ -21,6 +21,7 @@ public class CachedKeyDecryptorTest { private ExecutorService executor; + private TenantSecurityRequest encryptionService; private static final String TEST_EDEK = "test-edek-base64-string"; private static final String DIFFERENT_EDEK = "different-edek-base64-string"; private DocumentMetadata metadata = @@ -29,6 +30,9 @@ public class CachedKeyDecryptorTest { @BeforeClass public void setup() { executor = Executors.newFixedThreadPool(2); + // This endpoint doesn't exist, so we won't call `close` on the cached decryptor to avoid the + // report-operations request + encryptionService = new TenantSecurityRequest("http://localhost:0", "test-api-key", 1, 1000); } @AfterClass @@ -36,6 +40,13 @@ public void teardown() { if (executor != null) { executor.shutdown(); } + if (encryptionService != null) { + try { + encryptionService.close(); + } catch (Exception e) { + // ignore + } + } } private byte[] createValidDek() { @@ -44,70 +55,100 @@ private byte[] createValidDek() { return dek; } + private CachedKeyDecryptor createDecryptor() { + return new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor, encryptionService, + metadata); + } + // Constructor validation tests + @SuppressWarnings("resource") public void constructorRejectNullDek() { try { - new CachedKeyDecryptor(null, TEST_EDEK, executor); + new CachedKeyDecryptor(null, TEST_EDEK, executor, encryptionService, metadata); fail("Should have thrown IllegalArgumentException"); } catch (IllegalArgumentException e) { assertTrue(e.getMessage().contains("DEK must be exactly 32 bytes")); } } + @SuppressWarnings("resource") public void constructorRejectWrongSizeDek() { byte[] shortDek = new byte[16]; try { - new CachedKeyDecryptor(shortDek, TEST_EDEK, executor); + new CachedKeyDecryptor(shortDek, TEST_EDEK, executor, encryptionService, metadata); fail("Should have thrown IllegalArgumentException"); } catch (IllegalArgumentException e) { assertTrue(e.getMessage().contains("DEK must be exactly 32 bytes")); } } + @SuppressWarnings("resource") public void constructorRejectNullEdek() { try { - new CachedKeyDecryptor(createValidDek(), null, executor); + new CachedKeyDecryptor(createValidDek(), null, executor, encryptionService, metadata); fail("Should have thrown IllegalArgumentException"); } catch (IllegalArgumentException e) { assertTrue(e.getMessage().contains("EDEK must not be null or empty")); } } + @SuppressWarnings("resource") public void constructorRejectEmptyEdek() { try { - new CachedKeyDecryptor(createValidDek(), "", executor); + new CachedKeyDecryptor(createValidDek(), "", executor, encryptionService, metadata); fail("Should have thrown IllegalArgumentException"); } catch (IllegalArgumentException e) { assertTrue(e.getMessage().contains("EDEK must not be null or empty")); } } + @SuppressWarnings("resource") public void constructorRejectNullExecutor() { try { - new CachedKeyDecryptor(createValidDek(), TEST_EDEK, null); + new CachedKeyDecryptor(createValidDek(), TEST_EDEK, null, encryptionService, metadata); fail("Should have thrown IllegalArgumentException"); } catch (IllegalArgumentException e) { assertTrue(e.getMessage().contains("encryptionExecutor must not be null")); } } + @SuppressWarnings("resource") + public void constructorRejectNullEncryptionService() { + try { + new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor, null, metadata); + fail("Should have thrown IllegalArgumentException"); + } catch (IllegalArgumentException e) { + assertTrue(e.getMessage().contains("requestService must not be null")); + } + } + + @SuppressWarnings("resource") + public void constructorRejectNullMetadata() { + try { + new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor, encryptionService, null); + fail("Should have thrown IllegalArgumentException"); + } catch (IllegalArgumentException e) { + assertTrue(e.getMessage().contains("metadata must not be null")); + } + } + // Getter tests public void getEdekReturnsCorrectValue() { - CachedKeyDecryptor decryptor = new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor); + CachedKeyDecryptor decryptor = createDecryptor(); assertEquals(decryptor.getEdek(), TEST_EDEK); decryptor.close(); } public void isClosedReturnsFalseInitially() { - CachedKeyDecryptor decryptor = new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor); + CachedKeyDecryptor decryptor = createDecryptor(); assertFalse(decryptor.isClosed()); decryptor.close(); } public void isClosedReturnsTrueAfterClose() { - CachedKeyDecryptor decryptor = new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor); + CachedKeyDecryptor decryptor = createDecryptor(); decryptor.close(); assertTrue(decryptor.isClosed()); } @@ -115,7 +156,7 @@ public void isClosedReturnsTrueAfterClose() { // Close tests public void closeIsIdempotent() { - CachedKeyDecryptor decryptor = new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor); + CachedKeyDecryptor decryptor = createDecryptor(); decryptor.close(); assertTrue(decryptor.isClosed()); // Should not throw @@ -124,10 +165,18 @@ public void closeIsIdempotent() { assertTrue(decryptor.isClosed()); } + // Operation count tests + + public void operationCountStartsAtZero() { + CachedKeyDecryptor decryptor = createDecryptor(); + assertEquals(decryptor.getOperationCount(), 0); + decryptor.close(); + } + // Decrypt validation tests public void decryptFailsWhenClosed() { - CachedKeyDecryptor decryptor = new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor); + CachedKeyDecryptor decryptor = createDecryptor(); decryptor.close(); EncryptedDocument encDoc = new EncryptedDocument(java.util.Collections.emptyMap(), TEST_EDEK); @@ -142,7 +191,7 @@ public void decryptFailsWhenClosed() { } public void decryptFailsWhenEdekMismatch() { - CachedKeyDecryptor decryptor = new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor); + CachedKeyDecryptor decryptor = createDecryptor(); EncryptedDocument encDoc = new EncryptedDocument(java.util.Collections.emptyMap(), DIFFERENT_EDEK); @@ -161,7 +210,7 @@ public void decryptFailsWhenEdekMismatch() { // DecryptStream validation tests public void decryptStreamFailsWhenClosed() { - CachedKeyDecryptor decryptor = new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor); + CachedKeyDecryptor decryptor = createDecryptor(); decryptor.close(); ByteArrayInputStream input = new ByteArrayInputStream(new byte[0]); @@ -177,7 +226,7 @@ public void decryptStreamFailsWhenClosed() { } public void decryptStreamFailsWhenEdekMismatch() { - CachedKeyDecryptor decryptor = new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor); + CachedKeyDecryptor decryptor = createDecryptor(); ByteArrayInputStream input = new ByteArrayInputStream(new byte[0]); ByteArrayOutputStream output = new ByteArrayOutputStream(); @@ -197,7 +246,8 @@ public void decryptStreamFailsWhenEdekMismatch() { public void constructorCopiesDekToPreventExternalModification() throws Exception { byte[] originalDek = createValidDek(); - CachedKeyDecryptor decryptor = new CachedKeyDecryptor(originalDek, TEST_EDEK, executor); + CachedKeyDecryptor decryptor = + new CachedKeyDecryptor(originalDek, TEST_EDEK, executor, encryptionService, metadata); // Modify the original array Arrays.fill(originalDek, (byte) 0x00); diff --git a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptorTest.java b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptorTest.java new file mode 100644 index 0000000..3a1644c --- /dev/null +++ b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptorTest.java @@ -0,0 +1,247 @@ +package com.ironcorelabs.tenantsecurity.kms.v1; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.fail; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.lang.reflect.Field; +import java.security.SecureRandom; +import java.util.Arrays; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; +import com.ironcorelabs.tenantsecurity.kms.v1.exception.TscException; + +@Test(groups = {"unit"}) +public class CachedKeyEncryptorTest { + + private ExecutorService executor; + private SecureRandom secureRandom; + private TenantSecurityRequest encryptionService; + private static final String TEST_EDEK = "test-edek-base64-string"; + private DocumentMetadata metadata = + new DocumentMetadata("tenantId", "requestingUserOrServiceId", "dataLabel"); + + @BeforeClass + public void setup() { + executor = Executors.newFixedThreadPool(2); + secureRandom = new SecureRandom(); + // This endpoint doesn't exist, so we won't call `close` on the cached encryptor to avoid the + // report-operations request + encryptionService = new TenantSecurityRequest("http://localhost:0", "test-api-key", 1, 1000); + } + + @AfterClass + public void teardown() { + if (executor != null) { + executor.shutdown(); + } + if (encryptionService != null) { + try { + encryptionService.close(); + } catch (Exception e) { + // ignore + } + } + } + + private byte[] createValidDek() { + byte[] dek = new byte[32]; + Arrays.fill(dek, (byte) 0x42); + return dek; + } + + private CachedKeyEncryptor createEncryptor() { + return new CachedKeyEncryptor(createValidDek(), TEST_EDEK, executor, secureRandom, + encryptionService, metadata); + } + + // Constructor validation tests + + @SuppressWarnings("resource") + public void constructorRejectNullDek() { + try { + new CachedKeyEncryptor(null, TEST_EDEK, executor, secureRandom, encryptionService, metadata); + fail("Should have thrown IllegalArgumentException"); + } catch (IllegalArgumentException e) { + assertTrue(e.getMessage().contains("DEK must be exactly 32 bytes")); + } + } + + @SuppressWarnings("resource") + public void constructorRejectWrongSizeDek() { + byte[] shortDek = new byte[16]; + try { + new CachedKeyEncryptor(shortDek, TEST_EDEK, executor, secureRandom, encryptionService, + metadata); + fail("Should have thrown IllegalArgumentException"); + } catch (IllegalArgumentException e) { + assertTrue(e.getMessage().contains("DEK must be exactly 32 bytes")); + } + } + + @SuppressWarnings("resource") + public void constructorRejectNullEdek() { + try { + new CachedKeyEncryptor(createValidDek(), null, executor, secureRandom, encryptionService, + metadata); + fail("Should have thrown IllegalArgumentException"); + } catch (IllegalArgumentException e) { + assertTrue(e.getMessage().contains("EDEK must not be null or empty")); + } + } + + @SuppressWarnings("resource") + public void constructorRejectEmptyEdek() { + try { + new CachedKeyEncryptor(createValidDek(), "", executor, secureRandom, encryptionService, + metadata); + fail("Should have thrown IllegalArgumentException"); + } catch (IllegalArgumentException e) { + assertTrue(e.getMessage().contains("EDEK must not be null or empty")); + } + } + + @SuppressWarnings("resource") + public void constructorRejectNullExecutor() { + try { + new CachedKeyEncryptor(createValidDek(), TEST_EDEK, null, secureRandom, encryptionService, + metadata); + fail("Should have thrown IllegalArgumentException"); + } catch (IllegalArgumentException e) { + assertTrue(e.getMessage().contains("encryptionExecutor must not be null")); + } + } + + @SuppressWarnings("resource") + public void constructorRejectNullSecureRandom() { + try { + new CachedKeyEncryptor(createValidDek(), TEST_EDEK, executor, null, encryptionService, + metadata); + fail("Should have thrown IllegalArgumentException"); + } catch (IllegalArgumentException e) { + assertTrue(e.getMessage().contains("secureRandom must not be null")); + } + } + + @SuppressWarnings("resource") + public void constructorRejectNullEncryptionService() { + try { + new CachedKeyEncryptor(createValidDek(), TEST_EDEK, executor, secureRandom, null, metadata); + fail("Should have thrown IllegalArgumentException"); + } catch (IllegalArgumentException e) { + assertTrue(e.getMessage().contains("requestService must not be null")); + } + } + + @SuppressWarnings("resource") + public void constructorRejectNullMetadata() { + try { + new CachedKeyEncryptor(createValidDek(), TEST_EDEK, executor, secureRandom, encryptionService, + null); + fail("Should have thrown IllegalArgumentException"); + } catch (IllegalArgumentException e) { + assertTrue(e.getMessage().contains("metadata must not be null")); + } + } + + // Getter tests + + public void getEdekReturnsCorrectValue() { + CachedKeyEncryptor encryptor = createEncryptor(); + assertEquals(encryptor.getEdek(), TEST_EDEK); + encryptor.close(); + } + + public void isClosedReturnsFalseInitially() { + CachedKeyEncryptor encryptor = createEncryptor(); + assertFalse(encryptor.isClosed()); + encryptor.close(); + } + + public void isClosedReturnsTrueAfterClose() { + CachedKeyEncryptor encryptor = createEncryptor(); + encryptor.close(); + assertTrue(encryptor.isClosed()); + } + + // Close tests + + public void closeIsIdempotent() { + CachedKeyEncryptor encryptor = createEncryptor(); + encryptor.close(); + assertTrue(encryptor.isClosed()); + // Should not throw + encryptor.close(); + encryptor.close(); + assertTrue(encryptor.isClosed()); + } + + // Operation count tests + + public void operationCountStartsAtZero() { + CachedKeyEncryptor encryptor = createEncryptor(); + assertEquals(encryptor.getOperationCount(), 0); + encryptor.close(); + } + + // Encrypt validation tests + + public void encryptFailsWhenClosed() { + CachedKeyEncryptor encryptor = createEncryptor(); + encryptor.close(); + + try { + encryptor.encrypt(java.util.Collections.emptyMap(), metadata).join(); + fail("Should have thrown CompletionException"); + } catch (CompletionException e) { + assertTrue(e.getCause() instanceof TscException); + assertTrue(e.getCause().getMessage().contains("CachedKeyEncryptor has been closed")); + } + } + + public void encryptStreamFailsWhenClosed() { + CachedKeyEncryptor encryptor = createEncryptor(); + encryptor.close(); + + ByteArrayInputStream input = new ByteArrayInputStream(new byte[0]); + ByteArrayOutputStream output = new ByteArrayOutputStream(); + + try { + encryptor.encryptStream(input, output, metadata).join(); + fail("Should have thrown CompletionException"); + } catch (CompletionException e) { + assertTrue(e.getCause() instanceof TscException); + assertTrue(e.getCause().getMessage().contains("CachedKeyEncryptor has been closed")); + } + } + + // DEK copying test + + public void constructorCopiesDekToPreventExternalModification() throws Exception { + byte[] originalDek = createValidDek(); + CachedKeyEncryptor encryptor = new CachedKeyEncryptor(originalDek, TEST_EDEK, executor, + secureRandom, encryptionService, metadata); + + // Modify the original array + Arrays.fill(originalDek, (byte) 0x00); + + // Use reflection to verify internal DEK still has original values + Field dekField = CachedKeyEncryptor.class.getDeclaredField("dek"); + dekField.setAccessible(true); + byte[] internalDek = (byte[]) dekField.get(encryptor); + + // Internal DEK should still be 0x42, not 0x00 + for (byte b : internalDek) { + assertEquals(b, (byte) 0x42, "Internal DEK should not be affected by external modification"); + } + + encryptor.close(); + } +} diff --git a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyOpsRoundTrip.java b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyOpsRoundTrip.java new file mode 100644 index 0000000..7ee77e7 --- /dev/null +++ b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyOpsRoundTrip.java @@ -0,0 +1,379 @@ +package com.ironcorelabs.tenantsecurity.kms.v1; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import com.ironcorelabs.tenantsecurity.TestUtils; +import org.testng.annotations.Test; + +@Test(groups = {"local-integration"}) +public class CachedKeyOpsRoundTrip { + private static String TENANT_ID = "tenant-gcp"; + private static String API_KEY = "0WUaXesNgbTAuLwn"; + + private void assertEqualBytes(byte[] one, byte[] two) throws Exception { + assertEquals(new String(one, "UTF-8"), new String(two, "UTF-8")); + } + + private CompletableFuture getClient() { + Map envVars = System.getenv(); + String tsp_address = envVars.getOrDefault("TSP_ADDRESS", TestSettings.TSP_ADDRESS); + String tsp_port = + TestUtils.ensureLeadingColon(envVars.getOrDefault("TSP_PORT", TestSettings.TSP_PORT)); + String api_key = envVars.getOrDefault("API_KEY", API_KEY); + return TestUtils.createTscWithAllowInsecure(tsp_address + tsp_port, api_key); + } + + private DocumentMetadata getMetadata() { + Map envVars = System.getenv(); + String tenant_id = envVars.getOrDefault("TENANT_ID", TENANT_ID); + return new DocumentMetadata(tenant_id, "integrationTest", "cachedKeyOps"); + } + + private Map getDocumentFields() throws Exception { + Map documentMap = new HashMap<>(); + documentMap.put("field1", "First field data".getBytes("UTF-8")); + documentMap.put("field2", "Second field data".getBytes("UTF-8")); + documentMap.put("field3", "Third field data".getBytes("UTF-8")); + return documentMap; + } + + // === CachedKeyEncryptor tests === + + public void cachedEncryptorRoundTrip() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc1 = getDocumentFields(); + Map doc2 = new HashMap<>(); + doc2.put("other", "Other document data".getBytes("UTF-8")); + + TenantSecurityClient client = getClient().get(); + + try (CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get()) { + assertFalse(encryptor.isClosed()); + assertFalse(encryptor.isExpired()); + assertEquals(encryptor.getOperationCount(), 0); + + // Encrypt two documents with the cached encryptor + EncryptedDocument enc1 = encryptor.encrypt(doc1, metadata).get(); + EncryptedDocument enc2 = encryptor.encrypt(doc2, metadata).get(); + + assertEquals(encryptor.getOperationCount(), 2); + + // All documents should share the same EDEK + assertEquals(enc1.getEdek(), enc2.getEdek()); + assertEquals(enc1.getEdek(), encryptor.getEdek()); + + // Decrypt with standard client and verify roundtrip + PlaintextDocument dec1 = client.decrypt(enc1, metadata).get(); + PlaintextDocument dec2 = client.decrypt(enc2, metadata).get(); + + assertEqualBytes(dec1.getDecryptedFields().get("field1"), doc1.get("field1")); + assertEqualBytes(dec1.getDecryptedFields().get("field2"), doc1.get("field2")); + assertEqualBytes(dec1.getDecryptedFields().get("field3"), doc1.get("field3")); + assertEqualBytes(dec2.getDecryptedFields().get("other"), doc2.get("other")); + } + + client.close(); + } + + public void cachedEncryptorWithPattern() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc = getDocumentFields(); + + TenantSecurityClient client = getClient().get(); + + // Use the withCachedEncryptor pattern for automatic lifecycle management + EncryptedDocument encrypted = + client.withCachedEncryptor(metadata, encryptor -> encryptor.encrypt(doc, metadata)).get(); + + // Verify the encrypted document can be decrypted + PlaintextDocument decrypted = client.decrypt(encrypted, metadata).get(); + assertEqualBytes(decrypted.getDecryptedFields().get("field1"), doc.get("field1")); + assertEqualBytes(decrypted.getDecryptedFields().get("field2"), doc.get("field2")); + assertEqualBytes(decrypted.getDecryptedFields().get("field3"), doc.get("field3")); + + client.close(); + } + + // === CachedKeyDecryptor tests === + + public void cachedDecryptorRoundTrip() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc1 = getDocumentFields(); + Map doc2 = new HashMap<>(); + doc2.put("other", "Other document data".getBytes("UTF-8")); + + TenantSecurityClient client = getClient().get(); + + // Encrypt two documents with the same key (using cached encryptor) + EncryptedDocument enc1; + EncryptedDocument enc2; + try (CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get()) { + enc1 = encryptor.encrypt(doc1, metadata).get(); + enc2 = encryptor.encrypt(doc2, metadata).get(); + } + + // Decrypt both using a cached decryptor (single unwrap call) + try (CachedKeyDecryptor decryptor = + client.createCachedDecryptor(enc1.getEdek(), metadata).get()) { + assertFalse(decryptor.isClosed()); + assertFalse(decryptor.isExpired()); + assertEquals(decryptor.getOperationCount(), 0); + + PlaintextDocument dec1 = decryptor.decrypt(enc1, metadata).get(); + PlaintextDocument dec2 = decryptor.decrypt(enc2, metadata).get(); + + assertEquals(decryptor.getOperationCount(), 2); + + assertEqualBytes(dec1.getDecryptedFields().get("field1"), doc1.get("field1")); + assertEqualBytes(dec1.getDecryptedFields().get("field2"), doc1.get("field2")); + assertEqualBytes(dec1.getDecryptedFields().get("field3"), doc1.get("field3")); + assertEqualBytes(dec2.getDecryptedFields().get("other"), doc2.get("other")); + } + + client.close(); + } + + public void cachedDecryptorFromEncryptedDocument() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc = getDocumentFields(); + + TenantSecurityClient client = getClient().get(); + + EncryptedDocument encrypted = client.encrypt(doc, metadata).get(); + + // Create decryptor from EncryptedDocument directly + try (CachedKeyDecryptor decryptor = client.createCachedDecryptor(encrypted, metadata).get()) { + PlaintextDocument decrypted = decryptor.decrypt(encrypted, metadata).get(); + assertEqualBytes(decrypted.getDecryptedFields().get("field1"), doc.get("field1")); + assertEqualBytes(decrypted.getDecryptedFields().get("field2"), doc.get("field2")); + assertEqualBytes(decrypted.getDecryptedFields().get("field3"), doc.get("field3")); + } + + client.close(); + } + + public void cachedDecryptorWithPattern() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc = getDocumentFields(); + + TenantSecurityClient client = getClient().get(); + + EncryptedDocument encrypted = client.encrypt(doc, metadata).get(); + + // Use the withCachedDecryptor pattern for automatic lifecycle management + PlaintextDocument decrypted = client.withCachedDecryptor(encrypted, metadata, + decryptor -> decryptor.decrypt(encrypted, metadata)).get(); + + assertEqualBytes(decrypted.getDecryptedFields().get("field1"), doc.get("field1")); + assertEqualBytes(decrypted.getDecryptedFields().get("field2"), doc.get("field2")); + assertEqualBytes(decrypted.getDecryptedFields().get("field3"), doc.get("field3")); + + client.close(); + } + + // === Streaming tests === + + public void cachedEncryptorStreamRoundTrip() throws Exception { + DocumentMetadata metadata = getMetadata(); + byte[] plaintext = "Stream encrypt with cached key test data".getBytes("UTF-8"); + + TenantSecurityClient client = getClient().get(); + + ByteArrayOutputStream encryptedOutput = new ByteArrayOutputStream(); + String edek; + + try (CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get()) { + ByteArrayInputStream input = new ByteArrayInputStream(plaintext); + StreamingResponse response = encryptor.encryptStream(input, encryptedOutput, metadata).get(); + edek = response.getEdek(); + assertEquals(encryptor.getOperationCount(), 1); + // EDEK from streaming response should match the encryptor's EDEK + assertEquals(edek, encryptor.getEdek()); + } + + // Decrypt with standard client + ByteArrayInputStream encryptedInput = new ByteArrayInputStream(encryptedOutput.toByteArray()); + ByteArrayOutputStream decryptedOutput = new ByteArrayOutputStream(); + client.decryptStream(edek, encryptedInput, decryptedOutput, metadata).get(); + + assertEqualBytes(decryptedOutput.toByteArray(), plaintext); + + client.close(); + } + + public void cachedDecryptorStreamRoundTrip() throws Exception { + DocumentMetadata metadata = getMetadata(); + byte[] plaintext = "Stream decrypt with cached key test data".getBytes("UTF-8"); + + TenantSecurityClient client = getClient().get(); + + // Encrypt with standard client + ByteArrayOutputStream encryptedOutput = new ByteArrayOutputStream(); + ByteArrayInputStream input = new ByteArrayInputStream(plaintext); + StreamingResponse encResponse = client.encryptStream(input, encryptedOutput, metadata).get(); + String edek = encResponse.getEdek(); + + // Decrypt with cached decryptor + try (CachedKeyDecryptor decryptor = client.createCachedDecryptor(edek, metadata).get()) { + ByteArrayInputStream encryptedInput = new ByteArrayInputStream(encryptedOutput.toByteArray()); + ByteArrayOutputStream decryptedOutput = new ByteArrayOutputStream(); + decryptor.decryptStream(edek, encryptedInput, decryptedOutput, metadata).get(); + + assertEquals(decryptor.getOperationCount(), 1); + assertEqualBytes(decryptedOutput.toByteArray(), plaintext); + } + + client.close(); + } + + // === Full round-trip: cached encrypt -> cached decrypt === + + public void cachedEncryptToCachedDecryptRoundTrip() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc1 = getDocumentFields(); + Map doc2 = new HashMap<>(); + doc2.put("solo", "Solo field document".getBytes("UTF-8")); + + TenantSecurityClient client = getClient().get(); + + // Encrypt multiple docs with cached encryptor + EncryptedDocument enc1; + EncryptedDocument enc2; + try (CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get()) { + enc1 = encryptor.encrypt(doc1, metadata).get(); + enc2 = encryptor.encrypt(doc2, metadata).get(); + } + + // Decrypt all with cached decryptor (one unwrap call for all) + try (CachedKeyDecryptor decryptor = + client.createCachedDecryptor(enc1.getEdek(), metadata).get()) { + PlaintextDocument dec1 = decryptor.decrypt(enc1, metadata).get(); + PlaintextDocument dec2 = decryptor.decrypt(enc2, metadata).get(); + + assertEqualBytes(dec1.getDecryptedFields().get("field1"), doc1.get("field1")); + assertEqualBytes(dec1.getDecryptedFields().get("field2"), doc1.get("field2")); + assertEqualBytes(dec1.getDecryptedFields().get("field3"), doc1.get("field3")); + assertEqualBytes(dec2.getDecryptedFields().get("solo"), doc2.get("solo")); + } + + client.close(); + } + + public void cachedStreamEncryptToCachedStreamDecryptRoundTrip() throws Exception { + DocumentMetadata metadata = getMetadata(); + byte[] plaintext = "Full cached stream round-trip data".getBytes("UTF-8"); + + TenantSecurityClient client = getClient().get(); + + // Encrypt stream with cached encryptor + ByteArrayOutputStream encryptedOutput = new ByteArrayOutputStream(); + String edek; + try (CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get()) { + ByteArrayInputStream input = new ByteArrayInputStream(plaintext); + StreamingResponse response = encryptor.encryptStream(input, encryptedOutput, metadata).get(); + edek = response.getEdek(); + } + + // Decrypt stream with cached decryptor + try (CachedKeyDecryptor decryptor = client.createCachedDecryptor(edek, metadata).get()) { + ByteArrayInputStream encryptedInput = new ByteArrayInputStream(encryptedOutput.toByteArray()); + ByteArrayOutputStream decryptedOutput = new ByteArrayOutputStream(); + decryptor.decryptStream(edek, encryptedInput, decryptedOutput, metadata).get(); + assertEqualBytes(decryptedOutput.toByteArray(), plaintext); + } + + client.close(); + } + + // === Encryptor close behavior === + + public void cachedEncryptorRejectsAfterClose() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc = getDocumentFields(); + + TenantSecurityClient client = getClient().get(); + + CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get(); + // Encrypt once to verify it works + encryptor.encrypt(doc, metadata).get(); + assertEquals(encryptor.getOperationCount(), 1); + + // Close and verify it rejects + encryptor.close(); + assertTrue(encryptor.isClosed()); + + try { + encryptor.encrypt(doc, metadata).get(); + assertTrue(false, "Should have thrown after close"); + } catch (ExecutionException e) { + assertTrue(e.getCause().getMessage().contains("closed")); + } + + client.close(); + } + + public void cachedDecryptorRejectsAfterClose() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc = getDocumentFields(); + + TenantSecurityClient client = getClient().get(); + + EncryptedDocument encrypted = client.encrypt(doc, metadata).get(); + + CachedKeyDecryptor decryptor = + client.createCachedDecryptor(encrypted.getEdek(), metadata).get(); + // Decrypt once to verify it works + decryptor.decrypt(encrypted, metadata).get(); + assertEquals(decryptor.getOperationCount(), 1); + + // Close and verify it rejects + decryptor.close(); + assertTrue(decryptor.isClosed()); + + try { + decryptor.decrypt(encrypted, metadata).get(); + assertTrue(false, "Should have thrown after close"); + } catch (ExecutionException e) { + assertTrue(e.getCause().getMessage().contains("closed")); + } + + client.close(); + } + + // === EDEK mismatch === + + public void cachedDecryptorRejectsEdekMismatch() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc = getDocumentFields(); + + TenantSecurityClient client = getClient().get(); + + // Encrypt two documents with different keys + EncryptedDocument enc1 = client.encrypt(doc, metadata).get(); + EncryptedDocument enc2 = client.encrypt(doc, metadata).get(); + + // Create decryptor for enc1's EDEK + try (CachedKeyDecryptor decryptor = + client.createCachedDecryptor(enc1.getEdek(), metadata).get()) { + // Decrypting enc1 should work + decryptor.decrypt(enc1, metadata).get(); + + // Decrypting enc2 (different EDEK) should fail + try { + decryptor.decrypt(enc2, metadata).get(); + assertTrue(false, "Should have thrown for EDEK mismatch"); + } catch (ExecutionException e) { + assertTrue(e.getCause().getMessage().contains("EDEK does not match")); + } + } + + client.close(); + } +} From 44541ff08eb7efb520770cb9ca4110b5fb9ae5bc Mon Sep 17 00:00:00 2001 From: Colt Frederickson Date: Tue, 10 Mar 2026 09:26:27 -0600 Subject: [PATCH 5/8] formatting --- .../tenantsecurity/kms/v1/CachedKeyDecryptor.java | 4 ++-- .../tenantsecurity/kms/v1/TenantSecurityClient.java | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java index b3c47d1..b9e0d1d 100644 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java @@ -171,8 +171,8 @@ public CompletableFuture decrypt(EncryptedDocument encryptedD return checkUsable().thenCompose(unused -> { // Validate EDEK matches if (!edek.equals(encryptedDocument.getEdek())) { - return CompletableFuture - .failedFuture(new TscException(TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, + return CompletableFuture.failedFuture( + new TscException(TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "EncryptedDocument EDEK does not match the cached EDEK. " + "This decryptor can only decrypt documents with matching EDEKs.")); } diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java index ee61529..12c4652 100644 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java @@ -585,8 +585,8 @@ public CompletableFuture decrypt(EncryptedDocument encryptedD public CompletableFuture createCachedDecryptor(String edek, DocumentMetadata metadata) { return this.encryptionService.unwrapKey(edek, metadata).thenApply(dekBytes -> { - CachedKeyDecryptor decryptor = new CachedKeyDecryptor(dekBytes, edek, - this.encryptionExecutor, this.encryptionService, metadata); + CachedKeyDecryptor decryptor = new CachedKeyDecryptor(dekBytes, edek, this.encryptionExecutor, + this.encryptionService, metadata); Arrays.fill(dekBytes, (byte) 0); return decryptor; }); From fd2f8fee6c3f8bf6a1da881fd315e358815762cc Mon Sep 17 00:00:00 2001 From: Craig Colegrove Date: Tue, 10 Mar 2026 11:36:27 -0700 Subject: [PATCH 6/8] Checkpoint --- .../kms/v1/CachedKeyDecryptor.java | 134 +++++++--------- .../kms/v1/CachedKeyEncryptor.java | 83 +++++----- .../kms/v1/DocumentCryptoOps.java | 76 +++++++++ .../kms/v1/DocumentDecryptor.java | 11 ++ .../kms/v1/DocumentEncryptor.java | 10 ++ .../kms/v1/TenantSecurityClient.java | 122 +++------------ .../kms/v1/CachedKeyDecryptorTest.java | 41 +++++ .../kms/v1/CachedKeyEncryptorTest.java | 20 +++ .../kms/v1/CachedKeyOpsRoundTrip.java | 148 ++++++++++++++++++ .../kms/v1/DocumentCryptoOpsTest.java | 69 ++++++++ 10 files changed, 496 insertions(+), 218 deletions(-) create mode 100644 src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentCryptoOps.java create mode 100644 src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentCryptoOpsTest.java diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java index b9e0d1d..beef95a 100644 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java @@ -8,12 +8,15 @@ import java.util.Arrays; import java.util.Map; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; -import java.util.stream.Collectors; +import java.util.function.Supplier; +import java.util.function.ToIntFunction; +import com.ironcorelabs.tenantsecurity.kms.v1.exception.TenantSecurityException; import com.ironcorelabs.tenantsecurity.kms.v1.exception.TscException; -import com.ironcorelabs.tenantsecurity.utils.CompletableFutures; /** * Holds a cached DEK (Document Encryption Key) for repeated decrypt operations without making @@ -139,10 +142,14 @@ public int getOperationCount() { } /** - * Check if this decryptor is usable (not closed and not expired). Returns a failed future if not, - * or a completed future if usable. + * Guard an operation with usability checks and operation counting. Verifies the decryptor is not + * closed or expired before running the operation, and increments the operation count on success. + * + * @param operation The operation to perform + * @param countOps Extracts the number of successful operations from the result */ - private CompletableFuture checkUsable() { + private CompletableFuture executeAndIncrement(Supplier> operation, + ToIntFunction countOps) { if (closed.get()) { return CompletableFuture.failedFuture(new TscException( TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "CachedKeyDecryptor has been closed")); @@ -151,95 +158,74 @@ private CompletableFuture checkUsable() { return CompletableFuture.failedFuture(new TscException( TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "CachedKeyDecryptor has expired")); } - return CompletableFuture.completedFuture(null); + return operation.get().thenApply(result -> { + operationCount.addAndGet(countOps.applyAsInt(result)); + return result; + }); + } + + private CompletableFuture validateEdekAndDecrypt( + EncryptedDocument encryptedDocument) { + if (!edek.equals(encryptedDocument.getEdek())) { + return CompletableFuture.failedFuture( + new TscException(TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, + "EncryptedDocument EDEK does not match the cached EDEK. " + + "This decryptor can only decrypt documents with matching EDEKs.")); + } + return DocumentCryptoOps.decryptFields(encryptedDocument.getEncryptedFields(), dek, + encryptedDocument.getEdek(), encryptionExecutor); } - /** - * Decrypt the provided EncryptedDocument using the cached DEK. - * - *

- * The document's EDEK must match the EDEK used to create this decryptor, otherwise an error is - * returned. - * - * @param encryptedDocument Document to decrypt - * @param metadata Metadata about the document being decrypted (used for audit/logging) - * @return CompletableFuture resolving to PlaintextDocument - */ @Override public CompletableFuture decrypt(EncryptedDocument encryptedDocument, DocumentMetadata metadata) { - return checkUsable().thenCompose(unused -> { - // Validate EDEK matches - if (!edek.equals(encryptedDocument.getEdek())) { - return CompletableFuture.failedFuture( - new TscException(TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, - "EncryptedDocument EDEK does not match the cached EDEK. " - + "This decryptor can only decrypt documents with matching EDEKs.")); - } - - return decryptFields(encryptedDocument.getEncryptedFields(), encryptedDocument.getEdek()) - .thenApply(result -> { - operationCount.incrementAndGet(); - return result; - }); - }); + return executeAndIncrement(() -> validateEdekAndDecrypt(encryptedDocument), result -> 1); } - /** - * Decrypt a stream using the cached DEK. - * - *

- * The provided EDEK must match the EDEK used to create this decryptor, otherwise an error is - * returned. - * - * @param edek Encrypted document encryption key - must match this decryptor's EDEK - * @param input A stream representing the encrypted document - * @param output An output stream to write the decrypted document to - * @param metadata Metadata about the document being decrypted - * @return Future which will complete when input has been decrypted - */ @Override public CompletableFuture decryptStream(String edek, InputStream input, OutputStream output, DocumentMetadata metadata) { - return checkUsable().thenCompose(unused -> { - // Validate EDEK matches + return executeAndIncrement(() -> { if (!this.edek.equals(edek)) { return CompletableFuture .failedFuture(new TscException(TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "Provided EDEK does not match the cached EDEK. " + "This decryptor can only decrypt documents with matching EDEKs.")); } - return CompletableFuture .supplyAsync(() -> CryptoUtils.decryptStreamInternal(this.dek, input, output).join(), - encryptionExecutor) - .thenApply(result -> { - operationCount.incrementAndGet(); - return result; - }); - }); + encryptionExecutor); + }, result -> 1); } - /** - * Decrypt all fields in the document using the cached DEK. Pattern follows - * TenantSecurityClient.decryptFields(). - */ - // Caller must call checkUsable() before invoking this method. - private CompletableFuture decryptFields(Map document, - String documentEdek) { - // Parallel decrypt each field - Map> decryptOps = document.entrySet().stream() - .collect(Collectors.toMap(Map.Entry::getKey, - entry -> CompletableFuture.supplyAsync( - () -> CryptoUtils.decryptDocument(entry.getValue(), dek).join(), - encryptionExecutor))); - - // Join all futures and build result - return CompletableFutures.tryCatchNonFatal(() -> { - Map decryptedBytes = decryptOps.entrySet().stream() - .collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().join())); - return new PlaintextDocument(decryptedBytes, documentEdek); - }); + @Override + public CompletableFuture> decryptBatch( + Map encryptedDocuments, DocumentMetadata metadata) { + return executeAndIncrement(() -> { + ConcurrentMap> ops = new ConcurrentHashMap<>(); + ConcurrentMap edekMismatches = new ConcurrentHashMap<>(); + + encryptedDocuments.forEach((id, encDoc) -> { + if (!edek.equals(encDoc.getEdek())) { + edekMismatches.put(id, + new TscException(TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, + "EncryptedDocument EDEK does not match the cached EDEK. " + + "This decryptor can only decrypt documents with matching EDEKs.")); + } else { + ops.put(id, DocumentCryptoOps.decryptFields(encDoc.getEncryptedFields(), dek, + encDoc.getEdek(), encryptionExecutor)); + } + }); + + return CompletableFuture.supplyAsync(() -> { + BatchResult result = DocumentCryptoOps.cryptoOperationToBatchResult(ops, + TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED); + ConcurrentMap allFailures = + new ConcurrentHashMap<>(result.getFailures()); + allFailures.putAll(edekMismatches); + return new BatchResult<>(result.getSuccesses(), allFailures); + }); + }, result -> result.getSuccesses().size()); } /** diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptor.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptor.java index 25f4799..046002f 100644 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptor.java +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptor.java @@ -9,12 +9,14 @@ import java.util.Arrays; import java.util.Map; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; -import java.util.stream.Collectors; +import java.util.function.Supplier; +import java.util.function.ToIntFunction; import com.ironcorelabs.tenantsecurity.kms.v1.exception.TscException; -import com.ironcorelabs.tenantsecurity.utils.CompletableFutures; /** * Holds a cached DEK (Document Encryption Key) for repeated encrypt operations without making @@ -147,10 +149,14 @@ public int getOperationCount() { } /** - * Check if this encryptor is usable (not closed and not expired). Returns a failed future if not, - * or a completed future if usable. + * Guard an operation with usability checks and operation counting. Verifies the encryptor is not + * closed or expired before running the operation, and increments the operation count on success. + * + * @param operation The operation to perform + * @param countOps Extracts the number of successful operations from the result */ - private CompletableFuture checkUsable() { + private CompletableFuture executeAndIncrement(Supplier> operation, + ToIntFunction countOps) { if (closed.get()) { return CompletableFuture.failedFuture(new TscException( TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED, "CachedKeyEncryptor has been closed")); @@ -159,60 +165,43 @@ private CompletableFuture checkUsable() { return CompletableFuture.failedFuture(new TscException( TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED, "CachedKeyEncryptor has expired")); } - return CompletableFuture.completedFuture(null); + return operation.get().thenApply(result -> { + operationCount.addAndGet(countOps.applyAsInt(result)); + return result; + }); } - /** - * Encrypt the provided document fields using the cached DEK. - * - * @param document Map of field names to plaintext bytes to encrypt. - * @param metadata Metadata about the document being encrypted. - * @return CompletableFuture resolving to EncryptedDocument with encrypted field bytes and EDEK. - */ @Override public CompletableFuture encrypt(Map document, DocumentMetadata metadata) { - return checkUsable() - .thenCompose(unused -> encryptFields(document, metadata).thenApply(result -> { - operationCount.incrementAndGet(); - return result; - })); + return executeAndIncrement( + () -> DocumentCryptoOps.encryptFields(document, metadata, dek, edek, encryptionExecutor, + secureRandom), + result -> 1); } - /** - * Encrypt a stream of bytes using the cached DEK. - * - * @param input The input stream of plaintext bytes to encrypt. - * @param output The output stream to write encrypted bytes to. - * @param metadata Metadata about the document being encrypted. - * @return CompletableFuture resolving to StreamingResponse containing the EDEK. - */ @Override public CompletableFuture encryptStream(InputStream input, OutputStream output, DocumentMetadata metadata) { - return checkUsable().thenCompose(unused -> CompletableFuture.supplyAsync( - () -> CryptoUtils.encryptStreamInternal(dek, metadata, input, output, secureRandom).join(), - encryptionExecutor).thenApply(v -> { - operationCount.incrementAndGet(); - return new StreamingResponse(edek); - })); + return executeAndIncrement( + () -> CompletableFuture.supplyAsync( + () -> CryptoUtils.encryptStreamInternal(dek, metadata, input, output, secureRandom) + .join(), + encryptionExecutor).thenApply(v -> new StreamingResponse(edek)), + result -> 1); } - // Caller must call checkUsable() before invoking this method. - private CompletableFuture encryptFields(Map document, - DocumentMetadata metadata) { - // Parallel encrypt each field - Map> encryptOps = document.entrySet().stream() - .collect(Collectors.toMap(Map.Entry::getKey, entry -> CompletableFuture.supplyAsync( - () -> CryptoUtils.encryptBytes(entry.getValue(), metadata, dek, secureRandom).join(), - encryptionExecutor))); - - // Join all futures and build result - return CompletableFutures.tryCatchNonFatal(() -> { - Map encryptedBytes = encryptOps.entrySet().stream() - .collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().join())); - return new EncryptedDocument(encryptedBytes, edek); - }); + @Override + public CompletableFuture> encryptBatch( + Map> plaintextDocuments, DocumentMetadata metadata) { + return executeAndIncrement(() -> { + ConcurrentMap> ops = new ConcurrentHashMap<>(); + plaintextDocuments.forEach((id, doc) -> ops.put(id, + DocumentCryptoOps.encryptFields(doc, metadata, dek, edek, encryptionExecutor, + secureRandom))); + return CompletableFuture.supplyAsync(() -> DocumentCryptoOps.cryptoOperationToBatchResult(ops, + TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED)); + }, result -> result.getSuccesses().size()); } /** diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentCryptoOps.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentCryptoOps.java new file mode 100644 index 0000000..4c78847 --- /dev/null +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentCryptoOps.java @@ -0,0 +1,76 @@ +package com.ironcorelabs.tenantsecurity.kms.v1; + +import java.security.SecureRandom; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.ExecutorService; +import java.util.stream.Collectors; +import com.ironcorelabs.tenantsecurity.kms.v1.exception.TenantSecurityException; +import com.ironcorelabs.tenantsecurity.kms.v1.exception.TscException; +import com.ironcorelabs.tenantsecurity.utils.CompletableFutures; + +/** + * Package-private static helper for field-level encryption/decryption and batch result aggregation. + */ +final class DocumentCryptoOps { + + private DocumentCryptoOps() {} + + /** + * Encrypt all fields in the document using the provided DEK. Each field is encrypted in parallel + * on the provided executor. + */ + static CompletableFuture encryptFields(Map document, + DocumentMetadata metadata, byte[] dek, String edek, ExecutorService executor, + SecureRandom secureRandom) { + Map> encryptOps = document.entrySet().stream() + .collect(Collectors.toMap(Map.Entry::getKey, entry -> CompletableFuture.supplyAsync( + () -> CryptoUtils.encryptBytes(entry.getValue(), metadata, dek, secureRandom).join(), + executor))); + + return CompletableFutures.tryCatchNonFatal(() -> { + Map encryptedBytes = encryptOps.entrySet().stream() + .collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().join())); + return new EncryptedDocument(encryptedBytes, edek); + }); + } + + /** + * Decrypt all fields in the document using the provided DEK. Each field is decrypted in parallel + * on the provided executor. + */ + static CompletableFuture decryptFields(Map document, + byte[] dek, String edek, ExecutorService executor) { + Map> decryptOps = document.entrySet().stream() + .collect(Collectors.toMap(Map.Entry::getKey, entry -> CompletableFuture.supplyAsync( + () -> CryptoUtils.decryptDocument(entry.getValue(), dek).join(), executor))); + + return CompletableFutures.tryCatchNonFatal(() -> { + Map decryptedBytes = decryptOps.entrySet().stream() + .collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().join())); + return new PlaintextDocument(decryptedBytes, edek); + }); + } + + /** + * Collect a map from String to CompletableFuture into a BatchResult. Futures that complete + * exceptionally are wrapped in TscExceptions with the provided errorCode. + */ + static BatchResult cryptoOperationToBatchResult( + ConcurrentMap> operationResults, + TenantSecurityErrorCodes errorCode) { + ConcurrentMap successes = new ConcurrentHashMap<>(operationResults.size()); + ConcurrentMap failures = new ConcurrentHashMap<>(); + operationResults.entrySet().parallelStream().forEach(entry -> { + try { + T doc = entry.getValue().join(); + successes.put(entry.getKey(), doc); + } catch (Exception e) { + failures.put(entry.getKey(), new TscException(errorCode, e)); + } + }); + return new BatchResult(successes, failures); + } +} diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentDecryptor.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentDecryptor.java index 8883a74..b92d444 100644 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentDecryptor.java +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentDecryptor.java @@ -2,6 +2,7 @@ import java.io.InputStream; import java.io.OutputStream; +import java.util.Map; import java.util.concurrent.CompletableFuture; /** @@ -34,4 +35,14 @@ CompletableFuture decrypt(EncryptedDocument encryptedDocument */ CompletableFuture decryptStream(String edek, InputStream input, OutputStream output, DocumentMetadata metadata); + + /** + * Decrypt a batch of encrypted documents. Supports partial failure via {@link BatchResult}. + * + * @param encryptedDocuments Map of document ID to EncryptedDocument to decrypt. + * @param metadata Metadata about all of the documents being decrypted. + * @return Collection of successes and failures that occurred during operation. + */ + CompletableFuture> decryptBatch( + Map encryptedDocuments, DocumentMetadata metadata); } diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentEncryptor.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentEncryptor.java index da97ee0..8779903 100644 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentEncryptor.java +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentEncryptor.java @@ -32,4 +32,14 @@ CompletableFuture encrypt(Map document, */ CompletableFuture encryptStream(InputStream input, OutputStream output, DocumentMetadata metadata); + + /** + * Encrypt a batch of documents. Supports partial failure via {@link BatchResult}. + * + * @param plaintextDocuments Map of document ID to map of fields to encrypt. + * @param metadata Metadata about all of the documents being encrypted. + * @return Collection of successes and failures that occurred during operation. + */ + CompletableFuture> encryptBatch( + Map> plaintextDocuments, DocumentMetadata metadata); } diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java index 12c4652..d15757b 100644 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java @@ -19,7 +19,6 @@ import java.util.stream.Collectors; import java.util.stream.Stream; import com.ironcorelabs.tenantsecurity.kms.v1.exception.TenantSecurityException; -import com.ironcorelabs.tenantsecurity.kms.v1.exception.TscException; import com.ironcorelabs.tenantsecurity.logdriver.v1.EventMetadata; import com.ironcorelabs.tenantsecurity.logdriver.v1.SecurityEvent; import com.ironcorelabs.tenantsecurity.utils.CompletableFutures; @@ -233,61 +232,6 @@ public static CompletableFuture create(String tspDomain, S .tryCatchNonFatal(() -> new TenantSecurityClient.Builder(tspDomain, apiKey).build()); } - /** - * Encrypt the provided map of fields using the provided encryption key (DEK) and return the - * resulting encrypted document. - */ - private CompletableFuture encryptFields(Map document, - DocumentMetadata metadata, byte[] dek, String edek) { - // First, iterate over the map of documents and kick off the encrypt operation - // Future for each one. As part of doing this, we kick off the operation on to - // another thread so they run in parallel. - Map> encryptOps = - document.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, entry -> { - // Do this mapping in the .collect because we can just map the value. If we - // tried doing this in a .map above the .collect we'd have to return another - // Entry which is more complicated - return CompletableFuture.supplyAsync(() -> CryptoUtils - .encryptBytes(entry.getValue(), metadata, dek, this.secureRandom).join(), - encryptionExecutor); - })); - - return CompletableFutures.tryCatchNonFatal(() -> { - // Now iterate over our map of keys to Futures and call join on all of them. We - // do this in a separate stream() because if we called join() above it'd block - // each iteration and cause them to be run in CompletableFutures.sequence. - Map encryptedBytes = encryptOps.entrySet().stream() - .collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().join())); - return new EncryptedDocument(encryptedBytes, edek); - }); - } - - /** - * Encrypt the provided map of encrypted fields using the provided DEK and return the resulting - * decrypted document. - */ - private CompletableFuture decryptFields(Map document, - byte[] dek, String edek) { - // First map over the encrypted document map and convert the values from - // encrypted bytes to Futures of decrypted bytes. Make sure each decrypt happens - // on it's own thread to run them in parallel. - Map> decryptOps = - document.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, entry -> - // Do this mapping in the .collect because we can just map the value. If we - // tried doing this in a .map above the .collect we'd have to return another - // Entry which is more complicated - CompletableFuture.supplyAsync( - () -> CryptoUtils.decryptDocument(entry.getValue(), dek).join(), encryptionExecutor))); - // Then iterate over the map of Futures and join them to get the decrypted bytes - // out. Return the map with the same keys passed in, but the values will now be - // decrypted. - return CompletableFutures.tryCatchNonFatal(() -> { - Map decryptedBytes = decryptOps.entrySet().stream() - .collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().join())); - return new PlaintextDocument(decryptedBytes, edek); - }); - } - /** * Given a map of document IDs to plaintext bytes to encrypt and a map of document IDs to a fresh * DEK, iterate over the DEKs and encrypt the document with the same key. @@ -300,34 +244,14 @@ private BatchResult encryptBatchOfDocuments( .collect(Collectors.toConcurrentMap(ConcurrentMap.Entry::getKey, dekResult -> { String documentId = dekResult.getKey(); WrappedDocumentKey documentKeys = dekResult.getValue(); - return encryptFields(documents.get(documentId), metadata, documentKeys.getDekBytes(), - documentKeys.getEdek()); + return DocumentCryptoOps.encryptFields(documents.get(documentId), metadata, + documentKeys.getDekBytes(), documentKeys.getEdek(), encryptionExecutor, + secureRandom); })); - return cryptoOperationToBatchResult(encryptResults, + return DocumentCryptoOps.cryptoOperationToBatchResult(encryptResults, TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED); } - /** - * Collect a map from String to CompletableFuture into a BatchResult. T will be either an - * EncryptedDocument or a PlaintextDocument. CompletableFuture failures will be wrapped in - * TscExceptions with the provided errorCode and the underlying Throwable cause. - */ - private BatchResult cryptoOperationToBatchResult( - ConcurrentMap> operationResults, - TenantSecurityErrorCodes errorCode) { - ConcurrentMap successes = new ConcurrentHashMap<>(operationResults.size()); - ConcurrentMap failures = new ConcurrentHashMap<>(); - operationResults.entrySet().parallelStream().forEach(entry -> { - try { - T doc = entry.getValue().join(); - successes.put(entry.getKey(), doc); - } catch (Exception e) { - failures.put(entry.getKey(), new TscException(errorCode, e)); - } - }); - return new BatchResult(successes, failures); - } - /** * Given a map of document IDs to previously encrypted plaintext documents to re-encrypt and a map * of document IDs to the documents DEK, iterate over the DEKs and re-encrypt the document with @@ -341,10 +265,11 @@ private BatchResult encryptExistingBatchOfDocuments( .collect(Collectors.toConcurrentMap(ConcurrentMap.Entry::getKey, dekResult -> { String documentId = dekResult.getKey(); UnwrappedDocumentKey documentKeys = dekResult.getValue(); - return encryptFields(documents.get(documentId).getDecryptedFields(), metadata, - documentKeys.getDekBytes(), documents.get(documentId).getEdek()); + return DocumentCryptoOps.encryptFields(documents.get(documentId).getDecryptedFields(), + metadata, documentKeys.getDekBytes(), documents.get(documentId).getEdek(), + encryptionExecutor, secureRandom); })); - return cryptoOperationToBatchResult(encryptResults, + return DocumentCryptoOps.cryptoOperationToBatchResult(encryptResults, TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED); } @@ -361,10 +286,10 @@ private BatchResult decryptBatchDocuments( String documentId = dekResult.getKey(); UnwrappedDocumentKey documentKeys = dekResult.getValue(); EncryptedDocument eDoc = documents.get(documentId); - return decryptFields(eDoc.getEncryptedFields(), documentKeys.getDekBytes(), - eDoc.getEdek()); + return DocumentCryptoOps.decryptFields(eDoc.getEncryptedFields(), + documentKeys.getDekBytes(), eDoc.getEdek(), encryptionExecutor); })); - return cryptoOperationToBatchResult(decryptResults, + return DocumentCryptoOps.cryptoOperationToBatchResult(decryptResults, TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED); } @@ -461,8 +386,9 @@ public CompletableFuture decryptStream(String edek, InputStream input, Out public CompletableFuture encrypt(Map document, DocumentMetadata metadata) { return this.encryptionService.wrapKey(metadata) - .thenComposeAsync(newDocumentKeys -> encryptFields(document, metadata, - newDocumentKeys.getDekBytes(), newDocumentKeys.getEdek())); + .thenComposeAsync(newDocumentKeys -> DocumentCryptoOps.encryptFields(document, metadata, + newDocumentKeys.getDekBytes(), newDocumentKeys.getEdek(), encryptionExecutor, + secureRandom)); } /** @@ -483,7 +409,8 @@ public CompletableFuture encrypt(Map document public CompletableFuture encrypt(PlaintextDocument document, DocumentMetadata metadata) { return this.encryptionService.unwrapKey(document.getEdek(), metadata).thenComposeAsync( - dek -> encryptFields(document.getDecryptedFields(), metadata, dek, document.getEdek()), + dek -> DocumentCryptoOps.encryptFields(document.getDecryptedFields(), metadata, dek, + document.getEdek(), encryptionExecutor, secureRandom), encryptionExecutor); } @@ -499,6 +426,7 @@ public CompletableFuture encrypt(PlaintextDocument document, * @return Collection of successes and failures that occurred during operation. The keys of each * map returned will be the same keys provided in the original plaintextDocuments map. */ + @Override public CompletableFuture> encryptBatch( Map> plaintextDocuments, DocumentMetadata metadata) { return this.encryptionService.batchWrapKeys(plaintextDocuments.keySet(), metadata) @@ -555,8 +483,9 @@ public CompletableFuture> encryptExistingBatch( public CompletableFuture decrypt(EncryptedDocument encryptedDocument, DocumentMetadata metadata) { return this.encryptionService.unwrapKey(encryptedDocument.getEdek(), metadata).thenComposeAsync( - decryptedDocumentAESKey -> decryptFields(encryptedDocument.getEncryptedFields(), - decryptedDocumentAESKey, encryptedDocument.getEdek())); + decryptedDocumentAESKey -> DocumentCryptoOps.decryptFields( + encryptedDocument.getEncryptedFields(), decryptedDocumentAESKey, + encryptedDocument.getEdek(), encryptionExecutor)); } /** @@ -614,9 +543,8 @@ public CompletableFuture createCachedDecryptor( * This is the recommended pattern for using cached decryptors with CompletableFuture composition: * * 

-   * client.withCachedDecryptor(edek, metadata, decryptor ->
-   *     decryptor.decrypt(encDoc1, metadata)
-   *         .thenCompose(doc1 -> decryptor.decrypt(encDoc2, metadata)))
+   * client.withCachedDecryptor(edek, metadata, decryptor -> decryptor.decrypt(encDoc1, metadata)
+   *     .thenCompose(doc1 -> decryptor.decrypt(encDoc2, metadata)))
    *
* * @param The type returned by the operation @@ -688,9 +616,8 @@ public CompletableFuture createCachedEncryptor(DocumentMetad * This is the recommended pattern for using cached encryptors with CompletableFuture composition: * * 
-   * client.withCachedEncryptor(metadata, encryptor ->
-   *     encryptor.encrypt(doc1, metadata)
-   *         .thenCompose(enc1 -> encryptor.encrypt(doc2, metadata)))
+   * client.withCachedEncryptor(metadata, encryptor -> encryptor.encrypt(doc1, metadata)
+   *     .thenCompose(enc1 -> encryptor.encrypt(doc2, metadata)))
    *
* * @param The type returned by the operation @@ -732,6 +659,7 @@ public CompletableFuture rekeyEdek(String edek, DocumentMetadata metadat * @return Collection of successes and failures that occurred during operation. The keys of each * map returned will be the same keys provided in the original encryptedDocuments map. */ + @Override public CompletableFuture> decryptBatch( Map encryptedDocuments, DocumentMetadata metadata) { Map edekMap = encryptedDocuments.entrySet().stream() diff --git a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptorTest.java b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptorTest.java index 88231a4..9e4122c 100644 --- a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptorTest.java +++ b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptorTest.java @@ -9,6 +9,8 @@ import java.io.ByteArrayOutputStream; import java.lang.reflect.Field; import java.util.Arrays; +import java.util.HashMap; +import java.util.Map; import java.util.concurrent.CompletionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -242,6 +244,45 @@ public void decryptStreamFailsWhenEdekMismatch() { } } + // decryptBatch validation tests + + public void decryptBatchFailsWhenClosed() { + CachedKeyDecryptor decryptor = createDecryptor(); + decryptor.close(); + + Map docs = new HashMap<>(); + docs.put("doc1", new EncryptedDocument(java.util.Collections.emptyMap(), TEST_EDEK)); + + try { + decryptor.decryptBatch(docs, metadata).join(); + fail("Should have thrown CompletionException"); + } catch (CompletionException e) { + assertTrue(e.getCause() instanceof TscException); + assertTrue(e.getCause().getMessage().contains("CachedKeyDecryptor has been closed")); + } + } + + public void decryptBatchEdekMismatchGoesToFailures() { + CachedKeyDecryptor decryptor = createDecryptor(); + + Map docs = new HashMap<>(); + docs.put("matching", new EncryptedDocument(java.util.Collections.emptyMap(), TEST_EDEK)); + docs.put("mismatched", new EncryptedDocument(java.util.Collections.emptyMap(), DIFFERENT_EDEK)); + + BatchResult result = decryptor.decryptBatch(docs, metadata).join(); + + // The matching doc with empty fields should succeed (no fields to decrypt) + assertTrue(result.getSuccesses().containsKey("matching")); + // The mismatched doc should be in failures + assertTrue(result.getFailures().containsKey("mismatched")); + assertTrue( + result.getFailures().get("mismatched").getMessage().contains("EDEK does not match")); + // The matching doc should NOT be in failures + assertFalse(result.getFailures().containsKey("matching")); + + decryptor.close(); + } + // DEK copying test public void constructorCopiesDekToPreventExternalModification() throws Exception { diff --git a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptorTest.java b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptorTest.java index 3a1644c..7d4311a 100644 --- a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptorTest.java +++ b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptorTest.java @@ -10,6 +10,8 @@ import java.lang.reflect.Field; import java.security.SecureRandom; import java.util.Arrays; +import java.util.HashMap; +import java.util.Map; import java.util.concurrent.CompletionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -222,6 +224,24 @@ public void encryptStreamFailsWhenClosed() { } } + // encryptBatch validation tests + + public void encryptBatchFailsWhenClosed() { + CachedKeyEncryptor encryptor = createEncryptor(); + encryptor.close(); + + Map> docs = new HashMap<>(); + docs.put("doc1", java.util.Collections.singletonMap("field", new byte[] {1, 2, 3})); + + try { + encryptor.encryptBatch(docs, metadata).join(); + fail("Should have thrown CompletionException"); + } catch (CompletionException e) { + assertTrue(e.getCause() instanceof TscException); + assertTrue(e.getCause().getMessage().contains("CachedKeyEncryptor has been closed")); + } + } + // DEK copying test public void constructorCopiesDekToPreventExternalModification() throws Exception { diff --git a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyOpsRoundTrip.java b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyOpsRoundTrip.java index 7ee77e7..0438a21 100644 --- a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyOpsRoundTrip.java +++ b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyOpsRoundTrip.java @@ -349,6 +349,154 @@ public void cachedDecryptorRejectsAfterClose() throws Exception { // === EDEK mismatch === + // === Batch encrypt/decrypt tests === + + public void cachedEncryptorBatchRoundTrip() throws Exception { + DocumentMetadata metadata = getMetadata(); + + Map> docs = new HashMap<>(); + docs.put("doc1", getDocumentFields()); + Map doc2 = new HashMap<>(); + doc2.put("other", "Other data".getBytes("UTF-8")); + docs.put("doc2", doc2); + Map doc3 = new HashMap<>(); + doc3.put("solo", "Solo data".getBytes("UTF-8")); + docs.put("doc3", doc3); + + TenantSecurityClient client = getClient().get(); + + BatchResult encResult; + try (CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get()) { + encResult = encryptor.encryptBatch(docs, metadata).get(); + assertEquals(encryptor.getOperationCount(), 3); + } + + assertFalse(encResult.hasFailures()); + assertEquals(encResult.getSuccesses().size(), 3); + + // All encrypted docs share the same EDEK + String commonEdek = encResult.getSuccesses().values().iterator().next().getEdek(); + for (EncryptedDocument enc : encResult.getSuccesses().values()) { + assertEquals(enc.getEdek(), commonEdek); + } + + // Decrypt each with standard client and verify roundtrip + PlaintextDocument dec1 = client.decrypt(encResult.getSuccesses().get("doc1"), metadata).get(); + assertEqualBytes(dec1.getDecryptedFields().get("field1"), docs.get("doc1").get("field1")); + + PlaintextDocument dec2 = client.decrypt(encResult.getSuccesses().get("doc2"), metadata).get(); + assertEqualBytes(dec2.getDecryptedFields().get("other"), doc2.get("other")); + + PlaintextDocument dec3 = client.decrypt(encResult.getSuccesses().get("doc3"), metadata).get(); + assertEqualBytes(dec3.getDecryptedFields().get("solo"), doc3.get("solo")); + + client.close(); + } + + public void cachedDecryptorBatchRoundTrip() throws Exception { + DocumentMetadata metadata = getMetadata(); + + Map doc1 = getDocumentFields(); + Map doc2 = new HashMap<>(); + doc2.put("other", "Other data".getBytes("UTF-8")); + Map doc3 = new HashMap<>(); + doc3.put("solo", "Solo data".getBytes("UTF-8")); + + TenantSecurityClient client = getClient().get(); + + // Encrypt all 3 with cached encryptor (same key) + EncryptedDocument enc1, enc2, enc3; + try (CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get()) { + enc1 = encryptor.encrypt(doc1, metadata).get(); + enc2 = encryptor.encrypt(doc2, metadata).get(); + enc3 = encryptor.encrypt(doc3, metadata).get(); + } + + // Batch decrypt with cached decryptor + Map encDocs = new HashMap<>(); + encDocs.put("doc1", enc1); + encDocs.put("doc2", enc2); + encDocs.put("doc3", enc3); + + try (CachedKeyDecryptor decryptor = + client.createCachedDecryptor(enc1.getEdek(), metadata).get()) { + BatchResult result = decryptor.decryptBatch(encDocs, metadata).get(); + assertEquals(decryptor.getOperationCount(), 3); + + assertFalse(result.hasFailures()); + assertEquals(result.getSuccesses().size(), 3); + + assertEqualBytes(result.getSuccesses().get("doc1").getDecryptedFields().get("field1"), + doc1.get("field1")); + assertEqualBytes(result.getSuccesses().get("doc2").getDecryptedFields().get("other"), + doc2.get("other")); + assertEqualBytes(result.getSuccesses().get("doc3").getDecryptedFields().get("solo"), + doc3.get("solo")); + } + + client.close(); + } + + public void cachedDecryptorBatchEdekMismatchPartialFailure() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc = getDocumentFields(); + + TenantSecurityClient client = getClient().get(); + + // Encrypt two documents with different keys + EncryptedDocument enc1 = client.encrypt(doc, metadata).get(); + EncryptedDocument enc2 = client.encrypt(doc, metadata).get(); + + Map encDocs = new HashMap<>(); + encDocs.put("match", enc1); + encDocs.put("mismatch", enc2); + + // Create decryptor for enc1's key + try (CachedKeyDecryptor decryptor = + client.createCachedDecryptor(enc1.getEdek(), metadata).get()) { + BatchResult result = decryptor.decryptBatch(encDocs, metadata).get(); + + // match should succeed + assertTrue(result.getSuccesses().containsKey("match")); + assertEqualBytes(result.getSuccesses().get("match").getDecryptedFields().get("field1"), + doc.get("field1")); + + // mismatch should be in failures + assertTrue(result.getFailures().containsKey("mismatch")); + assertTrue( + result.getFailures().get("mismatch").getMessage().contains("EDEK does not match")); + + assertEquals(decryptor.getOperationCount(), 1); + } + + client.close(); + } + + public void cachedBatchOperationCount() throws Exception { + DocumentMetadata metadata = getMetadata(); + + Map> docs = new HashMap<>(); + docs.put("doc1", getDocumentFields()); + Map doc2 = new HashMap<>(); + doc2.put("field", "data".getBytes("UTF-8")); + docs.put("doc2", doc2); + + TenantSecurityClient client = getClient().get(); + + try (CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get()) { + assertEquals(encryptor.getOperationCount(), 0); + encryptor.encryptBatch(docs, metadata).get(); + assertEquals(encryptor.getOperationCount(), 2); + // Single encrypt should add 1 more + encryptor.encrypt(docs.get("doc1"), metadata).get(); + assertEquals(encryptor.getOperationCount(), 3); + } + + client.close(); + } + + // === EDEK mismatch === + public void cachedDecryptorRejectsEdekMismatch() throws Exception { DocumentMetadata metadata = getMetadata(); Map doc = getDocumentFields(); diff --git a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentCryptoOpsTest.java b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentCryptoOpsTest.java new file mode 100644 index 0000000..eeaaed6 --- /dev/null +++ b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentCryptoOpsTest.java @@ -0,0 +1,69 @@ +package com.ironcorelabs.tenantsecurity.kms.v1; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import org.testng.annotations.Test; +import com.ironcorelabs.tenantsecurity.kms.v1.exception.TscException; + +@Test(groups = {"unit"}) +public class DocumentCryptoOpsTest { + + public void cryptoOperationToBatchResultAllSuccess() { + ConcurrentMap> ops = new ConcurrentHashMap<>(); + ops.put("a", CompletableFuture.completedFuture("resultA")); + ops.put("b", CompletableFuture.completedFuture("resultB")); + + BatchResult result = DocumentCryptoOps.cryptoOperationToBatchResult(ops, + TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED); + + assertEquals(result.getSuccesses().size(), 2); + assertEquals(result.getSuccesses().get("a"), "resultA"); + assertEquals(result.getSuccesses().get("b"), "resultB"); + assertFalse(result.hasFailures()); + } + + public void cryptoOperationToBatchResultAllFailure() { + ConcurrentMap> ops = new ConcurrentHashMap<>(); + ops.put("a", CompletableFuture.failedFuture(new RuntimeException("fail A"))); + ops.put("b", CompletableFuture.failedFuture(new RuntimeException("fail B"))); + + BatchResult result = DocumentCryptoOps.cryptoOperationToBatchResult(ops, + TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED); + + assertFalse(result.hasSuccesses()); + assertEquals(result.getFailures().size(), 2); + assertTrue(result.getFailures().get("a") instanceof TscException); + assertTrue(result.getFailures().get("b") instanceof TscException); + } + + public void cryptoOperationToBatchResultMixed() { + ConcurrentMap> ops = new ConcurrentHashMap<>(); + ops.put("good", CompletableFuture.completedFuture("value")); + ops.put("bad", CompletableFuture.failedFuture(new RuntimeException("oops"))); + + BatchResult result = DocumentCryptoOps.cryptoOperationToBatchResult(ops, + TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED); + + assertTrue(result.hasSuccesses()); + assertTrue(result.hasFailures()); + assertEquals(result.getSuccesses().size(), 1); + assertEquals(result.getSuccesses().get("good"), "value"); + assertEquals(result.getFailures().size(), 1); + assertTrue(result.getFailures().containsKey("bad")); + } + + public void cryptoOperationToBatchResultEmpty() { + ConcurrentMap> ops = new ConcurrentHashMap<>(); + + BatchResult result = DocumentCryptoOps.cryptoOperationToBatchResult(ops, + TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED); + + assertFalse(result.hasSuccesses()); + assertFalse(result.hasFailures()); + } +} From a2f70a4fb77fadde9664fc3170dd954a0fcb54fb Mon Sep 17 00:00:00 2001 From: Craig Colegrove Date: Tue, 10 Mar 2026 15:17:51 -0700 Subject: [PATCH 7/8] Checkpoint after making CachedKey --- .../kms/v1/CachedDecryptor.java | 19 ++ .../kms/v1/CachedEncryptor.java | 18 ++ ...CachedKeyDecryptor.java => CachedKey.java} | 178 ++++++++---- .../kms/v1/CachedKeyEncryptor.java | 226 --------------- .../kms/v1/CachedKeyLifecycle.java | 45 +++ .../kms/v1/DocumentDecryptor.java | 4 +- .../kms/v1/DocumentEncryptor.java | 4 +- .../kms/v1/TenantSecurityClient.java | 188 +++++++++--- .../kms/v1/CachedKeyEncryptorTest.java | 267 ------------------ .../kms/v1/CachedKeyOpsRoundTrip.java | 112 +++++--- ...yDecryptorTest.java => CachedKeyTest.java} | 172 +++++++---- 11 files changed, 542 insertions(+), 691 deletions(-) create mode 100644 src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedDecryptor.java create mode 100644 src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedEncryptor.java rename src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/{CachedKeyDecryptor.java => CachedKey.java} (53%) delete mode 100644 src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptor.java create mode 100644 src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyLifecycle.java delete mode 100644 src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptorTest.java rename src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/{CachedKeyDecryptorTest.java => CachedKeyTest.java} (60%) diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedDecryptor.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedDecryptor.java new file mode 100644 index 0000000..9de4c30 --- /dev/null +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedDecryptor.java @@ -0,0 +1,19 @@ +package com.ironcorelabs.tenantsecurity.kms.v1; + +/** + * A cached document decryptor that holds a DEK for repeated decrypt operations without making + * additional TSP unwrap calls. Can only decrypt documents that were encrypted with the same + * DEK/EDEK pair. + * + *

+ * Instances are created via + * {@link TenantSecurityClient#createCachedDecryptor(String, DocumentMetadata)} or + * {@link TenantSecurityClient#withCachedDecryptor}. The cached key should be closed when done to + * securely zero the DEK. + * + * @see TenantSecurityClient#createCachedDecryptor(String, DocumentMetadata) + * @see TenantSecurityClient#withCachedDecryptor(String, DocumentMetadata, + * java.util.function.Function) + */ +public interface CachedDecryptor extends DocumentDecryptor, CachedKeyLifecycle { +} diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedEncryptor.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedEncryptor.java new file mode 100644 index 0000000..63fb820 --- /dev/null +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedEncryptor.java @@ -0,0 +1,18 @@ +package com.ironcorelabs.tenantsecurity.kms.v1; + +/** + * A cached document encryptor that holds a DEK for repeated encrypt operations without making + * additional TSP wrap calls. All documents encrypted with this instance share the same DEK/EDEK + * pair. + * + *

+ * Instances are created via {@link TenantSecurityClient#createCachedEncryptor(DocumentMetadata)} or + * {@link TenantSecurityClient#withCachedEncryptor}. The cached key should be closed when done to + * securely zero the DEK. + * + * @see TenantSecurityClient#createCachedEncryptor(DocumentMetadata) + * @see TenantSecurityClient#withCachedEncryptor(DocumentMetadata, + * java.util.function.Function) + */ +public interface CachedEncryptor extends DocumentEncryptor, CachedKeyLifecycle { +} diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKey.java similarity index 53% rename from src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java rename to src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKey.java index beef95a..b156a7a 100644 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKey.java @@ -1,8 +1,8 @@ package com.ironcorelabs.tenantsecurity.kms.v1; -import java.io.Closeable; import java.io.InputStream; import java.io.OutputStream; +import java.security.SecureRandom; import java.time.Duration; import java.time.Instant; import java.util.Arrays; @@ -19,42 +19,51 @@ import com.ironcorelabs.tenantsecurity.kms.v1.exception.TscException; /** - * Holds a cached DEK (Document Encryption Key) for repeated decrypt operations without making - * additional TSP unwrap calls. The DEK is securely zeroed when close() is called. + * Holds a cached DEK (Document Encryption Key) for repeated encrypt and decrypt operations without + * making additional TSP wrap/unwrap calls. All documents encrypted with this instance will share the + * same DEK/EDEK pair. The DEK is securely zeroed when close() is called. * *

- * This class is thread-safe and can be used concurrently for multiple decrypt operations. Once - * closed, all decrypt operations will fail. + * This class is thread-safe and can be used concurrently for multiple encrypt and decrypt + * operations. Once closed, all operations will fail. * *

- * Expiration: This decryptor automatically expires after a short time period. Caching a DEK + * Expiration: This cached key automatically expires after a short time period. Caching a DEK * for long-term use is not supported as it would undermine the security benefits of key wrapping. - * The decryptor is intended for short-lived batch operations where multiple documents sharing the - * same EDEK need to be decrypted in quick succession. Use {@link #isExpired()} to check expiration - * status. + * The cached key is intended for short-lived batch operations where multiple documents need to be + * encrypted or decrypted in quick succession with the same key. Use {@link #isExpired()} to check + * expiration status. * *

- * Instances are created via {@link TenantSecurityClient#createCachedDecryptor} or + * Instances are created via {@link TenantSecurityClient#createCachedEncryptor}, + * {@link TenantSecurityClient#createCachedDecryptor}, + * {@link TenantSecurityClient#withCachedEncryptor}, or * {@link TenantSecurityClient#withCachedDecryptor}. See those methods for usage examples. * + * @see TenantSecurityClient#createCachedEncryptor(DocumentMetadata) * @see TenantSecurityClient#createCachedDecryptor(String, DocumentMetadata) - * @see TenantSecurityClient#withCachedDecryptor(String, DocumentMetadata, - * java.util.function.Function) */ -public final class CachedKeyDecryptor implements DocumentDecryptor, Closeable { +public final class CachedKey implements CachedEncryptor, CachedDecryptor { - // Maximum time the decryptor can be used before it expires + // Maximum time the cached key can be used before it expires private static final Duration TIMEOUT = Duration.ofMinutes(1); + private static final String EDEK_MISMATCH_MESSAGE = + "EDEK does not match the cached EDEK. " + + "This CachedKey can only decrypt documents with matching EDEKs."; + // The cached DEK bytes - zeroed on close() private final byte[] dek; - // The EDEK that was used to derive the DEK - used for validation + // The EDEK associated with the cached DEK private final String edek; - // Executor for async field decryption operations + // Executor for async field encryption/decryption operations private final ExecutorService encryptionExecutor; + // Secure random for IV generation during encryption + private final SecureRandom secureRandom; + // For reporting operations on close private final TenantSecurityRequest requestService; private final DocumentMetadata metadata; @@ -62,23 +71,27 @@ public final class CachedKeyDecryptor implements DocumentDecryptor, Closeable { // Flag to track if close() has been called private final AtomicBoolean closed = new AtomicBoolean(false); - // When this decryptor was created - used for timeout enforcement + // When this cached key was created - used for timeout enforcement private final Instant createdAt; + // Count of successful encrypt operations performed + private final AtomicInteger encryptCount = new AtomicInteger(0); + // Count of successful decrypt operations performed - private final AtomicInteger operationCount = new AtomicInteger(0); + private final AtomicInteger decryptCount = new AtomicInteger(0); /** - * Package-private constructor. Use TenantSecurityClient.createCachedDecryptor() to create - * instances. + * Package-private constructor. Use TenantSecurityClient.createCachedEncryptor() or + * TenantSecurityClient.createCachedDecryptor() to create instances. * - * @param dek The unwrapped document encryption key bytes (will be copied) + * @param dek The document encryption key bytes (will be copied) * @param edek The encrypted document encryption key string - * @param encryptionExecutor Executor for async decryption operations + * @param encryptionExecutor Executor for async encryption/decryption operations + * @param secureRandom Secure random for IV generation * @param requestService TSP request service for reporting operations on close * @param metadata Document metadata for reporting operations on close */ - CachedKeyDecryptor(byte[] dek, String edek, ExecutorService encryptionExecutor, + CachedKey(byte[] dek, String edek, ExecutorService encryptionExecutor, SecureRandom secureRandom, TenantSecurityRequest requestService, DocumentMetadata metadata) { if (dek == null || dek.length != 32) { throw new IllegalArgumentException("DEK must be exactly 32 bytes"); @@ -89,6 +102,9 @@ public final class CachedKeyDecryptor implements DocumentDecryptor, Closeable { if (encryptionExecutor == null) { throw new IllegalArgumentException("encryptionExecutor must not be null"); } + if (secureRandom == null) { + throw new IllegalArgumentException("secureRandom must not be null"); + } if (requestService == null) { throw new IllegalArgumentException("requestService must not be null"); } @@ -99,14 +115,14 @@ public final class CachedKeyDecryptor implements DocumentDecryptor, Closeable { this.dek = Arrays.copyOf(dek, dek.length); this.edek = edek; this.encryptionExecutor = encryptionExecutor; + this.secureRandom = secureRandom; this.requestService = requestService; this.metadata = metadata; this.createdAt = Instant.now(); } /** - * Get the EDEK associated with this cached decryptor. Useful for verifying which documents can be - * decrypted with this instance. + * Get the EDEK associated with this cached key. * * @return The EDEK string */ @@ -115,7 +131,7 @@ public String getEdek() { } /** - * Check if this decryptor has been closed. + * Check if this cached key has been closed. * * @return true if close() has been called */ @@ -124,7 +140,7 @@ public boolean isClosed() { } /** - * Check if this decryptor has expired due to timeout. + * Check if this cached key has expired due to timeout. * * @return true if the timeout has elapsed since creation */ @@ -133,44 +149,103 @@ public boolean isExpired() { } /** - * Get the number of successful decrypt operations performed with this decryptor. + * Get the number of successful encrypt operations performed with this cached key. * - * @return The operation count + * @return The encrypt operation count + */ + public int getEncryptCount() { + return encryptCount.get(); + } + + /** + * Get the number of successful decrypt operations performed with this cached key. + * + * @return The decrypt operation count + */ + public int getDecryptCount() { + return decryptCount.get(); + } + + /** + * Get the total number of successful operations (encrypts + decrypts) performed with this cached + * key. + * + * @return The total operation count */ public int getOperationCount() { - return operationCount.get(); + return encryptCount.get() + decryptCount.get(); } /** - * Guard an operation with usability checks and operation counting. Verifies the decryptor is not + * Guard an operation with usability checks and operation counting. Verifies the cached key is not * closed or expired before running the operation, and increments the operation count on success. * * @param operation The operation to perform * @param countOps Extracts the number of successful operations from the result + * @param counter The counter to increment on success + * @param errorCode The error code to use for closed/expired failures */ private CompletableFuture executeAndIncrement(Supplier> operation, - ToIntFunction countOps) { + ToIntFunction countOps, AtomicInteger counter, + TenantSecurityErrorCodes errorCode) { if (closed.get()) { - return CompletableFuture.failedFuture(new TscException( - TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "CachedKeyDecryptor has been closed")); + return CompletableFuture + .failedFuture(new TscException(errorCode, "CachedKey has been closed")); } if (isExpired()) { - return CompletableFuture.failedFuture(new TscException( - TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, "CachedKeyDecryptor has expired")); + return CompletableFuture + .failedFuture(new TscException(errorCode, "CachedKey has expired")); } return operation.get().thenApply(result -> { - operationCount.addAndGet(countOps.applyAsInt(result)); + counter.addAndGet(countOps.applyAsInt(result)); return result; }); } + // === Encrypt operations === + + @Override + public CompletableFuture encrypt(Map document, + DocumentMetadata metadata) { + return executeAndIncrement( + () -> DocumentCryptoOps.encryptFields(document, metadata, dek, edek, encryptionExecutor, + secureRandom), + result -> 1, encryptCount, TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED); + } + + @Override + public CompletableFuture encryptStream(InputStream input, OutputStream output, + DocumentMetadata metadata) { + return executeAndIncrement( + () -> CompletableFuture.supplyAsync( + () -> CryptoUtils.encryptStreamInternal(dek, metadata, input, output, secureRandom) + .join(), + encryptionExecutor).thenApply(v -> new StreamingResponse(edek)), + result -> 1, encryptCount, TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED); + } + + @Override + public CompletableFuture> encryptBatch( + Map> plaintextDocuments, DocumentMetadata metadata) { + return executeAndIncrement(() -> { + ConcurrentMap> ops = new ConcurrentHashMap<>(); + plaintextDocuments.forEach((id, doc) -> ops.put(id, + DocumentCryptoOps.encryptFields(doc, metadata, dek, edek, encryptionExecutor, + secureRandom))); + return CompletableFuture.supplyAsync(() -> DocumentCryptoOps.cryptoOperationToBatchResult(ops, + TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED)); + }, result -> result.getSuccesses().size(), encryptCount, + TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED); + } + + // === Decrypt operations === + private CompletableFuture validateEdekAndDecrypt( EncryptedDocument encryptedDocument) { if (!edek.equals(encryptedDocument.getEdek())) { return CompletableFuture.failedFuture( new TscException(TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, - "EncryptedDocument EDEK does not match the cached EDEK. " - + "This decryptor can only decrypt documents with matching EDEKs.")); + EDEK_MISMATCH_MESSAGE)); } return DocumentCryptoOps.decryptFields(encryptedDocument.getEncryptedFields(), dek, encryptedDocument.getEdek(), encryptionExecutor); @@ -179,7 +254,8 @@ private CompletableFuture validateEdekAndDecrypt( @Override public CompletableFuture decrypt(EncryptedDocument encryptedDocument, DocumentMetadata metadata) { - return executeAndIncrement(() -> validateEdekAndDecrypt(encryptedDocument), result -> 1); + return executeAndIncrement(() -> validateEdekAndDecrypt(encryptedDocument), result -> 1, + decryptCount, TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED); } @Override @@ -189,13 +265,12 @@ public CompletableFuture decryptStream(String edek, InputStream input, Out if (!this.edek.equals(edek)) { return CompletableFuture .failedFuture(new TscException(TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, - "Provided EDEK does not match the cached EDEK. " - + "This decryptor can only decrypt documents with matching EDEKs.")); + EDEK_MISMATCH_MESSAGE)); } return CompletableFuture .supplyAsync(() -> CryptoUtils.decryptStreamInternal(this.dek, input, output).join(), encryptionExecutor); - }, result -> 1); + }, result -> 1, decryptCount, TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED); } @Override @@ -209,8 +284,7 @@ public CompletableFuture> decryptBatch( if (!edek.equals(encDoc.getEdek())) { edekMismatches.put(id, new TscException(TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, - "EncryptedDocument EDEK does not match the cached EDEK. " - + "This decryptor can only decrypt documents with matching EDEKs.")); + EDEK_MISMATCH_MESSAGE)); } else { ops.put(id, DocumentCryptoOps.decryptFields(encDoc.getEncryptedFields(), dek, encDoc.getEdek(), encryptionExecutor)); @@ -225,12 +299,13 @@ public CompletableFuture> decryptBatch( allFailures.putAll(edekMismatches); return new BatchResult<>(result.getSuccesses(), allFailures); }); - }, result -> result.getSuccesses().size()); + }, result -> result.getSuccesses().size(), decryptCount, + TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED); } /** - * Securely zero the DEK bytes, report operations to the TSP, and mark this decryptor as closed. - * After calling close(), all decrypt operations will fail. + * Securely zero the DEK bytes, report operations to the TSP, and mark this cached key as closed. + * After calling close(), all encrypt and decrypt operations will fail. * *

* This method is idempotent - calling it multiple times has no additional effect. @@ -241,9 +316,10 @@ public void close() { // Zero out the DEK bytes for security Arrays.fill(dek, (byte) 0); // Report operations to TSP - int count = operationCount.get(); - if (count > 0) { - requestService.reportOperations(metadata, edek, 0, count); + int encrypts = encryptCount.get(); + int decrypts = decryptCount.get(); + if (encrypts > 0 || decrypts > 0) { + requestService.reportOperations(metadata, edek, encrypts, decrypts); } } } diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptor.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptor.java deleted file mode 100644 index 046002f..0000000 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptor.java +++ /dev/null @@ -1,226 +0,0 @@ -package com.ironcorelabs.tenantsecurity.kms.v1; - -import java.io.Closeable; -import java.io.InputStream; -import java.io.OutputStream; -import java.security.SecureRandom; -import java.time.Duration; -import java.time.Instant; -import java.util.Arrays; -import java.util.Map; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ConcurrentMap; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.function.Supplier; -import java.util.function.ToIntFunction; -import com.ironcorelabs.tenantsecurity.kms.v1.exception.TscException; - -/** - * Holds a cached DEK (Document Encryption Key) for repeated encrypt operations without making - * additional TSP wrap calls. All documents encrypted with this instance will share the same - * DEK/EDEK pair. The DEK is securely zeroed when close() is called. - * - *

- * This class is thread-safe and can be used concurrently for multiple encrypt operations. Once - * closed, all encrypt operations will fail. - * - *

- * Expiration: This encryptor automatically expires after a short time period. Caching a DEK - * for long-term use is not supported as it would undermine the security benefits of key wrapping. - * The encryptor is intended for short-lived batch operations where multiple documents need to be - * encrypted in quick succession with the same key. Use {@link #isExpired()} to check expiration - * status. - * - *

- * Instances are created via {@link TenantSecurityClient#createCachedEncryptor} or - * {@link TenantSecurityClient#withCachedEncryptor}. See those methods for usage examples. - * - * @see TenantSecurityClient#createCachedEncryptor(DocumentMetadata) - * @see TenantSecurityClient#withCachedEncryptor(DocumentMetadata, java.util.function.Function) - */ -public final class CachedKeyEncryptor implements DocumentEncryptor, Closeable { - - // Maximum time the encryptor can be used before it expires - private static final Duration TIMEOUT = Duration.ofMinutes(1); - - // The cached DEK bytes - zeroed on close() - private final byte[] dek; - - // The EDEK associated with the cached DEK - private final String edek; - - // Executor for async field encryption operations - private final ExecutorService encryptionExecutor; - - // Secure random for IV generation during encryption - private final SecureRandom secureRandom; - - // For reporting operations on close - private final TenantSecurityRequest requestService; - private final DocumentMetadata metadata; - - // Flag to track if close() has been called - private final AtomicBoolean closed = new AtomicBoolean(false); - - // When this encryptor was created - used for timeout enforcement - private final Instant createdAt; - - // Count of successful encrypt operations performed - private final AtomicInteger operationCount = new AtomicInteger(0); - - /** - * Package-private constructor. Use TenantSecurityClient.createCachedEncryptor() to create - * instances. - * - * @param dek The document encryption key bytes (will be copied) - * @param edek The encrypted document encryption key string - * @param encryptionExecutor Executor for async encryption operations - * @param secureRandom Secure random for IV generation - * @param requestService TSP request service for reporting operations on close - * @param metadata Document metadata for reporting operations on close - */ - CachedKeyEncryptor(byte[] dek, String edek, ExecutorService encryptionExecutor, - SecureRandom secureRandom, TenantSecurityRequest requestService, DocumentMetadata metadata) { - if (dek == null || dek.length != 32) { - throw new IllegalArgumentException("DEK must be exactly 32 bytes"); - } - if (edek == null || edek.isEmpty()) { - throw new IllegalArgumentException("EDEK must not be null or empty"); - } - if (encryptionExecutor == null) { - throw new IllegalArgumentException("encryptionExecutor must not be null"); - } - if (secureRandom == null) { - throw new IllegalArgumentException("secureRandom must not be null"); - } - if (requestService == null) { - throw new IllegalArgumentException("requestService must not be null"); - } - if (metadata == null) { - throw new IllegalArgumentException("metadata must not be null"); - } - // Copy DEK to prevent external modification - this.dek = Arrays.copyOf(dek, dek.length); - this.edek = edek; - this.encryptionExecutor = encryptionExecutor; - this.secureRandom = secureRandom; - this.requestService = requestService; - this.metadata = metadata; - this.createdAt = Instant.now(); - } - - /** - * Get the EDEK associated with this cached encryptor. - * - * @return The EDEK string - */ - public String getEdek() { - return edek; - } - - /** - * Check if this encryptor has been closed. - * - * @return true if close() has been called - */ - public boolean isClosed() { - return closed.get(); - } - - /** - * Check if this encryptor has expired due to timeout. - * - * @return true if the timeout has elapsed since creation - */ - public boolean isExpired() { - return Duration.between(createdAt, Instant.now()).compareTo(TIMEOUT) > 0; - } - - /** - * Get the number of successful encrypt operations performed with this encryptor. - * - * @return The operation count - */ - public int getOperationCount() { - return operationCount.get(); - } - - /** - * Guard an operation with usability checks and operation counting. Verifies the encryptor is not - * closed or expired before running the operation, and increments the operation count on success. - * - * @param operation The operation to perform - * @param countOps Extracts the number of successful operations from the result - */ - private CompletableFuture executeAndIncrement(Supplier> operation, - ToIntFunction countOps) { - if (closed.get()) { - return CompletableFuture.failedFuture(new TscException( - TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED, "CachedKeyEncryptor has been closed")); - } - if (isExpired()) { - return CompletableFuture.failedFuture(new TscException( - TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED, "CachedKeyEncryptor has expired")); - } - return operation.get().thenApply(result -> { - operationCount.addAndGet(countOps.applyAsInt(result)); - return result; - }); - } - - @Override - public CompletableFuture encrypt(Map document, - DocumentMetadata metadata) { - return executeAndIncrement( - () -> DocumentCryptoOps.encryptFields(document, metadata, dek, edek, encryptionExecutor, - secureRandom), - result -> 1); - } - - @Override - public CompletableFuture encryptStream(InputStream input, OutputStream output, - DocumentMetadata metadata) { - return executeAndIncrement( - () -> CompletableFuture.supplyAsync( - () -> CryptoUtils.encryptStreamInternal(dek, metadata, input, output, secureRandom) - .join(), - encryptionExecutor).thenApply(v -> new StreamingResponse(edek)), - result -> 1); - } - - @Override - public CompletableFuture> encryptBatch( - Map> plaintextDocuments, DocumentMetadata metadata) { - return executeAndIncrement(() -> { - ConcurrentMap> ops = new ConcurrentHashMap<>(); - plaintextDocuments.forEach((id, doc) -> ops.put(id, - DocumentCryptoOps.encryptFields(doc, metadata, dek, edek, encryptionExecutor, - secureRandom))); - return CompletableFuture.supplyAsync(() -> DocumentCryptoOps.cryptoOperationToBatchResult(ops, - TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED)); - }, result -> result.getSuccesses().size()); - } - - /** - * Securely zero the DEK bytes, report operations to the TSP, and mark this encryptor as closed. - * After calling close(), all encrypt operations will fail. - * - *

- * This method is idempotent - calling it multiple times has no additional effect. - */ - @Override - public void close() { - if (closed.compareAndSet(false, true)) { - // Zero out the DEK bytes for security - Arrays.fill(dek, (byte) 0); - // Report operations to TSP - int count = operationCount.get(); - if (count > 0) { - requestService.reportOperations(metadata, edek, count, 0); - } - } - } -} diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyLifecycle.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyLifecycle.java new file mode 100644 index 0000000..f490aab --- /dev/null +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyLifecycle.java @@ -0,0 +1,45 @@ +package com.ironcorelabs.tenantsecurity.kms.v1; + +import java.io.Closeable; + +/** + * Common lifecycle methods shared by {@link CachedEncryptor} and {@link CachedDecryptor}. Provides + * access to the cached EDEK, status checks, operation counting, and resource cleanup. + */ +public interface CachedKeyLifecycle extends Closeable { + + /** + * Get the EDEK associated with this cached key. + * + * @return The EDEK string + */ + String getEdek(); + + /** + * Check if this cached key has been closed. + * + * @return true if close() has been called + */ + boolean isClosed(); + + /** + * Check if this cached key has expired due to timeout. + * + * @return true if the timeout has elapsed since creation + */ + boolean isExpired(); + + /** + * Get the total number of successful operations performed with this cached key. + * + * @return The total operation count + */ + int getOperationCount(); + + /** + * Securely zero the DEK and release resources. After calling close(), all operations will fail. + * This override narrows the Closeable contract to not throw IOException. + */ + @Override + void close(); +} diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentDecryptor.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentDecryptor.java index b92d444..5e60814 100644 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentDecryptor.java +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentDecryptor.java @@ -7,8 +7,8 @@ /** * Interface for document decryption capabilities. Implemented by both TenantSecurityClient (for - * standard decrypt operations that unwrap the EDEK each time) and CachedKeyDecryptor (for repeated - * decrypts using a cached DEK). + * standard decrypt operations that unwrap the EDEK each time) and CachedKey (for repeated decrypts + * using a cached DEK). */ public interface DocumentDecryptor { diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentEncryptor.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentEncryptor.java index 8779903..fc396dc 100644 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentEncryptor.java +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentEncryptor.java @@ -7,8 +7,8 @@ /** * Interface for document encryption capabilities. Implemented by both TenantSecurityClient (for - * standard encrypt operations that wrap a new DEK each time) and CachedKeyEncryptor (for repeated - * encrypts using a cached DEK). + * standard encrypt operations that wrap a new DEK each time) and CachedKey (for repeated encrypts + * using a cached DEK). */ public interface DocumentEncryptor { diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java index d15757b..78e0d7e 100644 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java @@ -488,20 +488,54 @@ public CompletableFuture decrypt(EncryptedDocument encryptedD encryptedDocument.getEdek(), encryptionExecutor)); } + // === Private helpers for creating CachedKey instances === + + private CompletableFuture newCachedKeyFromUnwrap(String edek, + DocumentMetadata metadata) { + return this.encryptionService.unwrapKey(edek, metadata).thenApply(dekBytes -> { + CachedKey cachedKey = new CachedKey(dekBytes, edek, this.encryptionExecutor, + this.secureRandom, this.encryptionService, metadata); + Arrays.fill(dekBytes, (byte) 0); + return cachedKey; + }); + } + + private CompletableFuture newCachedKeyFromWrap(DocumentMetadata metadata) { + return this.encryptionService.wrapKey(metadata).thenApply(wrappedKey -> { + byte[] dekBytes = wrappedKey.getDekBytes(); + CachedKey cachedKey = new CachedKey(dekBytes, wrappedKey.getEdek(), + this.encryptionExecutor, this.secureRandom, this.encryptionService, metadata); + Arrays.fill(dekBytes, (byte) 0); + return cachedKey; + }); + } + + /** + * Execute an operation on a cached resource with automatic lifecycle management. The resource is + * closed (and DEK zeroed) when the operation completes, whether successfully or with an error. + */ + private CompletableFuture withCachedResource( + CompletableFuture resource, Function> operation) { + return resource.thenCompose( + k -> operation.apply(k).whenComplete((result, error) -> k.close())); + } + + // === Cached decryptor factory methods === + /** - * Create a CachedKeyDecryptor for repeated decrypt operations using the same DEK. This unwraps - * the EDEK once and caches the resulting DEK for subsequent decrypts. + * Create a CachedDecryptor for repeated decrypt operations using the same DEK. This unwraps the + * EDEK once and caches the resulting DEK for subsequent decrypts. * *

* Use this when you need to decrypt multiple documents that share the same EDEK, to avoid * repeated TSP unwrap calls. * *

- * The returned decryptor implements AutoCloseable and should be used with try-with-resources to + * The returned CachedDecryptor implements Closeable and should be used with try-with-resources to * ensure the DEK is securely zeroed when done: * * 

-   * try (CachedKeyDecryptor decryptor = client.createCachedDecryptor(edek, metadata).get()) {
+   * try (CachedDecryptor decryptor = client.createCachedDecryptor(edek, metadata).get()) {
    *   PlaintextDocument doc1 = decryptor.decrypt(encDoc1, metadata).get();
    *   PlaintextDocument doc2 = decryptor.decrypt(encDoc2, metadata).get();
    * }
@@ -509,35 +543,30 @@ public CompletableFuture decrypt(EncryptedDocument encryptedD
    *
    * @param edek The encrypted document encryption key to unwrap
    * @param metadata Metadata for the unwrap operation
-   * @return CompletableFuture resolving to a CachedKeyDecryptor
+   * @return CompletableFuture resolving to a CachedDecryptor
    */
-  public CompletableFuture createCachedDecryptor(String edek,
+  public CompletableFuture createCachedDecryptor(String edek,
       DocumentMetadata metadata) {
-    return this.encryptionService.unwrapKey(edek, metadata).thenApply(dekBytes -> {
-      CachedKeyDecryptor decryptor = new CachedKeyDecryptor(dekBytes, edek, this.encryptionExecutor,
-          this.encryptionService, metadata);
-      Arrays.fill(dekBytes, (byte) 0);
-      return decryptor;
-    });
+    return newCachedKeyFromUnwrap(edek, metadata).thenApply(k -> k);
   }
 
   /**
-   * Create a CachedKeyDecryptor from an existing EncryptedDocument. Convenience method that
-   * extracts the EDEK from the document.
+   * Create a CachedDecryptor from an existing EncryptedDocument. Convenience method that extracts
+   * the EDEK from the document.
    *
    * @param encryptedDocument The encrypted document whose EDEK should be unwrapped
    * @param metadata Metadata for the unwrap operation
-   * @return CompletableFuture resolving to a CachedKeyDecryptor
+   * @return CompletableFuture resolving to a CachedDecryptor
    */
-  public CompletableFuture createCachedDecryptor(
+  public CompletableFuture createCachedDecryptor(
       EncryptedDocument encryptedDocument, DocumentMetadata metadata) {
     return createCachedDecryptor(encryptedDocument.getEdek(), metadata);
   }
 
   /**
-   * Execute an operation using a CachedKeyDecryptor with automatic lifecycle management. The
-   * decryptor is automatically closed (and DEK zeroed) when the operation completes, whether
-   * successfully or with an error.
+   * Execute an operation using a CachedDecryptor with automatic lifecycle management. The cached
+   * key is automatically closed (and DEK zeroed) when the operation completes, whether successfully
+   * or with an error.
    *
    * 

    * This is the recommended pattern for using cached decryptors with CompletableFuture composition:
@@ -550,33 +579,34 @@ public CompletableFuture createCachedDecryptor(
    * @param  The type returned by the operation
    * @param edek The encrypted document encryption key to unwrap
    * @param metadata Metadata for the unwrap operation
-   * @param operation Function that takes the decryptor and returns a CompletableFuture
+   * @param operation Function that takes the CachedDecryptor and returns a CompletableFuture
    * @return CompletableFuture resolving to the operation's result
    */
   public  CompletableFuture withCachedDecryptor(String edek, DocumentMetadata metadata,
-      Function> operation) {
-    return createCachedDecryptor(edek, metadata).thenCompose(
-        decryptor -> operation.apply(decryptor).whenComplete((result, error) -> decryptor.close()));
+      Function> operation) {
+    return withCachedResource(createCachedDecryptor(edek, metadata), operation);
   }
 
   /**
-   * Execute an operation using a CachedKeyDecryptor with automatic lifecycle management.
-   * Convenience method that extracts the EDEK from the document.
+   * Execute an operation using a CachedDecryptor with automatic lifecycle management. Convenience
+   * method that extracts the EDEK from the document.
    *
    * @param  The type returned by the operation
    * @param encryptedDocument The encrypted document whose EDEK should be unwrapped
    * @param metadata Metadata for the unwrap operation
-   * @param operation Function that takes the decryptor and returns a CompletableFuture
+   * @param operation Function that takes the CachedDecryptor and returns a CompletableFuture
    * @return CompletableFuture resolving to the operation's result
    */
   public  CompletableFuture withCachedDecryptor(EncryptedDocument encryptedDocument,
-      DocumentMetadata metadata, Function> operation) {
+      DocumentMetadata metadata, Function> operation) {
     return withCachedDecryptor(encryptedDocument.getEdek(), metadata, operation);
   }
 
+  // === Cached encryptor factory methods ===
+
   /**
-   * Create a CachedKeyEncryptor for repeated encrypt operations using the same DEK. This wraps a
-   * new key once and caches the resulting DEK/EDEK pair for subsequent encrypts. All documents
+   * Create a CachedEncryptor for repeated encrypt operations using the same DEK. This wraps a new
+   * key once and caches the resulting DEK/EDEK pair for subsequent encrypts. All documents
    * encrypted with this instance will share the same DEK/EDEK pair.
    *
    * 

@@ -584,33 +614,27 @@ public  CompletableFuture withCachedDecryptor(EncryptedDocument encryptedD
    * to avoid repeated TSP wrap calls.
    *
    * 

-   * The returned encryptor implements AutoCloseable and should be used with try-with-resources to
+   * The returned CachedEncryptor implements Closeable and should be used with try-with-resources to
    * ensure the DEK is securely zeroed when done:
    *
    * 
-   * try (CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get()) {
+   * try (CachedEncryptor encryptor = client.createCachedEncryptor(metadata).get()) {
    *   EncryptedDocument enc1 = encryptor.encrypt(doc1, metadata).get();
    *   EncryptedDocument enc2 = encryptor.encrypt(doc2, metadata).get();
    * }
    * 

    *
    * @param metadata Metadata for the wrap operation
-   * @return CompletableFuture resolving to a CachedKeyEncryptor
+   * @return CompletableFuture resolving to a CachedEncryptor
    */
-  public CompletableFuture createCachedEncryptor(DocumentMetadata metadata) {
-    return this.encryptionService.wrapKey(metadata).thenApply(wrappedKey -> {
-      byte[] dekBytes = wrappedKey.getDekBytes();
-      CachedKeyEncryptor encryptor = new CachedKeyEncryptor(dekBytes, wrappedKey.getEdek(),
-          this.encryptionExecutor, this.secureRandom, this.encryptionService, metadata);
-      Arrays.fill(dekBytes, (byte) 0);
-      return encryptor;
-    });
+  public CompletableFuture createCachedEncryptor(DocumentMetadata metadata) {
+    return newCachedKeyFromWrap(metadata).thenApply(k -> k);
   }
 
   /**
-   * Execute an operation using a CachedKeyEncryptor with automatic lifecycle management. The
-   * encryptor is automatically closed (and DEK zeroed) when the operation completes, whether
-   * successfully or with an error.
+   * Execute an operation using a CachedEncryptor with automatic lifecycle management. The cached
+   * key is automatically closed (and DEK zeroed) when the operation completes, whether successfully
+   * or with an error.
    *
    * 

    * This is the recommended pattern for using cached encryptors with CompletableFuture composition:
@@ -622,13 +646,83 @@ public CompletableFuture createCachedEncryptor(DocumentMetad
    *
    * @param  The type returned by the operation
    * @param metadata Metadata for the wrap operation
-   * @param operation Function that takes the encryptor and returns a CompletableFuture
+   * @param operation Function that takes the CachedEncryptor and returns a CompletableFuture
    * @return CompletableFuture resolving to the operation's result
    */
   public  CompletableFuture withCachedEncryptor(DocumentMetadata metadata,
-      Function> operation) {
-    return createCachedEncryptor(metadata).thenCompose(
-        encryptor -> operation.apply(encryptor).whenComplete((result, error) -> encryptor.close()));
+      Function> operation) {
+    return withCachedResource(createCachedEncryptor(metadata), operation);
+  }
+
+  // === CachedKey factory methods (full encrypt + decrypt access) ===
+
+  /**
+   * Create a CachedKey for both encrypt and decrypt operations. Wraps a new key and caches the
+   * resulting DEK/EDEK pair.
+   *
+   * 

+   * Use this when you need both encrypt and decrypt capabilities with the same cached key. If you
+   * only need encrypt or decrypt, prefer {@link #createCachedEncryptor(DocumentMetadata)} or
+   * {@link #createCachedDecryptor(String, DocumentMetadata)} for narrower type safety.
+   *
+   * @param metadata Metadata for the wrap operation
+   * @return CompletableFuture resolving to a CachedKey
+   */
+  public CompletableFuture createCachedKey(DocumentMetadata metadata) {
+    return newCachedKeyFromWrap(metadata);
+  }
+
+  /**
+   * Create a CachedKey for both encrypt and decrypt operations by unwrapping an existing EDEK.
+   *
+   * @param edek The encrypted document encryption key to unwrap
+   * @param metadata Metadata for the unwrap operation
+   * @return CompletableFuture resolving to a CachedKey
+   */
+  public CompletableFuture createCachedKey(String edek, DocumentMetadata metadata) {
+    return newCachedKeyFromUnwrap(edek, metadata);
+  }
+
+  /**
+   * Create a CachedKey for both encrypt and decrypt operations from an existing EncryptedDocument.
+   * Convenience method that extracts the EDEK from the document.
+   *
+   * @param encryptedDocument The encrypted document whose EDEK should be unwrapped
+   * @param metadata Metadata for the unwrap operation
+   * @return CompletableFuture resolving to a CachedKey
+   */
+  public CompletableFuture createCachedKey(EncryptedDocument encryptedDocument,
+      DocumentMetadata metadata) {
+    return createCachedKey(encryptedDocument.getEdek(), metadata);
+  }
+
+  /**
+   * Execute an operation using a CachedKey with automatic lifecycle management. Wraps a new key
+   * and provides full encrypt + decrypt access.
+   *
+   * @param  The type returned by the operation
+   * @param metadata Metadata for the wrap operation
+   * @param operation Function that takes the CachedKey and returns a CompletableFuture
+   * @return CompletableFuture resolving to the operation's result
+   */
+  public  CompletableFuture withCachedKey(DocumentMetadata metadata,
+      Function> operation) {
+    return withCachedResource(createCachedKey(metadata), operation);
+  }
+
+  /**
+   * Execute an operation using a CachedKey with automatic lifecycle management. Unwraps an
+   * existing EDEK and provides full encrypt + decrypt access.
+   *
+   * @param  The type returned by the operation
+   * @param edek The encrypted document encryption key to unwrap
+   * @param metadata Metadata for the unwrap operation
+   * @param operation Function that takes the CachedKey and returns a CompletableFuture
+   * @return CompletableFuture resolving to the operation's result
+   */
+  public  CompletableFuture withCachedKey(String edek, DocumentMetadata metadata,
+      Function> operation) {
+    return withCachedResource(createCachedKey(edek, metadata), operation);
   }
 
   /**
diff --git a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptorTest.java b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptorTest.java
deleted file mode 100644
index 7d4311a..0000000
--- a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyEncryptorTest.java
+++ /dev/null
@@ -1,267 +0,0 @@
-package com.ironcorelabs.tenantsecurity.kms.v1;
-
-import static org.testng.Assert.assertEquals;
-import static org.testng.Assert.assertFalse;
-import static org.testng.Assert.assertTrue;
-import static org.testng.Assert.fail;
-
-import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
-import java.lang.reflect.Field;
-import java.security.SecureRandom;
-import java.util.Arrays;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.concurrent.CompletionException;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import org.testng.annotations.AfterClass;
-import org.testng.annotations.BeforeClass;
-import org.testng.annotations.Test;
-import com.ironcorelabs.tenantsecurity.kms.v1.exception.TscException;
-
-@Test(groups = {"unit"})
-public class CachedKeyEncryptorTest {
-
-  private ExecutorService executor;
-  private SecureRandom secureRandom;
-  private TenantSecurityRequest encryptionService;
-  private static final String TEST_EDEK = "test-edek-base64-string";
-  private DocumentMetadata metadata =
-      new DocumentMetadata("tenantId", "requestingUserOrServiceId", "dataLabel");
-
-  @BeforeClass
-  public void setup() {
-    executor = Executors.newFixedThreadPool(2);
-    secureRandom = new SecureRandom();
-    // This endpoint doesn't exist, so we won't call `close` on the cached encryptor to avoid the
-    // report-operations request
-    encryptionService = new TenantSecurityRequest("http://localhost:0", "test-api-key", 1, 1000);
-  }
-
-  @AfterClass
-  public void teardown() {
-    if (executor != null) {
-      executor.shutdown();
-    }
-    if (encryptionService != null) {
-      try {
-        encryptionService.close();
-      } catch (Exception e) {
-        // ignore
-      }
-    }
-  }
-
-  private byte[] createValidDek() {
-    byte[] dek = new byte[32];
-    Arrays.fill(dek, (byte) 0x42);
-    return dek;
-  }
-
-  private CachedKeyEncryptor createEncryptor() {
-    return new CachedKeyEncryptor(createValidDek(), TEST_EDEK, executor, secureRandom,
-        encryptionService, metadata);
-  }
-
-  // Constructor validation tests
-
-  @SuppressWarnings("resource")
-  public void constructorRejectNullDek() {
-    try {
-      new CachedKeyEncryptor(null, TEST_EDEK, executor, secureRandom, encryptionService, metadata);
-      fail("Should have thrown IllegalArgumentException");
-    } catch (IllegalArgumentException e) {
-      assertTrue(e.getMessage().contains("DEK must be exactly 32 bytes"));
-    }
-  }
-
-  @SuppressWarnings("resource")
-  public void constructorRejectWrongSizeDek() {
-    byte[] shortDek = new byte[16];
-    try {
-      new CachedKeyEncryptor(shortDek, TEST_EDEK, executor, secureRandom, encryptionService,
-          metadata);
-      fail("Should have thrown IllegalArgumentException");
-    } catch (IllegalArgumentException e) {
-      assertTrue(e.getMessage().contains("DEK must be exactly 32 bytes"));
-    }
-  }
-
-  @SuppressWarnings("resource")
-  public void constructorRejectNullEdek() {
-    try {
-      new CachedKeyEncryptor(createValidDek(), null, executor, secureRandom, encryptionService,
-          metadata);
-      fail("Should have thrown IllegalArgumentException");
-    } catch (IllegalArgumentException e) {
-      assertTrue(e.getMessage().contains("EDEK must not be null or empty"));
-    }
-  }
-
-  @SuppressWarnings("resource")
-  public void constructorRejectEmptyEdek() {
-    try {
-      new CachedKeyEncryptor(createValidDek(), "", executor, secureRandom, encryptionService,
-          metadata);
-      fail("Should have thrown IllegalArgumentException");
-    } catch (IllegalArgumentException e) {
-      assertTrue(e.getMessage().contains("EDEK must not be null or empty"));
-    }
-  }
-
-  @SuppressWarnings("resource")
-  public void constructorRejectNullExecutor() {
-    try {
-      new CachedKeyEncryptor(createValidDek(), TEST_EDEK, null, secureRandom, encryptionService,
-          metadata);
-      fail("Should have thrown IllegalArgumentException");
-    } catch (IllegalArgumentException e) {
-      assertTrue(e.getMessage().contains("encryptionExecutor must not be null"));
-    }
-  }
-
-  @SuppressWarnings("resource")
-  public void constructorRejectNullSecureRandom() {
-    try {
-      new CachedKeyEncryptor(createValidDek(), TEST_EDEK, executor, null, encryptionService,
-          metadata);
-      fail("Should have thrown IllegalArgumentException");
-    } catch (IllegalArgumentException e) {
-      assertTrue(e.getMessage().contains("secureRandom must not be null"));
-    }
-  }
-
-  @SuppressWarnings("resource")
-  public void constructorRejectNullEncryptionService() {
-    try {
-      new CachedKeyEncryptor(createValidDek(), TEST_EDEK, executor, secureRandom, null, metadata);
-      fail("Should have thrown IllegalArgumentException");
-    } catch (IllegalArgumentException e) {
-      assertTrue(e.getMessage().contains("requestService must not be null"));
-    }
-  }
-
-  @SuppressWarnings("resource")
-  public void constructorRejectNullMetadata() {
-    try {
-      new CachedKeyEncryptor(createValidDek(), TEST_EDEK, executor, secureRandom, encryptionService,
-          null);
-      fail("Should have thrown IllegalArgumentException");
-    } catch (IllegalArgumentException e) {
-      assertTrue(e.getMessage().contains("metadata must not be null"));
-    }
-  }
-
-  // Getter tests
-
-  public void getEdekReturnsCorrectValue() {
-    CachedKeyEncryptor encryptor = createEncryptor();
-    assertEquals(encryptor.getEdek(), TEST_EDEK);
-    encryptor.close();
-  }
-
-  public void isClosedReturnsFalseInitially() {
-    CachedKeyEncryptor encryptor = createEncryptor();
-    assertFalse(encryptor.isClosed());
-    encryptor.close();
-  }
-
-  public void isClosedReturnsTrueAfterClose() {
-    CachedKeyEncryptor encryptor = createEncryptor();
-    encryptor.close();
-    assertTrue(encryptor.isClosed());
-  }
-
-  // Close tests
-
-  public void closeIsIdempotent() {
-    CachedKeyEncryptor encryptor = createEncryptor();
-    encryptor.close();
-    assertTrue(encryptor.isClosed());
-    // Should not throw
-    encryptor.close();
-    encryptor.close();
-    assertTrue(encryptor.isClosed());
-  }
-
-  // Operation count tests
-
-  public void operationCountStartsAtZero() {
-    CachedKeyEncryptor encryptor = createEncryptor();
-    assertEquals(encryptor.getOperationCount(), 0);
-    encryptor.close();
-  }
-
-  // Encrypt validation tests
-
-  public void encryptFailsWhenClosed() {
-    CachedKeyEncryptor encryptor = createEncryptor();
-    encryptor.close();
-
-    try {
-      encryptor.encrypt(java.util.Collections.emptyMap(), metadata).join();
-      fail("Should have thrown CompletionException");
-    } catch (CompletionException e) {
-      assertTrue(e.getCause() instanceof TscException);
-      assertTrue(e.getCause().getMessage().contains("CachedKeyEncryptor has been closed"));
-    }
-  }
-
-  public void encryptStreamFailsWhenClosed() {
-    CachedKeyEncryptor encryptor = createEncryptor();
-    encryptor.close();
-
-    ByteArrayInputStream input = new ByteArrayInputStream(new byte[0]);
-    ByteArrayOutputStream output = new ByteArrayOutputStream();
-
-    try {
-      encryptor.encryptStream(input, output, metadata).join();
-      fail("Should have thrown CompletionException");
-    } catch (CompletionException e) {
-      assertTrue(e.getCause() instanceof TscException);
-      assertTrue(e.getCause().getMessage().contains("CachedKeyEncryptor has been closed"));
-    }
-  }
-
-  // encryptBatch validation tests
-
-  public void encryptBatchFailsWhenClosed() {
-    CachedKeyEncryptor encryptor = createEncryptor();
-    encryptor.close();
-
-    Map> docs = new HashMap<>();
-    docs.put("doc1", java.util.Collections.singletonMap("field", new byte[] {1, 2, 3}));
-
-    try {
-      encryptor.encryptBatch(docs, metadata).join();
-      fail("Should have thrown CompletionException");
-    } catch (CompletionException e) {
-      assertTrue(e.getCause() instanceof TscException);
-      assertTrue(e.getCause().getMessage().contains("CachedKeyEncryptor has been closed"));
-    }
-  }
-
-  // DEK copying test
-
-  public void constructorCopiesDekToPreventExternalModification() throws Exception {
-    byte[] originalDek = createValidDek();
-    CachedKeyEncryptor encryptor = new CachedKeyEncryptor(originalDek, TEST_EDEK, executor,
-        secureRandom, encryptionService, metadata);
-
-    // Modify the original array
-    Arrays.fill(originalDek, (byte) 0x00);
-
-    // Use reflection to verify internal DEK still has original values
-    Field dekField = CachedKeyEncryptor.class.getDeclaredField("dek");
-    dekField.setAccessible(true);
-    byte[] internalDek = (byte[]) dekField.get(encryptor);
-
-    // Internal DEK should still be 0x42, not 0x00
-    for (byte b : internalDek) {
-      assertEquals(b, (byte) 0x42, "Internal DEK should not be affected by external modification");
-    }
-
-    encryptor.close();
-  }
-}
diff --git a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyOpsRoundTrip.java b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyOpsRoundTrip.java
index 0438a21..3896674 100644
--- a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyOpsRoundTrip.java
+++ b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyOpsRoundTrip.java
@@ -44,7 +44,7 @@ private Map getDocumentFields() throws Exception {
     return documentMap;
   }
 
-  // === CachedKeyEncryptor tests ===
+  // === CachedKey encrypt tests ===
 
   public void cachedEncryptorRoundTrip() throws Exception {
     DocumentMetadata metadata = getMetadata();
@@ -54,12 +54,12 @@ public void cachedEncryptorRoundTrip() throws Exception {
 
     TenantSecurityClient client = getClient().get();
 
-    try (CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get()) {
+    try (CachedEncryptor encryptor = client.createCachedEncryptor(metadata).get()) {
       assertFalse(encryptor.isClosed());
       assertFalse(encryptor.isExpired());
       assertEquals(encryptor.getOperationCount(), 0);
 
-      // Encrypt two documents with the cached encryptor
+      // Encrypt two documents with the cached key
       EncryptedDocument enc1 = encryptor.encrypt(doc1, metadata).get();
       EncryptedDocument enc2 = encryptor.encrypt(doc2, metadata).get();
 
@@ -90,7 +90,7 @@ public void cachedEncryptorWithPattern() throws Exception {
 
     // Use the withCachedEncryptor pattern for automatic lifecycle management
     EncryptedDocument encrypted =
-        client.withCachedEncryptor(metadata, encryptor -> encryptor.encrypt(doc, metadata)).get();
+        client.withCachedEncryptor(metadata, cachedKey -> cachedKey.encrypt(doc, metadata)).get();
 
     // Verify the encrypted document can be decrypted
     PlaintextDocument decrypted = client.decrypt(encrypted, metadata).get();
@@ -101,7 +101,7 @@ public void cachedEncryptorWithPattern() throws Exception {
     client.close();
   }
 
-  // === CachedKeyDecryptor tests ===
+  // === CachedKey decrypt tests ===
 
   public void cachedDecryptorRoundTrip() throws Exception {
     DocumentMetadata metadata = getMetadata();
@@ -111,17 +111,16 @@ public void cachedDecryptorRoundTrip() throws Exception {
 
     TenantSecurityClient client = getClient().get();
 
-    // Encrypt two documents with the same key (using cached encryptor)
+    // Encrypt two documents with the same key (using cached key)
     EncryptedDocument enc1;
     EncryptedDocument enc2;
-    try (CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get()) {
+    try (CachedEncryptor encryptor = client.createCachedEncryptor(metadata).get()) {
       enc1 = encryptor.encrypt(doc1, metadata).get();
       enc2 = encryptor.encrypt(doc2, metadata).get();
     }
 
     // Decrypt both using a cached decryptor (single unwrap call)
-    try (CachedKeyDecryptor decryptor =
-        client.createCachedDecryptor(enc1.getEdek(), metadata).get()) {
+    try (CachedDecryptor decryptor = client.createCachedDecryptor(enc1.getEdek(), metadata).get()) {
       assertFalse(decryptor.isClosed());
       assertFalse(decryptor.isExpired());
       assertEquals(decryptor.getOperationCount(), 0);
@@ -149,7 +148,7 @@ public void cachedDecryptorFromEncryptedDocument() throws Exception {
     EncryptedDocument encrypted = client.encrypt(doc, metadata).get();
 
     // Create decryptor from EncryptedDocument directly
-    try (CachedKeyDecryptor decryptor = client.createCachedDecryptor(encrypted, metadata).get()) {
+    try (CachedDecryptor decryptor = client.createCachedDecryptor(encrypted, metadata).get()) {
       PlaintextDocument decrypted = decryptor.decrypt(encrypted, metadata).get();
       assertEqualBytes(decrypted.getDecryptedFields().get("field1"), doc.get("field1"));
       assertEqualBytes(decrypted.getDecryptedFields().get("field2"), doc.get("field2"));
@@ -169,7 +168,7 @@ public void cachedDecryptorWithPattern() throws Exception {
 
     // Use the withCachedDecryptor pattern for automatic lifecycle management
     PlaintextDocument decrypted = client.withCachedDecryptor(encrypted, metadata,
-        decryptor -> decryptor.decrypt(encrypted, metadata)).get();
+        cachedKey -> cachedKey.decrypt(encrypted, metadata)).get();
 
     assertEqualBytes(decrypted.getDecryptedFields().get("field1"), doc.get("field1"));
     assertEqualBytes(decrypted.getDecryptedFields().get("field2"), doc.get("field2"));
@@ -189,12 +188,12 @@ public void cachedEncryptorStreamRoundTrip() throws Exception {
     ByteArrayOutputStream encryptedOutput = new ByteArrayOutputStream();
     String edek;
 
-    try (CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get()) {
+    try (CachedEncryptor encryptor = client.createCachedEncryptor(metadata).get()) {
       ByteArrayInputStream input = new ByteArrayInputStream(plaintext);
       StreamingResponse response = encryptor.encryptStream(input, encryptedOutput, metadata).get();
       edek = response.getEdek();
       assertEquals(encryptor.getOperationCount(), 1);
-      // EDEK from streaming response should match the encryptor's EDEK
+      // EDEK from streaming response should match the cached key's EDEK
       assertEquals(edek, encryptor.getEdek());
     }
 
@@ -220,8 +219,8 @@ public void cachedDecryptorStreamRoundTrip() throws Exception {
     StreamingResponse encResponse = client.encryptStream(input, encryptedOutput, metadata).get();
     String edek = encResponse.getEdek();
 
-    // Decrypt with cached decryptor
-    try (CachedKeyDecryptor decryptor = client.createCachedDecryptor(edek, metadata).get()) {
+    // Decrypt with cached key
+    try (CachedDecryptor decryptor = client.createCachedDecryptor(edek, metadata).get()) {
       ByteArrayInputStream encryptedInput = new ByteArrayInputStream(encryptedOutput.toByteArray());
       ByteArrayOutputStream decryptedOutput = new ByteArrayOutputStream();
       decryptor.decryptStream(edek, encryptedInput, decryptedOutput, metadata).get();
@@ -243,17 +242,16 @@ public void cachedEncryptToCachedDecryptRoundTrip() throws Exception {
 
     TenantSecurityClient client = getClient().get();
 
-    // Encrypt multiple docs with cached encryptor
+    // Encrypt multiple docs with cached key
     EncryptedDocument enc1;
     EncryptedDocument enc2;
-    try (CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get()) {
+    try (CachedEncryptor encryptor = client.createCachedEncryptor(metadata).get()) {
       enc1 = encryptor.encrypt(doc1, metadata).get();
       enc2 = encryptor.encrypt(doc2, metadata).get();
     }
 
-    // Decrypt all with cached decryptor (one unwrap call for all)
-    try (CachedKeyDecryptor decryptor =
-        client.createCachedDecryptor(enc1.getEdek(), metadata).get()) {
+    // Decrypt all with cached key (one unwrap call for all)
+    try (CachedDecryptor decryptor = client.createCachedDecryptor(enc1.getEdek(), metadata).get()) {
       PlaintextDocument dec1 = decryptor.decrypt(enc1, metadata).get();
       PlaintextDocument dec2 = decryptor.decrypt(enc2, metadata).get();
 
@@ -272,17 +270,17 @@ public void cachedStreamEncryptToCachedStreamDecryptRoundTrip() throws Exception
 
     TenantSecurityClient client = getClient().get();
 
-    // Encrypt stream with cached encryptor
+    // Encrypt stream with cached key
     ByteArrayOutputStream encryptedOutput = new ByteArrayOutputStream();
     String edek;
-    try (CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get()) {
+    try (CachedEncryptor encryptor = client.createCachedEncryptor(metadata).get()) {
       ByteArrayInputStream input = new ByteArrayInputStream(plaintext);
       StreamingResponse response = encryptor.encryptStream(input, encryptedOutput, metadata).get();
       edek = response.getEdek();
     }
 
-    // Decrypt stream with cached decryptor
-    try (CachedKeyDecryptor decryptor = client.createCachedDecryptor(edek, metadata).get()) {
+    // Decrypt stream with cached key
+    try (CachedDecryptor decryptor = client.createCachedDecryptor(edek, metadata).get()) {
       ByteArrayInputStream encryptedInput = new ByteArrayInputStream(encryptedOutput.toByteArray());
       ByteArrayOutputStream decryptedOutput = new ByteArrayOutputStream();
       decryptor.decryptStream(edek, encryptedInput, decryptedOutput, metadata).get();
@@ -292,7 +290,43 @@ public void cachedStreamEncryptToCachedStreamDecryptRoundTrip() throws Exception
     client.close();
   }
 
-  // === Encryptor close behavior ===
+  // === Single CachedKey for both encrypt and decrypt ===
+
+  public void singleCachedKeyEncryptAndDecrypt() throws Exception {
+    DocumentMetadata metadata = getMetadata();
+    Map doc1 = getDocumentFields();
+    Map doc2 = new HashMap<>();
+    doc2.put("other", "Other document data".getBytes("UTF-8"));
+
+    TenantSecurityClient client = getClient().get();
+
+    // Use createCachedKey for full encrypt + decrypt access
+    try (CachedKey cachedKey = client.createCachedKey(metadata).get()) {
+      // Encrypt
+      EncryptedDocument enc1 = cachedKey.encrypt(doc1, metadata).get();
+      EncryptedDocument enc2 = cachedKey.encrypt(doc2, metadata).get();
+
+      assertEquals(cachedKey.getEncryptCount(), 2);
+      assertEquals(cachedKey.getDecryptCount(), 0);
+
+      // Decrypt with the same CachedKey
+      PlaintextDocument dec1 = cachedKey.decrypt(enc1, metadata).get();
+      PlaintextDocument dec2 = cachedKey.decrypt(enc2, metadata).get();
+
+      assertEquals(cachedKey.getEncryptCount(), 2);
+      assertEquals(cachedKey.getDecryptCount(), 2);
+      assertEquals(cachedKey.getOperationCount(), 4);
+
+      assertEqualBytes(dec1.getDecryptedFields().get("field1"), doc1.get("field1"));
+      assertEqualBytes(dec1.getDecryptedFields().get("field2"), doc1.get("field2"));
+      assertEqualBytes(dec1.getDecryptedFields().get("field3"), doc1.get("field3"));
+      assertEqualBytes(dec2.getDecryptedFields().get("other"), doc2.get("other"));
+    }
+
+    client.close();
+  }
+
+  // === Close behavior ===
 
   public void cachedEncryptorRejectsAfterClose() throws Exception {
     DocumentMetadata metadata = getMetadata();
@@ -300,7 +334,7 @@ public void cachedEncryptorRejectsAfterClose() throws Exception {
 
     TenantSecurityClient client = getClient().get();
 
-    CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get();
+    CachedEncryptor encryptor = client.createCachedEncryptor(metadata).get();
     // Encrypt once to verify it works
     encryptor.encrypt(doc, metadata).get();
     assertEquals(encryptor.getOperationCount(), 1);
@@ -327,8 +361,7 @@ public void cachedDecryptorRejectsAfterClose() throws Exception {
 
     EncryptedDocument encrypted = client.encrypt(doc, metadata).get();
 
-    CachedKeyDecryptor decryptor =
-        client.createCachedDecryptor(encrypted.getEdek(), metadata).get();
+    CachedDecryptor decryptor = client.createCachedDecryptor(encrypted.getEdek(), metadata).get();
     // Decrypt once to verify it works
     decryptor.decrypt(encrypted, metadata).get();
     assertEquals(decryptor.getOperationCount(), 1);
@@ -366,7 +399,7 @@ public void cachedEncryptorBatchRoundTrip() throws Exception {
     TenantSecurityClient client = getClient().get();
 
     BatchResult encResult;
-    try (CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get()) {
+    try (CachedEncryptor encryptor = client.createCachedEncryptor(metadata).get()) {
       encResult = encryptor.encryptBatch(docs, metadata).get();
       assertEquals(encryptor.getOperationCount(), 3);
     }
@@ -404,22 +437,21 @@ public void cachedDecryptorBatchRoundTrip() throws Exception {
 
     TenantSecurityClient client = getClient().get();
 
-    // Encrypt all 3 with cached encryptor (same key)
+    // Encrypt all 3 with cached key (same key)
     EncryptedDocument enc1, enc2, enc3;
-    try (CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get()) {
+    try (CachedEncryptor encryptor = client.createCachedEncryptor(metadata).get()) {
       enc1 = encryptor.encrypt(doc1, metadata).get();
       enc2 = encryptor.encrypt(doc2, metadata).get();
       enc3 = encryptor.encrypt(doc3, metadata).get();
     }
 
-    // Batch decrypt with cached decryptor
+    // Batch decrypt with cached key
     Map encDocs = new HashMap<>();
     encDocs.put("doc1", enc1);
     encDocs.put("doc2", enc2);
     encDocs.put("doc3", enc3);
 
-    try (CachedKeyDecryptor decryptor =
-        client.createCachedDecryptor(enc1.getEdek(), metadata).get()) {
+    try (CachedDecryptor decryptor = client.createCachedDecryptor(enc1.getEdek(), metadata).get()) {
       BatchResult result = decryptor.decryptBatch(encDocs, metadata).get();
       assertEquals(decryptor.getOperationCount(), 3);
 
@@ -451,9 +483,8 @@ public void cachedDecryptorBatchEdekMismatchPartialFailure() throws Exception {
     encDocs.put("match", enc1);
     encDocs.put("mismatch", enc2);
 
-    // Create decryptor for enc1's key
-    try (CachedKeyDecryptor decryptor =
-        client.createCachedDecryptor(enc1.getEdek(), metadata).get()) {
+    // Create cached key for enc1's key
+    try (CachedDecryptor decryptor = client.createCachedDecryptor(enc1.getEdek(), metadata).get()) {
       BatchResult result = decryptor.decryptBatch(encDocs, metadata).get();
 
       // match should succeed
@@ -483,7 +514,7 @@ public void cachedBatchOperationCount() throws Exception {
 
     TenantSecurityClient client = getClient().get();
 
-    try (CachedKeyEncryptor encryptor = client.createCachedEncryptor(metadata).get()) {
+    try (CachedEncryptor encryptor = client.createCachedEncryptor(metadata).get()) {
       assertEquals(encryptor.getOperationCount(), 0);
       encryptor.encryptBatch(docs, metadata).get();
       assertEquals(encryptor.getOperationCount(), 2);
@@ -507,9 +538,8 @@ public void cachedDecryptorRejectsEdekMismatch() throws Exception {
     EncryptedDocument enc1 = client.encrypt(doc, metadata).get();
     EncryptedDocument enc2 = client.encrypt(doc, metadata).get();
 
-    // Create decryptor for enc1's EDEK
-    try (CachedKeyDecryptor decryptor =
-        client.createCachedDecryptor(enc1.getEdek(), metadata).get()) {
+    // Create cached key for enc1's EDEK
+    try (CachedDecryptor decryptor = client.createCachedDecryptor(enc1.getEdek(), metadata).get()) {
       // Decrypting enc1 should work
       decryptor.decrypt(enc1, metadata).get();
 
diff --git a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptorTest.java b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyTest.java
similarity index 60%
rename from src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptorTest.java
rename to src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyTest.java
index 9e4122c..154dce6 100644
--- a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptorTest.java
+++ b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyTest.java
@@ -8,6 +8,7 @@
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.lang.reflect.Field;
+import java.security.SecureRandom;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
@@ -20,9 +21,10 @@
 import com.ironcorelabs.tenantsecurity.kms.v1.exception.TscException;
 
 @Test(groups = {"unit"})
-public class CachedKeyDecryptorTest {
+public class CachedKeyTest {
 
   private ExecutorService executor;
+  private SecureRandom secureRandom;
   private TenantSecurityRequest encryptionService;
   private static final String TEST_EDEK = "test-edek-base64-string";
   private static final String DIFFERENT_EDEK = "different-edek-base64-string";
@@ -32,7 +34,8 @@ public class CachedKeyDecryptorTest {
   @BeforeClass
   public void setup() {
     executor = Executors.newFixedThreadPool(2);
-    // This endpoint doesn't exist, so we won't call `close` on the cached decryptor to avoid the
+    secureRandom = new SecureRandom();
+    // This endpoint doesn't exist, so we won't call `close` on the cached key to avoid the
     // report-operations request
     encryptionService = new TenantSecurityRequest("http://localhost:0", "test-api-key", 1, 1000);
   }
@@ -57,8 +60,8 @@ private byte[] createValidDek() {
     return dek;
   }
 
-  private CachedKeyDecryptor createDecryptor() {
-    return new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor, encryptionService,
+  private CachedKey createCachedKey() {
+    return new CachedKey(createValidDek(), TEST_EDEK, executor, secureRandom, encryptionService,
         metadata);
   }
 
@@ -67,7 +70,7 @@ private CachedKeyDecryptor createDecryptor() {
   @SuppressWarnings("resource")
   public void constructorRejectNullDek() {
     try {
-      new CachedKeyDecryptor(null, TEST_EDEK, executor, encryptionService, metadata);
+      new CachedKey(null, TEST_EDEK, executor, secureRandom, encryptionService, metadata);
       fail("Should have thrown IllegalArgumentException");
     } catch (IllegalArgumentException e) {
       assertTrue(e.getMessage().contains("DEK must be exactly 32 bytes"));
@@ -78,7 +81,7 @@ public void constructorRejectNullDek() {
   public void constructorRejectWrongSizeDek() {
     byte[] shortDek = new byte[16];
     try {
-      new CachedKeyDecryptor(shortDek, TEST_EDEK, executor, encryptionService, metadata);
+      new CachedKey(shortDek, TEST_EDEK, executor, secureRandom, encryptionService, metadata);
       fail("Should have thrown IllegalArgumentException");
     } catch (IllegalArgumentException e) {
       assertTrue(e.getMessage().contains("DEK must be exactly 32 bytes"));
@@ -88,7 +91,7 @@ public void constructorRejectWrongSizeDek() {
   @SuppressWarnings("resource")
   public void constructorRejectNullEdek() {
     try {
-      new CachedKeyDecryptor(createValidDek(), null, executor, encryptionService, metadata);
+      new CachedKey(createValidDek(), null, executor, secureRandom, encryptionService, metadata);
       fail("Should have thrown IllegalArgumentException");
     } catch (IllegalArgumentException e) {
       assertTrue(e.getMessage().contains("EDEK must not be null or empty"));
@@ -98,7 +101,7 @@ public void constructorRejectNullEdek() {
   @SuppressWarnings("resource")
   public void constructorRejectEmptyEdek() {
     try {
-      new CachedKeyDecryptor(createValidDek(), "", executor, encryptionService, metadata);
+      new CachedKey(createValidDek(), "", executor, secureRandom, encryptionService, metadata);
       fail("Should have thrown IllegalArgumentException");
     } catch (IllegalArgumentException e) {
       assertTrue(e.getMessage().contains("EDEK must not be null or empty"));
@@ -108,17 +111,27 @@ public void constructorRejectEmptyEdek() {
   @SuppressWarnings("resource")
   public void constructorRejectNullExecutor() {
     try {
-      new CachedKeyDecryptor(createValidDek(), TEST_EDEK, null, encryptionService, metadata);
+      new CachedKey(createValidDek(), TEST_EDEK, null, secureRandom, encryptionService, metadata);
       fail("Should have thrown IllegalArgumentException");
     } catch (IllegalArgumentException e) {
       assertTrue(e.getMessage().contains("encryptionExecutor must not be null"));
     }
   }
 
+  @SuppressWarnings("resource")
+  public void constructorRejectNullSecureRandom() {
+    try {
+      new CachedKey(createValidDek(), TEST_EDEK, executor, null, encryptionService, metadata);
+      fail("Should have thrown IllegalArgumentException");
+    } catch (IllegalArgumentException e) {
+      assertTrue(e.getMessage().contains("secureRandom must not be null"));
+    }
+  }
+
   @SuppressWarnings("resource")
   public void constructorRejectNullEncryptionService() {
     try {
-      new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor, null, metadata);
+      new CachedKey(createValidDek(), TEST_EDEK, executor, secureRandom, null, metadata);
       fail("Should have thrown IllegalArgumentException");
     } catch (IllegalArgumentException e) {
       assertTrue(e.getMessage().contains("requestService must not be null"));
@@ -128,7 +141,7 @@ public void constructorRejectNullEncryptionService() {
   @SuppressWarnings("resource")
   public void constructorRejectNullMetadata() {
     try {
-      new CachedKeyDecryptor(createValidDek(), TEST_EDEK, executor, encryptionService, null);
+      new CachedKey(createValidDek(), TEST_EDEK, executor, secureRandom, encryptionService, null);
       fail("Should have thrown IllegalArgumentException");
     } catch (IllegalArgumentException e) {
       assertTrue(e.getMessage().contains("metadata must not be null"));
@@ -138,138 +151,187 @@ public void constructorRejectNullMetadata() {
   // Getter tests
 
   public void getEdekReturnsCorrectValue() {
-    CachedKeyDecryptor decryptor = createDecryptor();
-    assertEquals(decryptor.getEdek(), TEST_EDEK);
-    decryptor.close();
+    CachedKey cachedKey = createCachedKey();
+    assertEquals(cachedKey.getEdek(), TEST_EDEK);
+    cachedKey.close();
   }
 
   public void isClosedReturnsFalseInitially() {
-    CachedKeyDecryptor decryptor = createDecryptor();
-    assertFalse(decryptor.isClosed());
-    decryptor.close();
+    CachedKey cachedKey = createCachedKey();
+    assertFalse(cachedKey.isClosed());
+    cachedKey.close();
   }
 
   public void isClosedReturnsTrueAfterClose() {
-    CachedKeyDecryptor decryptor = createDecryptor();
-    decryptor.close();
-    assertTrue(decryptor.isClosed());
+    CachedKey cachedKey = createCachedKey();
+    cachedKey.close();
+    assertTrue(cachedKey.isClosed());
   }
 
   // Close tests
 
   public void closeIsIdempotent() {
-    CachedKeyDecryptor decryptor = createDecryptor();
-    decryptor.close();
-    assertTrue(decryptor.isClosed());
+    CachedKey cachedKey = createCachedKey();
+    cachedKey.close();
+    assertTrue(cachedKey.isClosed());
     // Should not throw
-    decryptor.close();
-    decryptor.close();
-    assertTrue(decryptor.isClosed());
+    cachedKey.close();
+    cachedKey.close();
+    assertTrue(cachedKey.isClosed());
   }
 
   // Operation count tests
 
   public void operationCountStartsAtZero() {
-    CachedKeyDecryptor decryptor = createDecryptor();
-    assertEquals(decryptor.getOperationCount(), 0);
-    decryptor.close();
+    CachedKey cachedKey = createCachedKey();
+    assertEquals(cachedKey.getOperationCount(), 0);
+    assertEquals(cachedKey.getEncryptCount(), 0);
+    assertEquals(cachedKey.getDecryptCount(), 0);
+    cachedKey.close();
+  }
+
+  // Encrypt validation tests
+
+  public void encryptFailsWhenClosed() {
+    CachedKey cachedKey = createCachedKey();
+    cachedKey.close();
+
+    try {
+      cachedKey.encrypt(java.util.Collections.emptyMap(), metadata).join();
+      fail("Should have thrown CompletionException");
+    } catch (CompletionException e) {
+      assertTrue(e.getCause() instanceof TscException);
+      assertTrue(e.getCause().getMessage().contains("CachedKey has been closed"));
+    }
+  }
+
+  public void encryptStreamFailsWhenClosed() {
+    CachedKey cachedKey = createCachedKey();
+    cachedKey.close();
+
+    ByteArrayInputStream input = new ByteArrayInputStream(new byte[0]);
+    ByteArrayOutputStream output = new ByteArrayOutputStream();
+
+    try {
+      cachedKey.encryptStream(input, output, metadata).join();
+      fail("Should have thrown CompletionException");
+    } catch (CompletionException e) {
+      assertTrue(e.getCause() instanceof TscException);
+      assertTrue(e.getCause().getMessage().contains("CachedKey has been closed"));
+    }
+  }
+
+  public void encryptBatchFailsWhenClosed() {
+    CachedKey cachedKey = createCachedKey();
+    cachedKey.close();
+
+    Map> docs = new HashMap<>();
+    docs.put("doc1", java.util.Collections.singletonMap("field", new byte[] {1, 2, 3}));
+
+    try {
+      cachedKey.encryptBatch(docs, metadata).join();
+      fail("Should have thrown CompletionException");
+    } catch (CompletionException e) {
+      assertTrue(e.getCause() instanceof TscException);
+      assertTrue(e.getCause().getMessage().contains("CachedKey has been closed"));
+    }
   }
 
   // Decrypt validation tests
 
   public void decryptFailsWhenClosed() {
-    CachedKeyDecryptor decryptor = createDecryptor();
-    decryptor.close();
+    CachedKey cachedKey = createCachedKey();
+    cachedKey.close();
 
     EncryptedDocument encDoc = new EncryptedDocument(java.util.Collections.emptyMap(), TEST_EDEK);
 
     try {
-      decryptor.decrypt(encDoc, metadata).join();
+      cachedKey.decrypt(encDoc, metadata).join();
       fail("Should have thrown CompletionException");
     } catch (CompletionException e) {
       assertTrue(e.getCause() instanceof TscException);
-      assertTrue(e.getCause().getMessage().contains("CachedKeyDecryptor has been closed"));
+      assertTrue(e.getCause().getMessage().contains("CachedKey has been closed"));
     }
   }
 
   public void decryptFailsWhenEdekMismatch() {
-    CachedKeyDecryptor decryptor = createDecryptor();
+    CachedKey cachedKey = createCachedKey();
 
     EncryptedDocument encDoc =
         new EncryptedDocument(java.util.Collections.emptyMap(), DIFFERENT_EDEK);
 
     try {
-      decryptor.decrypt(encDoc, metadata).join();
+      cachedKey.decrypt(encDoc, metadata).join();
       fail("Should have thrown CompletionException");
     } catch (CompletionException e) {
       assertTrue(e.getCause() instanceof TscException);
       assertTrue(e.getCause().getMessage().contains("EDEK does not match"));
     } finally {
-      decryptor.close();
+      cachedKey.close();
     }
   }
 
   // DecryptStream validation tests
 
   public void decryptStreamFailsWhenClosed() {
-    CachedKeyDecryptor decryptor = createDecryptor();
-    decryptor.close();
+    CachedKey cachedKey = createCachedKey();
+    cachedKey.close();
 
     ByteArrayInputStream input = new ByteArrayInputStream(new byte[0]);
     ByteArrayOutputStream output = new ByteArrayOutputStream();
 
     try {
-      decryptor.decryptStream(TEST_EDEK, input, output, metadata).join();
+      cachedKey.decryptStream(TEST_EDEK, input, output, metadata).join();
       fail("Should have thrown CompletionException");
     } catch (CompletionException e) {
       assertTrue(e.getCause() instanceof TscException);
-      assertTrue(e.getCause().getMessage().contains("CachedKeyDecryptor has been closed"));
+      assertTrue(e.getCause().getMessage().contains("CachedKey has been closed"));
     }
   }
 
   public void decryptStreamFailsWhenEdekMismatch() {
-    CachedKeyDecryptor decryptor = createDecryptor();
+    CachedKey cachedKey = createCachedKey();
 
     ByteArrayInputStream input = new ByteArrayInputStream(new byte[0]);
     ByteArrayOutputStream output = new ByteArrayOutputStream();
 
     try {
-      decryptor.decryptStream(DIFFERENT_EDEK, input, output, metadata).join();
+      cachedKey.decryptStream(DIFFERENT_EDEK, input, output, metadata).join();
       fail("Should have thrown CompletionException");
     } catch (CompletionException e) {
       assertTrue(e.getCause() instanceof TscException);
       assertTrue(e.getCause().getMessage().contains("EDEK does not match"));
     } finally {
-      decryptor.close();
+      cachedKey.close();
     }
   }
 
   // decryptBatch validation tests
 
   public void decryptBatchFailsWhenClosed() {
-    CachedKeyDecryptor decryptor = createDecryptor();
-    decryptor.close();
+    CachedKey cachedKey = createCachedKey();
+    cachedKey.close();
 
     Map docs = new HashMap<>();
     docs.put("doc1", new EncryptedDocument(java.util.Collections.emptyMap(), TEST_EDEK));
 
     try {
-      decryptor.decryptBatch(docs, metadata).join();
+      cachedKey.decryptBatch(docs, metadata).join();
       fail("Should have thrown CompletionException");
     } catch (CompletionException e) {
       assertTrue(e.getCause() instanceof TscException);
-      assertTrue(e.getCause().getMessage().contains("CachedKeyDecryptor has been closed"));
+      assertTrue(e.getCause().getMessage().contains("CachedKey has been closed"));
     }
   }
 
   public void decryptBatchEdekMismatchGoesToFailures() {
-    CachedKeyDecryptor decryptor = createDecryptor();
+    CachedKey cachedKey = createCachedKey();
 
     Map docs = new HashMap<>();
     docs.put("matching", new EncryptedDocument(java.util.Collections.emptyMap(), TEST_EDEK));
     docs.put("mismatched", new EncryptedDocument(java.util.Collections.emptyMap(), DIFFERENT_EDEK));
 
-    BatchResult result = decryptor.decryptBatch(docs, metadata).join();
+    BatchResult result = cachedKey.decryptBatch(docs, metadata).join();
 
     // The matching doc with empty fields should succeed (no fields to decrypt)
     assertTrue(result.getSuccesses().containsKey("matching"));
@@ -280,29 +342,29 @@ public void decryptBatchEdekMismatchGoesToFailures() {
     // The matching doc should NOT be in failures
     assertFalse(result.getFailures().containsKey("matching"));
 
-    decryptor.close();
+    cachedKey.close();
   }
 
   // DEK copying test
 
   public void constructorCopiesDekToPreventExternalModification() throws Exception {
     byte[] originalDek = createValidDek();
-    CachedKeyDecryptor decryptor =
-        new CachedKeyDecryptor(originalDek, TEST_EDEK, executor, encryptionService, metadata);
+    CachedKey cachedKey = new CachedKey(originalDek, TEST_EDEK, executor, secureRandom,
+        encryptionService, metadata);
 
     // Modify the original array
     Arrays.fill(originalDek, (byte) 0x00);
 
     // Use reflection to verify internal DEK still has original values
-    Field dekField = CachedKeyDecryptor.class.getDeclaredField("dek");
+    Field dekField = CachedKey.class.getDeclaredField("dek");
     dekField.setAccessible(true);
-    byte[] internalDek = (byte[]) dekField.get(decryptor);
+    byte[] internalDek = (byte[]) dekField.get(cachedKey);
 
     // Internal DEK should still be 0x42, not 0x00
     for (byte b : internalDek) {
       assertEquals(b, (byte) 0x42, "Internal DEK should not be affected by external modification");
     }
 
-    decryptor.close();
+    cachedKey.close();
   }
 }

From 57217fedbf089b3324b701b5b7f03dd0a99219f8 Mon Sep 17 00:00:00 2001
From: Craig Colegrove 
Date: Tue, 10 Mar 2026 15:44:29 -0700
Subject: [PATCH 8/8] Retry reportOperations once in some conditions

---
 .../kms/v1/TenantSecurityRequest.java         | 22 ++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityRequest.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityRequest.java
index 9511ced..a96b297 100644
--- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityRequest.java
+++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityRequest.java
@@ -317,7 +317,27 @@ CompletableFuture reportOperations(DocumentMetadata metadata, String edek,
     String error = String.format(
         "Unable to make request to Tenant Security Proxy report-operations endpoint. Endpoint requested: %s",
         this.reportOperationsEndpoint);
-    return this.makeRequestAndParseFailure(this.reportOperationsEndpoint, postData, error);
+    return this.makeRequestAndParseFailure(this.reportOperationsEndpoint, postData, error)
+        .exceptionallyCompose(t -> {
+          if (isRetryable(t)) {
+            return this.makeRequestAndParseFailure(this.reportOperationsEndpoint, postData, error);
+          }
+          return CompletableFuture.failedFuture(t);
+        });
+  }
+
+  /**
+   * Check if a failure is retryable. Server errors (5xx) and connection failures (status 0) are
+   * retryable. Client errors (4xx) are not — the request itself is wrong and will fail identically
+   * on retry.
+   */
+  private static boolean isRetryable(Throwable t) {
+    Throwable cause = t instanceof CompletionException ? t.getCause() : t;
+    if (cause instanceof TspServiceException) {
+      int status = ((TspServiceException) cause).getHttpResponseCode();
+      return status == 0 || status >= 500;
+    }
+    return false;
   }
 
   /**