Merge pull request #16972 from kstrijbos

* pr/16972:
  Polish "Make it easier to set bufferRequestBody"
  Make it easier to set bufferRequestBody

Closes gh-16972

spring-boot-project/spring-boot/src/main/java/org/springframework/boot/web/client/RestTemplateBuilder.java

@@ -33,8 +33,11 @@ import java.util.function.Supplier;
 
 import org.springframework.beans.BeanUtils;
 import org.springframework.http.client.AbstractClientHttpRequestFactoryWrapper;
+import org.springframework.http.client.ClientHttpRequest;
 import org.springframework.http.client.ClientHttpRequestFactory;
 import org.springframework.http.client.ClientHttpRequestInterceptor;
+import org.springframework.http.client.HttpComponentsClientHttpRequestFactory;
+import org.springframework.http.client.SimpleClientHttpRequestFactory;
 import org.springframework.http.converter.HttpMessageConverter;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
@@ -60,6 +63,7 @@ import org.springframework.web.util.UriTemplateHandler;
  * @author Andy Wilkinson
  * @author Brian Clozel
  * @author Dmytro Nosan
+ * @author Kevin Strijbos
  * @since 1.4.0
  */
 public class RestTemplateBuilder {
@@ -506,6 +510,24 @@ public class RestTemplateBuilder {
 				this.interceptors);
 	}
 
