Add reference support in Transform FPE.

Original pull request: gh-897
Closes gh-894

Signed-off-by: Roopesh Chandran <[email protected]>

Author: rchandran80

spring-vault-core/src/main/java/org/springframework/vault/core/VaultTransformTemplate.java

@@ -181,6 +181,10 @@ private static void applyTransformOptions(VaultTransformContext context, Map<Str
 		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 @@ private static List<VaultTransformDecodeResult> toDecryptionResults(VaultRespons
 
 		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 @@ private static VaultTransformDecodeResult getDecryptionResult(Map<String, String
 
 		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 @@ private static TransformCiphertext toCiphertext(Map<String, ?> data, VaultTransf
 
 		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();
 		}
 

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 byte[] getTweak() {
 		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 static class VaultTransformRequestBuilder {
 
 		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 VaultTransformRequestBuilder tweak(byte[] tweak) {
 			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);
 		}
 
 	}

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 @@ void batchEncodeAndDecodeYieldsStartingResultsForInternalWithNoContext() {
 		}
 	}
 
-}
+	@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");
+	}
+
+}

spring-vault-core/src/test/java/org/springframework/vault/support/VaultTransformContextUnitTests.java

@@ -58,4 +58,22 @@ void createsFromTweak() {
 		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);
+	}
+
 }