diff --git a/spring-vault-core/src/main/java/org/springframework/vault/core/VaultTransformTemplate.java b/spring-vault-core/src/main/java/org/springframework/vault/core/VaultTransformTemplate.java index 6d0ea0be..08271602 100644 --- a/spring-vault-core/src/main/java/org/springframework/vault/core/VaultTransformTemplate.java +++ b/spring-vault-core/src/main/java/org/springframework/vault/core/VaultTransformTemplate.java @@ -181,6 +181,10 @@ public class VaultTransformTemplate implements VaultTransformOperations { if (!ObjectUtils.isEmpty(context.getTweak())) { request.put("tweak", Base64.getEncoder().encodeToString(context.getTweak())); } + // NEW: pass "reference" in each item, if present + if (StringUtils.hasText(context.getReference())) { + request.put("reference", context.getReference()); + } } private static List toEncodedResults(VaultResponse vaultResponse, @@ -221,17 +225,17 @@ public class VaultTransformTemplate implements VaultTransformOperations { for (int i = 0; i < batchRequest.size(); i++) { - VaultTransformDecodeResult encrypted; + VaultTransformDecodeResult decodeResult; // Renamed from "encrypted" TransformCiphertext ciphertext = batchRequest.get(i); if (batchData.size() > i) { - encrypted = getDecryptionResult(batchData.get(i), ciphertext); + decodeResult = getDecryptionResult(batchData.get(i), ciphertext); } else { - encrypted = new VaultTransformDecodeResult(new VaultException("No result for ciphertext #" + i)); + decodeResult = new VaultTransformDecodeResult(new VaultException("No result for ciphertext #" + i)); } - result.add(encrypted); + result.add(decodeResult); } return result; @@ -246,8 +250,29 @@ public class VaultTransformTemplate implements VaultTransformOperations { if (StringUtils.hasText(data.get("decoded_value"))) { - return new VaultTransformDecodeResult( - TransformPlaintext.of(data.get("decoded_value")).with(ciphertext.getContext())); + // 1. Read reference from Vault's response (if present). + String returnedRef = data.get("reference"); + + // 2. Build an updated context that merges the existing transformation/tweak + // with the newly-returned reference. If no reference is returned, keep the + // old one. Note:- Relying on reference from originalContext is aimed at + // providing a + // fallback strategy, if vault does not return the reference, in any + // circumstance. + VaultTransformContext originalContext = ciphertext.getContext(); + VaultTransformContext updatedContext = VaultTransformContext.builder() + .transformation(originalContext.getTransformation()) + .tweak(originalContext.getTweak()) + .reference(returnedRef != null ? returnedRef : originalContext.getReference()) + .build(); + + // 3. Attach that updated context to the newly decoded plaintext. + TransformPlaintext decodedPlaintext = TransformPlaintext.of(data.get("decoded_value")).with(updatedContext); + + return new VaultTransformDecodeResult(decodedPlaintext); + + // return new VaultTransformDecodeResult( + // TransformPlaintext.of(data.get("decoded_value")).with(ciphertext.getContext())); } return new VaultTransformDecodeResult(TransformPlaintext.empty().with(ciphertext.getContext())); @@ -257,12 +282,16 @@ public class VaultTransformTemplate implements VaultTransformOperations { String ciphertext = (String) data.get("encoded_value"); + // if Vault returns "reference" in batch_results,capturing it for co-relation. + String returnedRef = (String) data.get("reference"); + VaultTransformContext contextToUse = context; if (data.containsKey("tweak")) { byte[] tweak = Base64.getDecoder().decode((String) data.get("tweak")); contextToUse = VaultTransformContext.builder() .transformation(context.getTransformation()) .tweak(tweak) + .reference(returnedRef != null ? returnedRef : context.getReference()) .build(); } diff --git a/spring-vault-core/src/main/java/org/springframework/vault/support/VaultTransformContext.java b/spring-vault-core/src/main/java/org/springframework/vault/support/VaultTransformContext.java index 8b7427c1..ffc662f2 100644 --- a/spring-vault-core/src/main/java/org/springframework/vault/support/VaultTransformContext.java +++ b/spring-vault-core/src/main/java/org/springframework/vault/support/VaultTransformContext.java @@ -32,15 +32,18 @@ public class VaultTransformContext { * Empty (default) {@link VaultTransformContext} without a {@literal context} and * {@literal nonce}. */ - private static final VaultTransformContext EMPTY = new VaultTransformContext("", new byte[0]); + private static final VaultTransformContext EMPTY = new VaultTransformContext("", new byte[0], ""); private final String transformation; private final byte[] tweak; - private VaultTransformContext(String transformation, byte[] tweak) { + private final String reference; + + private VaultTransformContext(String transformation, byte[] tweak, String reference) { this.transformation = transformation; this.tweak = tweak; + this.reference = reference; } /** @@ -98,6 +101,13 @@ public class VaultTransformContext { return this.tweak; } + /** + * @return The reference identifier for batch operations + */ + public String getReference() { + return this.reference; + } + @Override public boolean equals(Object o) { if (this == o) @@ -123,6 +133,20 @@ public class VaultTransformContext { private byte[] tweak = new byte[0]; + /** + * A user-defined identifier that can be used to correlate items in a batch + * request with their corresponding results in Vault's {@code batch_results}. + *
+ *
+ * + *

+ * If set, Vault echoes this value in the response so clients can match inputs to + * outputs reliably. If Vault does not return the {@code reference}, the original + * client-supplied reference remains available for correlation. + *

+ */ + private String reference = ""; + private VaultTransformRequestBuilder() { } @@ -157,12 +181,24 @@ public class VaultTransformContext { return this; } + /** + * Set a user-defined reference identifier. This reference is placed into each + * item of a batch request and, if supported by Vault, echoed in the batch + * results. + * @param reference the correlation identifier; can be {@code null} or empty. + * @return {@code this} builder instance . + */ + public VaultTransformRequestBuilder reference(String reference) { + this.reference = reference; + return this; + } + /** * Build a new {@link VaultTransformContext} instance. * @return a new {@link VaultTransformContext}. */ public VaultTransformContext build() { - return new VaultTransformContext(this.transformation, this.tweak); + return new VaultTransformContext(this.transformation, this.tweak, this.reference); } } diff --git a/spring-vault-core/src/test/java/org/springframework/vault/core/VaultTransformTemplateIntegrationTests.java b/spring-vault-core/src/test/java/org/springframework/vault/core/VaultTransformTemplateIntegrationTests.java index 1eea959e..cbff83bb 100644 --- a/spring-vault-core/src/test/java/org/springframework/vault/core/VaultTransformTemplateIntegrationTests.java +++ b/spring-vault-core/src/test/java/org/springframework/vault/core/VaultTransformTemplateIntegrationTests.java @@ -15,6 +15,7 @@ */ package org.springframework.vault.core; +import java.util.ArrayList; import java.util.Arrays; import java.util.List; import java.util.stream.Collectors; @@ -248,4 +249,37 @@ class VaultTransformTemplateIntegrationTests extends IntegrationTestSupport { } } -} + @Test + void batchEncodeAndDecodeWithReference() { + // Prepare test data + List batch = new ArrayList<>(); + batch.add(TransformPlaintext.of("123-45-6789") + .with(VaultTransformContext.builder().transformation("myssn").reference("ref-1").build())); + batch.add(TransformPlaintext.of("234-56-7890") + .with(VaultTransformContext.builder().transformation("myssn").reference("ref-2").build())); + + // Encode + List encodeResults = transformOperations.encode("myrole", batch); + + // Verify encode results + assertThat(encodeResults).hasSize(2); + assertThat(encodeResults.get(0).isSuccessful()).isTrue(); + assertThat(encodeResults.get(1).isSuccessful()).isTrue(); + + // Prepare decode batch + List ciphertexts = new ArrayList<>(); + ciphertexts.add(encodeResults.get(0).get()); + ciphertexts.add(encodeResults.get(1).get()); + + // Decode + List decodeResults = transformOperations.decode("myrole", ciphertexts); + + // Verify decode results + assertThat(decodeResults).hasSize(2); + assertThat(decodeResults.get(0).get().asString()).isEqualTo("123-45-6789"); + assertThat(decodeResults.get(1).get().asString()).isEqualTo("234-56-7890"); + assertThat(decodeResults.get(0).get().getContext().getReference()).isEqualTo("ref-1"); + assertThat(decodeResults.get(1).get().getContext().getReference()).isEqualTo("ref-2"); + } + +} \ No newline at end of file diff --git a/spring-vault-core/src/test/java/org/springframework/vault/support/VaultTransformContextUnitTests.java b/spring-vault-core/src/test/java/org/springframework/vault/support/VaultTransformContextUnitTests.java index e5d1f7b3..459a20aa 100644 --- a/spring-vault-core/src/test/java/org/springframework/vault/support/VaultTransformContextUnitTests.java +++ b/spring-vault-core/src/test/java/org/springframework/vault/support/VaultTransformContextUnitTests.java @@ -58,4 +58,22 @@ class VaultTransformContextUnitTests { assertThat(context.getTransformation()).isEmpty(); } + @Test + void createsContextWithReference() { + + String transformName = "some_transformation"; + byte[] tweak = { 1, 2, 3, 4, 5, 6, 7 }; + String referenceValue = "my-reference"; + + VaultTransformContext context = VaultTransformContext.builder() + .transformation(transformName) + .tweak(tweak) + .reference(referenceValue) + .build(); + + assertThat(context.getTransformation()).isEqualTo(transformName); + assertThat(context.getTweak()).isEqualTo(tweak); + assertThat(context.getReference()).isEqualTo(referenceValue); + } + }