diff --git a/examples/cached-key-example/README.md b/examples/cached-key-example/README.md new file mode 100644 index 0000000..bed0919 --- /dev/null +++ b/examples/cached-key-example/README.md @@ -0,0 +1,63 @@ +# Cached Key Example + +In order to run this example, you need to be running a _Tenant Security Proxy_ (TSP) on your machine. +Check the [README.md](../README.md) file in the parent directory to see how to start the TSP, if you haven't done so +yet. + +Once the TSP is running, you can experiment with this example Java program. It demonstrates using +`CachedEncryptor` and `CachedDecryptor` to encrypt and decrypt multiple records while minimizing +calls to the TSP. The example code shows two scenarios: + +- encrypting three of a customer's records using a single cached key (one TSP wrap call) +- decrypting all three records using a single cached key (one TSP unwrap call) + +## Why use a cached key? + +A normal `encrypt()` call wraps a new DEK through the TSP on every invocation. If you're encrypting +several records in quick succession (like inside a database transaction), each call +adds a network round trip. + +A `CachedEncryptor` wraps the DEK once, then all subsequent `encrypt()` calls are purely local +CPU work. This means you can safely encrypt inside a database transaction without adding network +latency or external failure modes to the transaction. The same applies to `CachedDecryptor` for +reads. + +## Running the example + +To run the example, you will need to have a Java JRE 17+ and Maven installed on your computer. + +```bash +export API_KEY='0WUaXesNgbTAuLwn' +mvn package +java -cp target/cached-key-example-0.1.0.jar com.ironcorelabs.cachedkey.CachedKeyExample +``` + +We've assigned an API key for you, but in production you will make your own and edit the TSP +configuration with it. This should produce output like: + +``` +Using tenant tenant-gcp-l +Encrypted 3 records with one TSP call +Decrypted: Jim Bridger / 000-12-2345 +Decrypted: John Colter / 000-45-6789 +Decrypted: Sacagawea / 000-98-7654 +Decrypted 3 records with one TSP call +``` + +If you look at the TSP logs, you should see only two KMS operations: one wrap and one unwrap. +Without cached keys, the same work would have required six KMS operations (three wraps + three +unwraps). + +If you would like to experiment with a different tenant, just do: + +```bash +export TENANT_ID= +java -cp target/cached-key-example-0.1.0.jar com.ironcorelabs.cachedkey.CachedKeyExample +``` + +The list of available tenants is listed in the [README.md](../README.md) in the parent directory. + +## Additional Resources + +If you would like some more in-depth information, our website features a section of technical +documentation about the [SaaS Shield product](https://ironcorelabs.com/docs/saas-shield/). diff --git a/examples/cached-key-example/pom.xml b/examples/cached-key-example/pom.xml new file mode 100644 index 0000000..ff44087 --- /dev/null +++ b/examples/cached-key-example/pom.xml @@ -0,0 +1,85 @@ + + + + 4.0.0 + + com.ironcorelabs + cached-key-example + jar + 0.1.0 + + cached-key-example + https://www.docs.ironcorelabs.com + + + Apache-2 + https://opensource.org/licenses/Apache-2.0 + repo + + + + + UTF-8 + 1.8 + 1.8 + + + + + com.ironcorelabs + tenant-security-java + 8.1.0 + + + + + + + org.apache.maven.plugins + maven-gpg-plugin + 1.6 + + + sign-artifacts + verify + + sign + + + + + + org.apache.maven.plugins + maven-compiler-plugin + 3.3 + + 1.8 + 1.8 + 1.8 + 1.8 + -Xlint:unchecked + + + + org.apache.maven.plugins + maven-source-plugin + 3.0.1 + + + org.apache.maven.plugins + maven-shade-plugin + 3.2.1 + + + package + + shade + + + + + + + \ No newline at end of file diff --git a/examples/cached-key-example/src/main/java/com/ironcorelabs/cachedkey/CachedKeyExample.java b/examples/cached-key-example/src/main/java/com/ironcorelabs/cachedkey/CachedKeyExample.java new file mode 100644 index 0000000..6e6db82 --- /dev/null +++ b/examples/cached-key-example/src/main/java/com/ironcorelabs/cachedkey/CachedKeyExample.java @@ -0,0 +1,93 @@ +package com.ironcorelabs.cachedkey; + +import com.ironcorelabs.tenantsecurity.kms.v1.*; +import com.ironcorelabs.tenantsecurity.kms.v1.exception.TenantSecurityException; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ExecutionException; + +/** + * Demonstrates using CachedEncryptor and CachedDecryptor to encrypt/decrypt multiple records with a + * single TSP call. + */ +public class CachedKeyExample { + + private static final String TSP_ADDR = "http://localhost:7777"; + + public static void main(String[] args) throws Exception { + + String API_KEY = System.getenv("API_KEY"); + if (API_KEY == null) { + System.out.println("Must set the API_KEY environment variable."); + System.exit(1); + } + + String tenantId = System.getenv("TENANT_ID"); + if (tenantId == null) { + tenantId = "tenant-gcp-l"; + } + System.out.println("Using tenant " + tenantId); + + TenantSecurityClient client = + new TenantSecurityClient.Builder(TSP_ADDR, API_KEY).allowInsecureHttp(true).build(); + + DocumentMetadata metadata = new DocumentMetadata(tenantId, "serviceOrUserId", "PII"); + + // Simulate a database table: each row has an encrypted record and its EDEK + List> encryptedRows = new ArrayList<>(); + String sharedEdek; + + // Encrypt: one TSP call, then N local encrypts + // + // In a real application this block would be inside a database transaction. The + // createCachedEncryptor call is the only network round trip — every encrypt() after + // that is purely local CPU work, so it won't add latency or failure modes to the + // transaction. + try (CachedEncryptor encryptor = client.createCachedEncryptor(metadata).get()) { + String[][] customers = + {{"000-12-2345", "2825-519 Stone Creek Rd, Bozeman, MT 59715", "Jim Bridger"}, + {"000-45-6789", "100 Main St, Helena, MT 59601", "John Colter"}, + {"000-98-7654", "742 Evergreen Terrace, Missoula, MT 59801", "Sacagawea"},}; + + for (String[] customer : customers) { + Map record = new HashMap<>(); + record.put("ssn", customer[0].getBytes(StandardCharsets.UTF_8)); + record.put("address", customer[1].getBytes(StandardCharsets.UTF_8)); + record.put("name", customer[2].getBytes(StandardCharsets.UTF_8)); + + // This encrypt is local — no TSP call + EncryptedDocument encrypted = encryptor.encrypt(record, metadata).get(); + encryptedRows.add(encrypted.getEncryptedFields()); + } + + // All rows share this EDEK; store it alongside the rows (or once per batch) + sharedEdek = encryptor.getEdek(); + + System.out + .println("Encrypted " + encryptor.getOperationCount() + " records with one TSP call"); + } + // leaving the `try` block zeroes the DEK and reports usage to the TSP + + // Decrypt: one TSP call, then N local decrypts + try (CachedDecryptor decryptor = client.createCachedDecryptor(sharedEdek, metadata).get()) { + for (Map row : encryptedRows) { + EncryptedDocument doc = new EncryptedDocument(row, sharedEdek); + + // This decrypt is local — no TSP call + PlaintextDocument plaintext = decryptor.decrypt(doc, metadata).get(); + Map fields = plaintext.getDecryptedFields(); + + System.out.println("Decrypted: " + new String(fields.get("name"), StandardCharsets.UTF_8) + + " / " + new String(fields.get("ssn"), StandardCharsets.UTF_8)); + } + + System.out + .println("Decrypted " + decryptor.getOperationCount() + " records with one TSP call"); + } + + System.exit(0); + } +} diff --git a/examples/large-documents/README.md b/examples/large-documents/README.md index f94bae7..853911c 100644 --- a/examples/large-documents/README.md +++ b/examples/large-documents/README.md @@ -11,8 +11,7 @@ fields that share a DEK. The example code shows two scenarios: - encrypting a large document as many subdocs, using the disk for persistence - retrieving and decrypting subdocs individually -To run the example, you will need to have Java and Maven installed on your computer. Try a `java -version` to see -what version you are using. We tested the example code using 1.8. +To run the example, you will need to have Java JRE 17+ and Maven installed on your computer. If java is ready to go, execute these commands: diff --git a/examples/logging-example/README.md b/examples/logging-example/README.md index 8ae4248..0bb6e61 100644 --- a/examples/logging-example/README.md +++ b/examples/logging-example/README.md @@ -10,8 +10,7 @@ to use the Tenant Security Client (TSC) SDK to log security events. The example - logging a user create security event with minimal metadata - logging a login security event with additional metadata -To run the example, you will need to have Java and Maven installed on your computer. Try a `java -version` to see -what version you are using. We tested the example code using 1.8. +To run the example, you will need to have Java JRE 17+ and Maven installed on your computer. If java is ready to go, execute these commands: diff --git a/examples/rekey-example/README.md b/examples/rekey-example/README.md index d9b6d01..5369a9c 100644 --- a/examples/rekey-example/README.md +++ b/examples/rekey-example/README.md @@ -15,7 +15,7 @@ to use the Tenant Security Client (TSC) SDK to rekey data. The example code cont - Rekeying the encrypted record to a new tenant - Decrypting the encrypted record with the new tenant -To run the example, you will need to have a Java JRE 8+ and Maven installed on your computer. +To run the example, you will need to have a Java JRE 17+ and Maven installed on your computer. ```bash export API_KEY='0WUaXesNgbTAuLwn' diff --git a/examples/simple-roundtrip/README.md b/examples/simple-roundtrip/README.md index 27c9111..29a439f 100644 --- a/examples/simple-roundtrip/README.md +++ b/examples/simple-roundtrip/README.md @@ -10,7 +10,7 @@ to use the Tenant Security Client (TSC) SDK to encrypt and decrypt data. The exa - encryption and decryption of a record that you might store in a key-value store or a database - encryption and decryption of a file, using the file-system for storage -To run the example, you will need to have a Java JRE 8+ and Maven installed on your computer. +To run the example, you will need to have a Java JRE 17+ and Maven installed on your computer. ```bash export API_KEY='0WUaXesNgbTAuLwn' 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/pom.xml b/pom.xml index 2b997f8..9cd1519 100644 --- a/pom.xml +++ b/pom.xml @@ -8,7 +8,7 @@ com.ironcorelabs tenant-security-java jar - 8.0.1 + 8.1.0 tenant-security-java https://ironcorelabs.com/docs Java client library for the IronCore Labs Tenant Security Proxy. @@ -257,4 +257,4 @@ - + \ No newline at end of file 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..9b38cc6 --- /dev/null +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedEncryptor.java @@ -0,0 +1,17 @@ +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/CachedKey.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKey.java new file mode 100644 index 0000000..d4fb053 --- /dev/null +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKey.java @@ -0,0 +1,343 @@ +package com.ironcorelabs.tenantsecurity.kms.v1; + +import java.io.InputStream; +import java.io.OutputStream; +import java.security.SecureRandom; +import java.time.Duration; +import java.time.Instant; +import java.lang.invoke.MethodHandles; +import java.lang.invoke.VarHandle; +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.Function; +import java.util.function.ToIntFunction; +import com.ironcorelabs.tenantsecurity.kms.v1.exception.TenantSecurityException; +import com.ironcorelabs.tenantsecurity.kms.v1.exception.TscException; + +/** + * 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 encrypt and decrypt + * operations. Once closed, all operations will fail. + * + *

