From 4a22f8c67cc79baaa9cddab10521eb24573f6ff0 Mon Sep 17 00:00:00 2001
From: Mark Paluch <mpaluch@pivotal.io>
Date: Sun, 9 Dec 2018 16:46:39 +0100
Subject: [PATCH] Remove code duplicates from VaultWrappingTemplate.

Resolves gh-345.
---
 .../vault/core/VaultWrappingTemplate.java     | 46 +++++++++----------
 1 file changed, 22 insertions(+), 24 deletions(-)

diff --git a/spring-vault-core/src/main/java/org/springframework/vault/core/VaultWrappingTemplate.java b/spring-vault-core/src/main/java/org/springframework/vault/core/VaultWrappingTemplate.java
index 057e6ccc..bac7f5af 100644
--- a/spring-vault-core/src/main/java/org/springframework/vault/core/VaultWrappingTemplate.java
+++ b/spring-vault-core/src/main/java/org/springframework/vault/core/VaultWrappingTemplate.java
@@ -21,6 +21,7 @@ import java.time.format.DateTimeFormatter;
 import java.time.temporal.TemporalAccessor;
 import java.util.Collections;
 import java.util.Map;
+import java.util.function.BiFunction;
 
 import org.springframework.core.ParameterizedTypeReference;
 import org.springframework.http.HttpEntity;
@@ -38,6 +39,7 @@ import org.springframework.vault.support.VaultResponseSupport;
 import org.springframework.vault.support.VaultToken;
 import org.springframework.vault.support.WrappedMetadata;
 import org.springframework.web.client.HttpStatusCodeException;
+import org.springframework.web.client.RestOperations;
 
 /**
  * @author Mark Paluch
@@ -89,27 +91,12 @@ public class VaultWrappingTemplate implements VaultWrappingOperations {
 	@Override
 	public VaultResponse read(VaultToken token) {
 
-		return vaultOperations.doWithVault(restOperations -> {
-
-			HttpHeaders headers = VaultHttpHeaders.from(token);
-			try {
-				return restOperations.exchange("sys/wrapping/unwrap", HttpMethod.POST,
-						new HttpEntity<>(headers), VaultResponse.class).getBody();
-			}
-			catch (HttpStatusCodeException e) {
-
-				if (e.getStatusCode() == HttpStatus.NOT_FOUND) {
-					return null;
-				}
-
-				if (e.getStatusCode() == HttpStatus.BAD_REQUEST
-						&& e.getResponseBodyAsString().contains("does not exist")) {
-					return null;
-				}
-
-				throw VaultResponses.buildException(e, "sys/wrapping/unwrap");
-			}
-		});
+		return doUnwrap(
+				token,
+				(restOperations, entity) -> {
+					return restOperations.exchange("sys/wrapping/unwrap",
+							HttpMethod.POST, entity, VaultResponse.class).getBody();
+				});
 	}
 
 	@Nullable
@@ -119,12 +106,23 @@ public class VaultWrappingTemplate implements VaultWrappingOperations {
 		ParameterizedTypeReference<VaultResponseSupport<T>> ref = VaultResponses
 				.getTypeReference(responseType);
 
+		return doUnwrap(
+				token,
+				(restOperations, entity) -> {
+					return restOperations.exchange("sys/wrapping/unwrap",
+							HttpMethod.POST, entity, ref).getBody();
+				});
+	}
+
+	@Nullable
+	private <T extends VaultResponseSupport<?>> T doUnwrap(VaultToken token,
+			BiFunction<RestOperations, HttpEntity<?>, T> requestFunction) {
+
 		return vaultOperations.doWithVault(restOperations -> {
 
-			HttpHeaders headers = VaultHttpHeaders.from(token);
 			try {
-				return restOperations.exchange("sys/wrapping/unwrap", HttpMethod.POST,
-						new HttpEntity<>(headers), ref).getBody();
+				return requestFunction.apply(restOperations, new HttpEntity<>(
+						VaultHttpHeaders.from(token)));
 			}
 			catch (HttpStatusCodeException e) {