Polishing.

Use PropertyMapper to map optional properties. Introduce mutate() method to VaultTransformContext to avoid external copy usage.

Original pull request: gh-897
See gh-894

Author: mp911de

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

@@ -99,7 +99,7 @@ public <T> Source<T> from(Supplier<T> supplier) {
 	 * @param value the value
 	 * @return a {@link Source} that can be used to complete the mapping
 	 */
-	public <T> Source<T> from(T value) {
+	public <T> Source<T> from(@Nullable T value) {
 		return from(() -> value);
 	}
 

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

@@ -23,7 +23,6 @@
 import java.util.Map;
 
 import org.springframework.util.Assert;
-import org.springframework.util.ObjectUtils;
 import org.springframework.util.StringUtils;
 import org.springframework.vault.VaultException;
 import org.springframework.vault.support.TransformCiphertext;
@@ -38,6 +37,7 @@
  *
  * @author Lauren Voswinkel
  * @author Mark Paluch
+ * @author Roopesh Chandran
  * @since 2.3
  */
 public class VaultTransformTemplate implements VaultTransformOperations {
@@ -174,17 +174,11 @@ public List<VaultTransformDecodeResult> decode(String roleName, List<TransformCi
 
 	private static void applyTransformOptions(VaultTransformContext context, Map<String, String> request) {
 
-		if (!ObjectUtils.isEmpty(context.getTransformation())) {
-			request.put("transformation", context.getTransformation());
-		}
+		PropertyMapper mapper = PropertyMapper.get();
 
-		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());
-		}
+		mapper.from(context.getTransformation()).whenNotEmpty().to("transformation", request);
+		mapper.from(context.getTweak()).whenNotEmpty().as(Base64.getEncoder()::encodeToString).to("tweak", request);
+		mapper.from(context.getReference()).whenNotEmpty().to("reference", request);
 	}
 
 	private static List<VaultTransformEncodeResult> toEncodedResults(VaultResponse vaultResponse,
@@ -250,29 +244,10 @@ private static VaultTransformDecodeResult getDecryptionResult(Map<String, String
 
 		if (StringUtils.hasText(data.get("decoded_value"))) {
 
-			// 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.
+			VaultTransformContext updatedContext = postProcessTransformContext(data, ciphertext.getContext());
 			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()));
@@ -281,24 +256,29 @@ private static VaultTransformDecodeResult getDecryptionResult(Map<String, String
 	private static TransformCiphertext toCiphertext(Map<String, ?> data, VaultTransformContext context) {
 
 		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();
-		}
+		VaultTransformContext contextToUse = postProcessTransformContext(data, context);
 
 		return contextToUse.isEmpty() ? TransformCiphertext.of(ciphertext)
 				: TransformCiphertext.of(ciphertext).with(contextToUse);
 	}
 
+	private static VaultTransformContext postProcessTransformContext(Map<String, ?> data,
+			VaultTransformContext context) {
+
+		if (data.containsKey("tweak") || data.containsKey("reference")) {
+
+			PropertyMapper mapper = PropertyMapper.get();
+			VaultTransformContext.VaultTransformRequestBuilder builder = context.mutate();
+
+			mapper.from((String) data.get("tweak")).whenNotEmpty().as(Base64.getDecoder()::decode).to(builder::tweak);
+			mapper.from((String) data.get("reference")).whenNotEmpty().to(builder::reference);
+
+			return builder.build();
+		}
+
+		return context;
+	}
+
 	@SuppressWarnings("unchecked")
 	private static List<Map<String, String>> getBatchData(VaultResponse vaultResponse) {
 		return (List<Map<String, String>>) vaultResponse.getRequiredData().get("batch_results");

spring-vault-core/src/main/java/org/springframework/vault/support/VaultTransformContext.java

@@ -17,13 +17,15 @@
 
 import java.util.Arrays;
 
+import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.util.ObjectUtils;
 
 /**
  * Transform backend encode/decode context object.
  *
  * @author Lauren Voswinkel
+ * @author Roopesh Chandran
  * @since 2.3
  */
 public class VaultTransformContext {
@@ -32,15 +34,15 @@ 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], null);
 
 	private final String transformation;
 
 	private final byte[] tweak;
 
-	private final String reference;
+	private final @Nullable String reference;
 
-	private VaultTransformContext(String transformation, byte[] tweak, String reference) {
+	private VaultTransformContext(String transformation, byte[] tweak, @Nullable String reference) {
 		this.transformation = transformation;
 		this.tweak = tweak;
 		this.reference = reference;
@@ -84,7 +86,8 @@ public static VaultTransformContext fromTweak(byte[] tweak) {
 	 * @return {@code true} if this object is empty.
 	 */
 	public boolean isEmpty() {
-		return ObjectUtils.isEmpty(this.transformation) && ObjectUtils.isEmpty(this.tweak);
+		return ObjectUtils.isEmpty(this.transformation) && ObjectUtils.isEmpty(this.tweak)
+				&& ObjectUtils.isEmpty(this.reference);
 	}
 
 	/**
@@ -102,25 +105,36 @@ public byte[] getTweak() {
 	}
 
 	/**
-	 * @return The reference identifier for batch operations
+	 * @return the reference identifier for batch operations.
+	 * @since 3.2
 	 */
-	public String getReference() {
+	public @Nullable String getReference() {
 		return this.reference;
 	}
 
+	/**
+	 * Return a builder to create a new {@code VaultTransformContext} whose settings are
+	 * replicated from the current {@code VaultTransformContext}.
+	 */
+	public VaultTransformRequestBuilder mutate() {
+		return new VaultTransformRequestBuilder(this);
+	}
+
 	@Override
 	public boolean equals(Object o) {
 		if (this == o)
 			return true;
 		if (!(o instanceof VaultTransformContext that))
 			return false;
-		return this.transformation.equals(that.transformation) && Arrays.equals(this.tweak, that.tweak);
+		return this.transformation.equals(that.transformation) && Arrays.equals(this.tweak, that.tweak)
+				&& ObjectUtils.nullSafeEquals(this.reference, that.reference);
 	}
 
 	@Override
 	public int hashCode() {
 		int result = this.transformation.hashCode();
 		result = 31 * result + Arrays.hashCode(this.tweak);
+		result = 31 * result + ObjectUtils.nullSafeHash(this.reference);
 		return result;
 	}
 
@@ -136,20 +150,22 @@ public static class VaultTransformRequestBuilder {
 		/**
 		 * 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 @Nullable String reference;
 
 		private VaultTransformRequestBuilder() {
 		}
 
+		private VaultTransformRequestBuilder(VaultTransformContext context) {
+			this.transformation = context.transformation;
+			this.tweak = Arrays.copyOf(context.tweak, context.tweak.length);
+			this.reference = context.reference;
+		}
+
 		/**
 		 * Configure a transformation to be used with the {@code transform} operation.
 		 * @param transformation name, provided as a String.
@@ -185,8 +201,9 @@ public VaultTransformRequestBuilder tweak(byte[] tweak) {
 		 * 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 .
+		 * @param reference the correlation identifier; can be or empty.
+		 * @return {@code this} builder instance.
+		 * @since 3.2
 		 */
 		public VaultTransformRequestBuilder reference(String reference) {
 			this.reference = reference;

spring-vault-core/src/test/java/org/springframework/vault/core/VaultTransformTemplateIntegrationTests.java

@@ -251,6 +251,7 @@ void batchEncodeAndDecodeYieldsStartingResultsForInternalWithNoContext() {
 
 	@Test
 	void batchEncodeAndDecodeWithReference() {
+
 		// Prepare test data
 		List<TransformPlaintext> batch = new ArrayList<>();
 		batch.add(TransformPlaintext.of("123-45-6789")
@@ -282,4 +283,4 @@ void batchEncodeAndDecodeWithReference() {
 		assertThat(decodeResults.get(1).get().getContext().getReference()).isEqualTo("ref-2");
 	}
 
-}
+}

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

@@ -23,6 +23,8 @@
  * Unit tests for {@link VaultTransitContext}.
  *
  * @author Lauren Voswinkel
+ * @author Roopesh Chandran
+ * @author Mark Paluch
  */
 class VaultTransformContextUnitTests {
 
@@ -61,6 +63,19 @@ void createsFromTweak() {
 	@Test
 	void createsContextWithReference() {
 
+		String referenceValue = "my-reference";
+
+		VaultTransformContext context = VaultTransformContext.builder().reference(referenceValue).build();
+
+		assertThat(context.getReference()).isEqualTo(referenceValue);
+		assertThat(context).isEqualTo(context).isNotEqualTo(VaultTransformContext.fromTweak(new byte[] { 1 }));
+		assertThat(context).hasSameHashCodeAs(context)
+			.doesNotHaveSameHashCodeAs(VaultTransformContext.fromTweak(new byte[] { 1 }));
+	}
+
+	@Test
+	void appliesMutation() {
+
 		String transformName = "some_transformation";
 		byte[] tweak = { 1, 2, 3, 4, 5, 6, 7 };
 		String referenceValue = "my-reference";
@@ -69,6 +84,8 @@ void createsContextWithReference() {
 			.transformation(transformName)
 			.tweak(tweak)
 			.reference(referenceValue)
+			.build()
+			.mutate()
 			.build();
 
 		assertThat(context.getTransformation()).isEqualTo(transformName);