+ * 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 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#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) + */ +public final class CachedKey implements CachedEncryptor, CachedDecryptor { + + // 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 associated with the cached DEK + private final String edek; + + // 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; + + // Flag to track if close() has been called + private final AtomicBoolean closed = new AtomicBoolean(false); + + // Protects the DEK from being read (copied) while close() is zeroing it + private final Object dekLock = new Object(); + + // 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 decryptCount = new AtomicInteger(0); + + /** + * Package-private constructor. Use TenantSecurityClient.createCachedEncryptor() or + * TenantSecurityClient.createCachedDecryptor() 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/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 + */ + 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"); + } + 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 key. + * + * @return The EDEK string + */ + public String getEdek() { + return edek; + } + + /** + * Check if this cached key has been closed. + * + * @return true if close() has been called + */ + public boolean isClosed() { + return closed.get(); + } + + /** + * Check if this cached key 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 cached key. + * + * @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 encryptCount.get() + decryptCount.get(); + } + + /** + * Guard an operation with usability checks, DEK copy safety, and operation counting. Atomically + * checks that the cached key is not closed and copies the DEK under a lock so that close() cannot + * zero the DEK mid-copy. The copy is zeroed after the operation completes (success or failure). + * + * @param operation The operation to perform, receiving a safe copy of the DEK + * @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( + Function> operation, ToIntFunction countOps, + AtomicInteger counter, TenantSecurityErrorCodes errorCode) { + byte[] dekCopy; + synchronized (dekLock) { + if (closed.get()) { + return CompletableFuture + .failedFuture(new TscException(errorCode, "CachedKey has been closed")); + } + dekCopy = Arrays.copyOf(dek, dek.length); + } + if (isExpired()) { + zeroDek(dekCopy); + return CompletableFuture.failedFuture(new TscException(errorCode, "CachedKey has expired")); + } + return operation.apply(dekCopy).whenComplete((result, ex) -> zeroDek(dekCopy)) + .thenApply(result -> { + counter.addAndGet(countOps.applyAsInt(result)); + return result; + }); + } + + @Override + public CompletableFuture encrypt(Map document, + DocumentMetadata metadata) { + return executeAndIncrement( + dekCopy -> DocumentCryptoOps.encryptFields(document, metadata, dekCopy, edek, + encryptionExecutor, secureRandom), + result -> 1, encryptCount, TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED); + } + + @Override + public CompletableFuture encryptStream(InputStream input, OutputStream output, + DocumentMetadata metadata) { + return executeAndIncrement( + dekCopy -> CompletableFuture + .supplyAsync( + () -> CryptoUtils + .encryptStreamInternal(dekCopy, 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(dekCopy -> { + ConcurrentMap> ops = new ConcurrentHashMap<>(); + plaintextDocuments.forEach((id, doc) -> ops.put(id, DocumentCryptoOps.encryptFields(doc, + metadata, dekCopy, edek, encryptionExecutor, secureRandom))); + return CompletableFuture.supplyAsync(() -> DocumentCryptoOps.cryptoOperationToBatchResult(ops, + TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED)); + }, result -> result.getSuccesses().size(), encryptCount, + TenantSecurityErrorCodes.DOCUMENT_ENCRYPT_FAILED); + } + + private CompletableFuture validateEdekAndDecrypt( + EncryptedDocument encryptedDocument, byte[] dekCopy) { + if (!edek.equals(encryptedDocument.getEdek())) { + return CompletableFuture.failedFuture(new TscException( + TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, EDEK_MISMATCH_MESSAGE)); + } + return DocumentCryptoOps.decryptFields(encryptedDocument.getEncryptedFields(), dekCopy, + encryptedDocument.getEdek(), encryptionExecutor); + } + + @Override + public CompletableFuture decrypt(EncryptedDocument encryptedDocument, + DocumentMetadata metadata) { + return executeAndIncrement(dekCopy -> validateEdekAndDecrypt(encryptedDocument, dekCopy), + result -> 1, decryptCount, TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED); + } + + @Override + public CompletableFuture decryptStream(String edek, InputStream input, OutputStream output, + DocumentMetadata metadata) { + return executeAndIncrement(dekCopy -> { + if (!this.edek.equals(edek)) { + return CompletableFuture.failedFuture(new TscException( + TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED, EDEK_MISMATCH_MESSAGE)); + } + return CompletableFuture.supplyAsync( + () -> CryptoUtils.decryptStreamInternal(dekCopy, input, output).join(), + encryptionExecutor); + }, result -> 1, decryptCount, TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED); + } + + @Override + public CompletableFuture> decryptBatch( + Map encryptedDocuments, DocumentMetadata metadata) { + return executeAndIncrement(dekCopy -> { + 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, + EDEK_MISMATCH_MESSAGE)); + } else { + ops.put(id, DocumentCryptoOps.decryptFields(encDoc.getEncryptedFields(), dekCopy, + 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(), decryptCount, + TenantSecurityErrorCodes.DOCUMENT_DECRYPT_FAILED); + } + + /** + * 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. + */ + @Override + public void close() { + if (closed.compareAndSet(false, true)) { + synchronized (dekLock) { + zeroDek(dek); + } + int encrypts = encryptCount.get(); + int decrypts = decryptCount.get(); + if (encrypts > 0 || decrypts > 0) { + requestService.reportOperations(metadata, edek, encrypts, decrypts); + } + } + } + + /** + * Zero a DEK byte array using opaque stores to prevent the JIT from eliminating the writes as + * dead stores. VarHandle.setOpaque is the lightest memory ordering mode that guarantees the + * writes actually happen per the Java Memory Model spec. + */ + private static final VarHandle BYTE_ARRAY = MethodHandles.arrayElementVarHandle(byte[].class); + + static void zeroDek(byte[] dek) { + for (int i = 0; i < dek.length; i++) { + BYTE_ARRAY.setOpaque(dek, i, (byte) 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/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 new file mode 100644 index 0000000..5e60814 --- /dev/null +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentDecryptor.java @@ -0,0 +1,48 @@ +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 decryption capabilities. Implemented by both TenantSecurityClient (for + * standard decrypt operations that unwrap the EDEK each time) and CachedKey (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 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); + + /** + * 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 new file mode 100644 index 0000000..fc396dc --- /dev/null +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/DocumentEncryptor.java @@ -0,0 +1,45 @@ +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 CachedKey (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); + + /** + * 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 8bf9d28..0db468d 100644 --- a/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java +++ b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityClient.java @@ -14,10 +14,10 @@ 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; -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; @@ -27,7 +27,7 @@ * * @author IronCore Labs */ -public final class TenantSecurityClient implements Closeable { +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 @@ -231,61 +231,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. @@ -298,34 +243,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 @@ -339,10 +264,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); } @@ -359,10 +285,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); } @@ -408,6 +334,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( @@ -434,6 +361,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( @@ -453,11 +381,13 @@ 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) - .thenComposeAsync(newDocumentKeys -> encryptFields(document, metadata, - newDocumentKeys.getDekBytes(), newDocumentKeys.getEdek())); + .thenComposeAsync(newDocumentKeys -> DocumentCryptoOps.encryptFields(document, metadata, + newDocumentKeys.getDekBytes(), newDocumentKeys.getEdek(), encryptionExecutor, + secureRandom)); } /** @@ -477,9 +407,10 @@ 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()), - encryptionExecutor); + return this.encryptionService.unwrapKey(document.getEdek(), metadata) + .thenComposeAsync(dek -> DocumentCryptoOps.encryptFields(document.getDecryptedFields(), + metadata, dek, document.getEdek(), encryptionExecutor, secureRandom), + encryptionExecutor); } /** @@ -494,6 +425,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) @@ -546,11 +478,187 @@ 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( - decryptedDocumentAESKey -> decryptFields(encryptedDocument.getEncryptedFields(), - decryptedDocumentAESKey, encryptedDocument.getEdek())); + return this.encryptionService.unwrapKey(encryptedDocument.getEdek(), metadata) + .thenComposeAsync(decryptedDocumentAESKey -> DocumentCryptoOps.decryptFields( + encryptedDocument.getEncryptedFields(), decryptedDocumentAESKey, + encryptedDocument.getEdek(), encryptionExecutor)); + } + + 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); + CachedKey.zeroDek(dekBytes); + 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); + CachedKey.zeroDek(dekBytes); + 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())); + } + + /** + * 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 CachedDecryptor implements Closeable and should be used with try-with-resources to + * ensure the DEK is securely zeroed when done. + * + * @param edek The encrypted document encryption key to unwrap + * @param metadata Metadata for the unwrap operation + * @return CompletableFuture resolving to a CachedDecryptor + */ + public CompletableFuture createCachedDecryptor(String edek, + DocumentMetadata metadata) { + return newCachedKeyFromUnwrap(edek, metadata) + // narrow the returned type to be a CachedDecryptor instead of a full CachedKey + .thenApply(k -> k); + } + + /** + * 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 CachedDecryptor + */ + public CompletableFuture createCachedDecryptor( + EncryptedDocument encryptedDocument, DocumentMetadata metadata) { + return createCachedDecryptor(encryptedDocument.getEdek(), metadata); + } + + /** + * 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. + * + * @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 CachedDecryptor and returns a CompletableFuture + * @return CompletableFuture resolving to the operation's result + */ + public CompletableFuture withCachedDecryptor(String edek, DocumentMetadata metadata, + Function> operation) { + return withCachedResource(createCachedDecryptor(edek, metadata), operation); + } + + /** + * 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. + * + *

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

