Add reference support in Transform FPE.

Original pull request: gh-897
Closes gh-894

Signed-off-by: Roopesh Chandran <148393707+rchandran80@users.noreply.github.com>
This commit is contained in:
Roopesh Chandran
2025-02-14 15:12:02 -05:00
committed by Mark Paluch
parent 51d88751fb
commit 596e3b4d79
4 changed files with 127 additions and 10 deletions

View File

@@ -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<VaultTransformEncodeResult> 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();
}

View File

@@ -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}.
* <br/>
* <br/>
*
* <p>
* 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.
* </p>
*/
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);
}
}

View File

@@ -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<TransformPlaintext> 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<VaultTransformEncodeResult> 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<TransformCiphertext> ciphertexts = new ArrayList<>();
ciphertexts.add(encodeResults.get(0).get());
ciphertexts.add(encodeResults.get(1).get());
// Decode
List<VaultTransformDecodeResult> 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");
}
}

View File

@@ -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);
}
}