Add cached key encryptor/decryptor

src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyDecryptor.java

@@ -0,0 +1,231 @@
+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.
+ *
+ * <p>
+ * This class is thread-safe and can be used concurrently for multiple decrypt operations. Once
+ * closed, all decrypt operations will fail.
+ *
+ * <p>
+ * <b>Expiration:</b> 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.
+ *
+ * <p>
+ * Instances are created via {@link TenantSecurityClient#createCachedDecryptor} or
+ * {@link TenantSecurityClient#withCachedDecryptor}. See those methods for usage examples.
+ *
+ * @see TenantSecurityClient#createCachedDecryptor(String, DocumentMetadata)
+ * @see TenantSecurityClient#withCachedDecryptor(String, DocumentMetadata,
+ *      java.util.function.Function)
+ */
+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.
+   *
+   * <p>
+   * 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<PlaintextDocument> 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.
+   *
+   * <p>
+   * 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<Void> 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<PlaintextDocument> decryptFields(Map<String, byte[]> 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<String, CompletableFuture<byte[]>> 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<String, byte[]> 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.
+   *
+   * <p>
+   * 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);
+    }
+  }
+}

src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentDecryptor.java

@@ -0,0 +1,37 @@
+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).
+ */
+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<PlaintextDocument> 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<Void> decryptStream(String edek, InputStream input, OutputStream output,
+      DocumentMetadata metadata);
+}

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<StreamingResponse> 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<Void> decryptStream(String edek, InputStream input, OutputStream output,
       DocumentMetadata metadata) {
     return this.encryptionService.unwrapKey(edek, metadata).thenApplyAsync(
@@ -546,13 +548,97 @@ public CompletableFuture<BatchResult<EncryptedDocument>> encryptExistingBatch(
    * @param metadata Metadata about the document being decrypted.
    * @return PlaintextDocument which contains each documents decrypted field bytes.
    */
+  @Override
   public CompletableFuture<PlaintextDocument> decrypt(EncryptedDocument encryptedDocument,
       DocumentMetadata metadata) {
     return this.encryptionService.unwrapKey(encryptedDocument.getEdek(), metadata).thenComposeAsync(
         decryptedDocumentAESKey -> decryptFields(encryptedDocument.getEncryptedFields(),
             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.
+   *
+   * <p>
+   * Use this when you need to decrypt multiple documents that share the same EDEK, to avoid
+   * repeated TSP unwrap calls.
+   *
+   * <p>
+   * The returned decryptor implements AutoCloseable and should be used with try-with-resources to
+   * ensure the DEK is securely zeroed when done:
+   *
+   * <pre>
+   * try (CachedKeyDecryptor decryptor = client.createCachedDecryptor(edek, metadata).get()) {
+   *   PlaintextDocument doc1 = decryptor.decrypt(encDoc1, metadata).get();
+   *   PlaintextDocument doc2 = decryptor.decrypt(encDoc2, metadata).get();
+   * }
+   * </pre>
+   *
+   * @param edek The encrypted document encryption key to unwrap
+   * @param metadata Metadata for the unwrap operation
+   * @return CompletableFuture resolving to a CachedKeyDecryptor
+   */
+  public CompletableFuture<CachedKeyDecryptor> 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<CachedKeyDecryptor> 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.
+   *
+   * <p>
+   * This is the recommended pattern for using cached decryptors with CompletableFuture composition:
+   *
+   * <pre>
+   * client.withCachedDecryptor(edek, metadata, decryptor -&gt;
+   *     decryptor.decrypt(encDoc1, metadata)
+   *         .thenCompose(doc1 -&gt; decryptor.decrypt(encDoc2, metadata)))
+   * </pre>
+   *
+   * @param <T> 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 <T> CompletableFuture<T> withCachedDecryptor(String edek, DocumentMetadata metadata,
+      Function<CachedKeyDecryptor, CompletableFuture<T>> 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 <T> 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 <T> CompletableFuture<T> withCachedDecryptor(EncryptedDocument encryptedDocument,
+      DocumentMetadata metadata, Function<CachedKeyDecryptor, CompletableFuture<T>> 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