+ * The returned CachedEncryptor implements Closeable and should be used with try-with-resources to + * ensure the DEK is securely zeroed when done. + * + * @param metadata Metadata for the wrap operation + * @return CompletableFuture resolving to a CachedEncryptor + */ + public CompletableFuture createCachedEncryptor(DocumentMetadata metadata) { + return newCachedKeyFromWrap(metadata) + // narrow the returned type to be a CachedEncryptor instead of a full CachedKey + .thenApply(k -> k); + } + + /** + * 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. + * + * @param The type returned by the operation + * @param metadata Metadata for the wrap operation + * @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 withCachedResource(createCachedEncryptor(metadata), operation); + } + + /** + * 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); + } + + /** + * 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); } /** @@ -581,6 +689,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/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityRequest.java b/src/main/java/com/ironcorelabs/tenantsecurity/kms/v1/TenantSecurityRequest.java index ec88381..a96b297 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,51 @@ 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) + .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; + } + /** * Request derive key endpoint. */ 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..4cd464c --- /dev/null +++ b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyOpsRoundTrip.java @@ -0,0 +1,508 @@ +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; + } + + public void cachedEncryptorRoundTrip() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc1 = getDocumentFields(); + Map doc2 = new HashMap<>(); + doc2.put("other", "Other document data".getBytes("UTF-8")); + + try (TenantSecurityClient client = getClient().get();) { + try (CachedEncryptor encryptor = client.createCachedEncryptor(metadata).get()) { + assertFalse(encryptor.isClosed()); + assertFalse(encryptor.isExpired()); + assertEquals(encryptor.getOperationCount(), 0); + + // Encrypt two documents with the cached key + 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")); + } + } + } + + public void cachedEncryptorWithPattern() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc = getDocumentFields(); + + try (TenantSecurityClient client = getClient().get()) { + // Use the withCachedEncryptor pattern for automatic lifecycle management + EncryptedDocument encrypted = + client.withCachedEncryptor(metadata, cachedKey -> cachedKey.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")); + } + } + + public void cachedDecryptorRoundTrip() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc1 = getDocumentFields(); + Map doc2 = new HashMap<>(); + doc2.put("other", "Other document data".getBytes("UTF-8")); + + try (TenantSecurityClient client = getClient().get()) { + // Encrypt two documents with the same key (using cached key) + EncryptedDocument enc1; + EncryptedDocument enc2; + 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 (CachedDecryptor 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")); + } + } + } + + public void cachedDecryptorFromEncryptedDocument() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc = getDocumentFields(); + + try (TenantSecurityClient client = getClient().get()) { + EncryptedDocument encrypted = client.encrypt(doc, metadata).get(); + + // Create decryptor from EncryptedDocument directly + 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")); + assertEqualBytes(decrypted.getDecryptedFields().get("field3"), doc.get("field3")); + } + } + } + + public void cachedDecryptorWithPattern() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc = getDocumentFields(); + + try (TenantSecurityClient client = getClient().get()) { + EncryptedDocument encrypted = client.encrypt(doc, metadata).get(); + + // Use the withCachedDecryptor pattern for automatic lifecycle management + PlaintextDocument decrypted = client.withCachedDecryptor(encrypted.getEdek(), metadata, + cachedKey -> cachedKey.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")); + } + } + + public void cachedEncryptorStreamRoundTrip() throws Exception { + DocumentMetadata metadata = getMetadata(); + byte[] plaintext = "Stream encrypt with cached key test data".getBytes("UTF-8"); + + try (TenantSecurityClient client = getClient().get()) { + ByteArrayOutputStream encryptedOutput = new ByteArrayOutputStream(); + String edek; + + 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 cached key'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); + } + } + + public void cachedDecryptorStreamRoundTrip() throws Exception { + DocumentMetadata metadata = getMetadata(); + byte[] plaintext = "Stream decrypt with cached key test data".getBytes("UTF-8"); + + try (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 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(); + + assertEquals(decryptor.getOperationCount(), 1); + assertEqualBytes(decryptedOutput.toByteArray(), plaintext); + } + } + } + + public void cachedEncryptToCachedDecryptRoundTrip() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc1 = getDocumentFields(); + Map doc2 = new HashMap<>(); + doc2.put("solo", "Solo field document".getBytes("UTF-8")); + + try (TenantSecurityClient client = getClient().get()) { + // Encrypt multiple docs with cached key + EncryptedDocument enc1; + EncryptedDocument enc2; + try (CachedEncryptor encryptor = client.createCachedEncryptor(metadata).get()) { + enc1 = encryptor.encrypt(doc1, metadata).get(); + enc2 = encryptor.encrypt(doc2, 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(); + + 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")); + } + } + } + + public void cachedStreamEncryptToCachedStreamDecryptRoundTrip() throws Exception { + DocumentMetadata metadata = getMetadata(); + byte[] plaintext = "Full cached stream round-trip data".getBytes("UTF-8"); + + try (TenantSecurityClient client = getClient().get()) { + // Encrypt stream with cached key + ByteArrayOutputStream encryptedOutput = new ByteArrayOutputStream(); + String edek; + 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 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(); + assertEqualBytes(decryptedOutput.toByteArray(), plaintext); + } + } + } + + public void singleCachedKeyEncryptAndDecrypt() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc1 = getDocumentFields(); + Map doc2 = new HashMap<>(); + doc2.put("other", "Other document data".getBytes("UTF-8")); + + try (TenantSecurityClient client = getClient().get(); + 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")); + } + } + + public void cachedEncryptorRejectsAfterClose() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc = getDocumentFields(); + + try (TenantSecurityClient client = getClient().get()) { + CachedEncryptor 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")); + } + } + } + + public void cachedDecryptorRejectsAfterClose() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc = getDocumentFields(); + + try (TenantSecurityClient client = getClient().get()) { + EncryptedDocument encrypted = client.encrypt(doc, 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); + + // 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")); + } + } + } + + 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); + BatchResult encResult; + + try (TenantSecurityClient client = getClient().get(); + CachedEncryptor 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")); + } + } + + 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")); + EncryptedDocument enc1, enc2, enc3; + + try (TenantSecurityClient client = getClient().get(); + CachedEncryptor encryptor = client.createCachedEncryptor(metadata).get()) { + // Encrypt all 3 with cached key (same key) + enc1 = encryptor.encrypt(doc1, metadata).get(); + enc2 = encryptor.encrypt(doc2, metadata).get(); + enc3 = encryptor.encrypt(doc3, metadata).get(); + + // Batch decrypt with cached key + Map encDocs = new HashMap<>(); + encDocs.put("doc1", enc1); + encDocs.put("doc2", enc2); + encDocs.put("doc3", enc3); + + try (CachedDecryptor 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")); + } + } + } + + public void cachedDecryptorBatchEdekMismatchPartialFailure() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc = getDocumentFields(); + + try (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 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 + 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); + } + } + } + + 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); + + try (TenantSecurityClient client = getClient().get(); + CachedEncryptor 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); + } + } + + public void cachedDecryptorRejectsEdekMismatch() throws Exception { + DocumentMetadata metadata = getMetadata(); + Map doc = getDocumentFields(); + + try (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 cached key for enc1's EDEK + try (CachedDecryptor 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")); + } + } + } + } +} diff --git a/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyTest.java b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyTest.java new file mode 100644 index 0000000..a819640 --- /dev/null +++ b/src/test/java/com/ironcorelabs/tenantsecurity/kms/v1/CachedKeyTest.java @@ -0,0 +1,337 @@ +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 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"; + 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 key 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 CachedKey createCachedKey() { + return new CachedKey(createValidDek(), TEST_EDEK, executor, secureRandom, encryptionService, + metadata); + } + + public void constructorRejectNullDek() { + try { + 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")); + } + } + + public void constructorRejectWrongSizeDek() { + byte[] shortDek = new byte[16]; + try { + 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")); + } + } + + public void constructorRejectNullEdek() { + try { + 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")); + } + } + + public void constructorRejectEmptyEdek() { + try { + 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")); + } + } + + public void constructorRejectNullExecutor() { + try { + 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")); + } + } + + 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")); + } + } + + public void constructorRejectNullEncryptionService() { + try { + 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")); + } + } + + public void constructorRejectNullMetadata() { + try { + 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")); + } + } + + public void getEdekReturnsCorrectValue() { + try (CachedKey cachedKey = createCachedKey()) { + assertEquals(cachedKey.getEdek(), TEST_EDEK); + } + } + + public void isClosedReturnsFalseInitially() { + try (CachedKey cachedKey = createCachedKey()) { + assertFalse(cachedKey.isClosed()); + } + } + + public void isClosedReturnsTrueAfterClose() { + CachedKey cachedKey = createCachedKey(); + cachedKey.close(); + assertTrue(cachedKey.isClosed()); + } + + public void closeIsIdempotent() { + CachedKey cachedKey = createCachedKey(); + cachedKey.close(); + assertTrue(cachedKey.isClosed()); + // Should not throw + cachedKey.close(); + cachedKey.close(); + assertTrue(cachedKey.isClosed()); + } + + public void operationCountStartsAtZero() { + try (CachedKey cachedKey = createCachedKey()) { + assertEquals(cachedKey.getOperationCount(), 0); + assertEquals(cachedKey.getEncryptCount(), 0); + assertEquals(cachedKey.getDecryptCount(), 0); + } + } + + 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")); + } + } + + public void decryptFailsWhenClosed() { + CachedKey cachedKey = createCachedKey(); + cachedKey.close(); + + EncryptedDocument encDoc = new EncryptedDocument(java.util.Collections.emptyMap(), TEST_EDEK); + + try { + cachedKey.decrypt(encDoc, 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 decryptFailsWhenEdekMismatch() { + try (CachedKey cachedKey = createCachedKey()) { + EncryptedDocument encDoc = + new EncryptedDocument(java.util.Collections.emptyMap(), DIFFERENT_EDEK); + + try { + 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")); + } + } + } + + public void decryptStreamFailsWhenClosed() { + CachedKey cachedKey = createCachedKey(); + cachedKey.close(); + + ByteArrayInputStream input = new ByteArrayInputStream(new byte[0]); + ByteArrayOutputStream output = new ByteArrayOutputStream(); + + try { + 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("CachedKey has been closed")); + } + } + + public void decryptStreamFailsWhenEdekMismatch() { + try (CachedKey cachedKey = createCachedKey()) { + ByteArrayInputStream input = new ByteArrayInputStream(new byte[0]); + ByteArrayOutputStream output = new ByteArrayOutputStream(); + + try { + 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")); + } + } + } + + public void decryptBatchFailsWhenClosed() { + CachedKey cachedKey = createCachedKey(); + cachedKey.close(); + + Map docs = new HashMap<>(); + docs.put("doc1", new EncryptedDocument(java.util.Collections.emptyMap(), TEST_EDEK)); + + try { + cachedKey.decryptBatch(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")); + } + } + + public void decryptBatchEdekMismatchGoesToFailures() { + try (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 = cachedKey.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")); + } + } + + public void constructorCopiesDekToPreventExternalModification() throws Exception { + byte[] originalDek = createValidDek(); + try (CachedKey cachedKey = new CachedKey(originalDek, TEST_EDEK, executor, secureRandom, + encryptionService, metadata)) { + // Modify the original array + Arrays.fill(originalDek, (byte) 0x00); + + Field dekField = CachedKey.class.getDeclaredField("dek"); + dekField.setAccessible(true); + 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"); + } + } + } +} 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()); + } +}