+	/**
+	 * Sets if the underling {@link ClientHttpRequestFactory} should buffer the
+	 * {@linkplain ClientHttpRequest#getBody() request body} internally.
+	 * @param bufferRequestBody value of the bufferRequestBody parameter
+	 * @return a new builder instance.
+	 * @since 2.2.0
+	 * @see SimpleClientHttpRequestFactory#setBufferRequestBody(boolean)
+	 * @see HttpComponentsClientHttpRequestFactory#setBufferRequestBody(boolean)
+	 */
+	public RestTemplateBuilder setBufferRequestBody(boolean bufferRequestBody) {
+		return new RestTemplateBuilder(this.detectRequestFactory, this.rootUri,
+				this.messageConverters, this.requestFactorySupplier,
+				this.uriTemplateHandler, this.errorHandler, this.basicAuthentication,
+				this.restTemplateCustomizers,
+				this.requestFactoryCustomizer.bufferRequestBody(bufferRequestBody),
+				this.interceptors);
+	}
+
 	/**
 	 * Build a new {@link RestTemplate} instance and configure it using this builder.
 	 * @return a configured {@link RestTemplate} instance.
@@ -617,21 +639,32 @@ public class RestTemplateBuilder {
 
 		private final Duration readTimeout;
 
+		private final Boolean bufferRequestBody;
+
 		RequestFactoryCustomizer() {
-			this(null, null);
+			this(null, null, null);
 		}
 
-		private RequestFactoryCustomizer(Duration connectTimeout, Duration readTimeout) {
+		private RequestFactoryCustomizer(Duration connectTimeout, Duration readTimeout,
+				Boolean bufferRequestBody) {
 			this.connectTimeout = connectTimeout;
 			this.readTimeout = readTimeout;
+			this.bufferRequestBody = bufferRequestBody;
 		}
 
 		public RequestFactoryCustomizer connectTimeout(Duration connectTimeout) {
-			return new RequestFactoryCustomizer(connectTimeout, this.readTimeout);
+			return new RequestFactoryCustomizer(connectTimeout, this.readTimeout,
+					this.bufferRequestBody);
 		}
 
 		public RequestFactoryCustomizer readTimeout(Duration readTimeout) {
-			return new RequestFactoryCustomizer(this.connectTimeout, readTimeout);
+			return new RequestFactoryCustomizer(this.connectTimeout, readTimeout,
+					this.bufferRequestBody);
+		}
+
+		public RequestFactoryCustomizer bufferRequestBody(boolean bufferRequestBody) {
+			return new RequestFactoryCustomizer(this.connectTimeout, this.readTimeout,
+					bufferRequestBody);
 		}
 
 		@Override
@@ -639,12 +672,13 @@ public class RestTemplateBuilder {
 			ClientHttpRequestFactory unwrappedRequestFactory = unwrapRequestFactoryIfNecessary(
 					requestFactory);
 			if (this.connectTimeout != null) {
-				new TimeoutRequestFactoryCustomizer(this.connectTimeout,
-						"setConnectTimeout").customize(unwrappedRequestFactory);
+				setConnectTimeout(unwrappedRequestFactory);
 			}
 			if (this.readTimeout != null) {
-				new TimeoutRequestFactoryCustomizer(this.readTimeout, "setReadTimeout")
-						.customize(unwrappedRequestFactory);
+				setReadTimeout(unwrappedRequestFactory);
+			}
+			if (this.bufferRequestBody != null) {
+				setBufferRequestBody(unwrappedRequestFactory);
 			}
 		}
 
@@ -664,35 +698,37 @@ public class RestTemplateBuilder {
 			return unwrappedRequestFactory;
 		}
 
-		/**
-		 * {@link ClientHttpRequestFactory} customizer to call a "set timeout" method.
-		 */
-		private static final class TimeoutRequestFactoryCustomizer {
-
-			private final Duration timeout;
-
-			private final String methodName;
+		private void setConnectTimeout(ClientHttpRequestFactory factory) {
+			Method method = findMethod(factory, "setConnectTimeout", int.class);
+			int timeout = Math.toIntExact(this.connectTimeout.toMillis());
+			invoke(factory, method, timeout);
+		}
 
-			TimeoutRequestFactoryCustomizer(Duration timeout, String methodName) {
-				this.timeout = timeout;
-				this.methodName = methodName;
-			}
+		private void setReadTimeout(ClientHttpRequestFactory factory) {
+			Method method = findMethod(factory, "setReadTimeout", int.class);
+			int timeout = Math.toIntExact(this.readTimeout.toMillis());
+			invoke(factory, method, timeout);
+		}
 
-			void customize(ClientHttpRequestFactory factory) {
-				ReflectionUtils.invokeMethod(findMethod(factory), factory,
-						Math.toIntExact(this.timeout.toMillis()));
-			}
+		private void setBufferRequestBody(ClientHttpRequestFactory factory) {
+			Method method = findMethod(factory, "setBufferRequestBody", boolean.class);
+			invoke(factory, method, this.bufferRequestBody);
+		}
 
-			private Method findMethod(ClientHttpRequestFactory factory) {
-				Method method = ReflectionUtils.findMethod(factory.getClass(),
-						this.methodName, int.class);
-				if (method != null) {
-					return method;
-				}
-				throw new IllegalStateException("Request factory " + factory.getClass()
-						+ " does not have a " + this.methodName + "(int) method");
+		private Method findMethod(ClientHttpRequestFactory requestFactory,
+				String methodName, Class<?>... parameters) {
+			Method method = ReflectionUtils.findMethod(requestFactory.getClass(),
+					methodName, parameters);
+			if (method != null) {
+				return method;
 			}
+			throw new IllegalStateException("Request factory " + requestFactory.getClass()
+					+ " does not have a suitable " + methodName + " method");
+		}
 
+		private void invoke(ClientHttpRequestFactory requestFactory, Method method,
+				Object... parameters) {
+			ReflectionUtils.invokeMethod(method, requestFactory, parameters);
 		}
 
 	}

spring-boot-project/spring-boot/src/test/java/org/springframework/boot/web/client/RestTemplateBuilderTests.java

@@ -47,6 +47,7 @@ import org.springframework.web.util.UriTemplateHandler;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
+import static org.assertj.core.api.Assertions.assertThatIllegalStateException;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
@@ -63,6 +64,7 @@ import static org.springframework.test.web.client.response.MockRestResponseCreat
  * @author Phillip Webb
  * @author Andy Wilkinson
  * @author Dmytro Nosan
+ * @author Kevin Strijbos
  */
 public class RestTemplateBuilderTests {
 
@@ -480,6 +482,23 @@ public class RestTemplateBuilderTests {
 				"requestConfig")).getSocketTimeout()).isEqualTo(1234);
 	}
 
+	@Test
+	public void bufferRequestBodyCanBeConfiguredOnHttpComponentsRequestFactory() {
+		ClientHttpRequestFactory requestFactory = this.builder
+				.requestFactory(HttpComponentsClientHttpRequestFactory.class)
+				.setBufferRequestBody(false).build().getRequestFactory();
+		assertThat(requestFactory).hasFieldOrPropertyWithValue("bufferRequestBody",
+				false);
+		requestFactory = this.builder
+				.requestFactory(HttpComponentsClientHttpRequestFactory.class)
+				.setBufferRequestBody(true).build().getRequestFactory();
+		assertThat(requestFactory).hasFieldOrPropertyWithValue("bufferRequestBody", true);
+		requestFactory = this.builder
+				.requestFactory(HttpComponentsClientHttpRequestFactory.class).build()
+				.getRequestFactory();
+		assertThat(requestFactory).hasFieldOrPropertyWithValue("bufferRequestBody", true);
+	}
+
 	@Test
 	public void connectTimeoutCanBeConfiguredOnSimpleRequestFactory() {
 		ClientHttpRequestFactory requestFactory = this.builder
@@ -496,6 +515,21 @@ public class RestTemplateBuilderTests {
 		assertThat(requestFactory).hasFieldOrPropertyWithValue("readTimeout", 1234);
 	}
 
+	@Test
+	public void bufferRequestBodyCanBeConfiguredOnSimpleRequestFactory() {
+		ClientHttpRequestFactory requestFactory = this.builder
+				.requestFactory(SimpleClientHttpRequestFactory.class)
+				.setBufferRequestBody(false).build().getRequestFactory();
+		assertThat(requestFactory).hasFieldOrPropertyWithValue("bufferRequestBody",
+				false);
+		requestFactory = this.builder.requestFactory(SimpleClientHttpRequestFactory.class)
+				.setBufferRequestBody(true).build().getRequestFactory();
+		assertThat(requestFactory).hasFieldOrPropertyWithValue("bufferRequestBody", true);
+		requestFactory = this.builder.requestFactory(SimpleClientHttpRequestFactory.class)
+				.build().getRequestFactory();
+		assertThat(requestFactory).hasFieldOrPropertyWithValue("bufferRequestBody", true);
+	}
+
 	@Test
 	public void connectTimeoutCanBeConfiguredOnOkHttp3RequestFactory() {
 		ClientHttpRequestFactory requestFactory = this.builder
@@ -516,6 +550,15 @@ public class RestTemplateBuilderTests {
 						.isEqualTo(1234);
 	}
 
+	@Test
+	public void bufferRequestBodyCanNotBeConfiguredOnOkHttp3RequestFactory() {
+		assertThatIllegalStateException()
+				.isThrownBy(() -> this.builder
+						.requestFactory(OkHttp3ClientHttpRequestFactory.class)
+						.setBufferRequestBody(false).build().getRequestFactory())
+				.withMessageContaining(OkHttp3ClientHttpRequestFactory.class.getName());
+	}
+
 	@Test
 	public void connectTimeoutCanBeConfiguredOnAWrappedRequestFactory() {
 		SimpleClientHttpRequestFactory requestFactory = new SimpleClientHttpRequestFactory();
@@ -536,6 +579,27 @@ public class RestTemplateBuilderTests {
 		assertThat(requestFactory).hasFieldOrPropertyWithValue("readTimeout", 1234);
 	}
 
+	@Test
+	public void bufferRequestBodyCanBeConfiguredOnAWrappedRequestFactory() {
+		SimpleClientHttpRequestFactory requestFactory = new SimpleClientHttpRequestFactory();
+		this.builder
+				.requestFactory(
+						() -> new BufferingClientHttpRequestFactory(requestFactory))
+				.setBufferRequestBody(false).build();
+		assertThat(requestFactory).hasFieldOrPropertyWithValue("bufferRequestBody",
+				false);
+		this.builder
+				.requestFactory(
+						() -> new BufferingClientHttpRequestFactory(requestFactory))
+				.setBufferRequestBody(true).build();
+		assertThat(requestFactory).hasFieldOrPropertyWithValue("bufferRequestBody", true);
+		this.builder
+				.requestFactory(
+						() -> new BufferingClientHttpRequestFactory(requestFactory))
+				.build();
+		assertThat(requestFactory).hasFieldOrPropertyWithValue("bufferRequestBody", true);
+	}
+
 	@Test
 	public void unwrappingDoesNotAffectRequestFactoryThatIsSetOnTheBuiltTemplate() {
 		SimpleClientHttpRequestFactory requestFactory = new SimpleClientHttpRequestFactory();