Merge branch '4.1.x'

Author: spencergibb

spring-cloud-gateway-server-mvc/src/main/java/org/springframework/cloud/gateway/server/mvc/GatewayServerMvcAutoConfiguration.java

@@ -88,8 +88,9 @@ public RestClientCustomizer gatewayRestClientCustomizer(ClientHttpRequestFactory
 
 	@Bean
 	@ConditionalOnMissingBean(ProxyExchange.class)
-	public RestClientProxyExchange restClientProxyExchange(RestClient.Builder restClientBuilder) {
-		return new RestClientProxyExchange(restClientBuilder.build());
+	public RestClientProxyExchange restClientProxyExchange(RestClient.Builder restClientBuilder,
+			GatewayMvcProperties properties) {
+		return new RestClientProxyExchange(restClientBuilder.build(), properties);
 	}
 
 	@Bean

spring-cloud-gateway-server-mvc/src/main/java/org/springframework/cloud/gateway/server/mvc/common/AbstractProxyExchange.java

@@ -0,0 +1,78 @@
+/*
+ * Copyright 2013-2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.gateway.server.mvc.common;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+
+import org.springframework.cloud.gateway.server.mvc.config.GatewayMvcProperties;
+import org.springframework.cloud.gateway.server.mvc.handler.ProxyExchange;
+import org.springframework.http.client.ClientHttpResponse;
+import org.springframework.util.Assert;
+import org.springframework.util.StreamUtils;
+
+public abstract class AbstractProxyExchange implements ProxyExchange {
+
+	private final GatewayMvcProperties properties;
+
+	protected AbstractProxyExchange(GatewayMvcProperties properties) {
+		this.properties = properties;
+	}
+
+	protected int copyResponseBody(ClientHttpResponse clientResponse, InputStream inputStream,
+			OutputStream outputStream) throws IOException {
+		Assert.notNull(clientResponse, "No ClientResponse specified");
+		Assert.notNull(inputStream, "No InputStream specified");
+		Assert.notNull(outputStream, "No OutputStream specified");
+
+		int transferredBytes;
+
+		if (properties.getStreamingMediaTypes().contains(clientResponse.getHeaders().getContentType())) {
+			transferredBytes = copyResponseBodyWithFlushing(inputStream, outputStream);
+		}
+		else {
+			transferredBytes = StreamUtils.copy(inputStream, outputStream);
+		}
+
+		return transferredBytes;
+	}
+
+	private int copyResponseBodyWithFlushing(InputStream inputStream, OutputStream outputStream) throws IOException {
+		int readBytes;
+		var totalReadBytes = 0;
+		var buffer = new byte[properties.getStreamingBufferSize()];
+
+		while ((readBytes = inputStream.read(buffer)) != -1) {
+			outputStream.write(buffer, 0, readBytes);
+			outputStream.flush();
+			if (totalReadBytes < Integer.MAX_VALUE) {
+				try {
+					totalReadBytes = Math.addExact(totalReadBytes, readBytes);
+				}
+				catch (ArithmeticException e) {
+					totalReadBytes = Integer.MAX_VALUE;
+				}
+			}
+		}
+
+		outputStream.flush();
+
+		return totalReadBytes;
+	}
+
+}

spring-cloud-gateway-server-mvc/src/main/java/org/springframework/cloud/gateway/server/mvc/config/GatewayMvcProperties.java

@@ -18,6 +18,7 @@
 
 import java.time.Duration;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.LinkedHashMap;
 import java.util.List;
 
@@ -26,6 +27,7 @@
 
 import org.springframework.boot.context.properties.ConfigurationProperties;
 import org.springframework.core.style.ToStringCreator;
+import org.springframework.http.MediaType;
 
 @ConfigurationProperties(GatewayMvcProperties.PREFIX)
 public class GatewayMvcProperties {
@@ -51,6 +53,18 @@ public class GatewayMvcProperties {
 
 	private HttpClient httpClient = new HttpClient();
 
+	/**
+	 * Mime-types that are streaming.
+	 */
+	private List<MediaType> streamingMediaTypes = Arrays.asList(MediaType.TEXT_EVENT_STREAM,
+			new MediaType("application", "stream+json"), new MediaType("application", "grpc"),
+			new MediaType("application", "grpc+protobuf"), new MediaType("application", "grpc+json"));
+
+	/**
+	 * Buffer size for streaming media mime-types.
+	 */
+	private int streamingBufferSize = 16384;
+
 	public List<RouteProperties> getRoutes() {
 		return routes;
 	}
@@ -71,11 +85,29 @@ public HttpClient getHttpClient() {
 		return httpClient;
 	}
 
+	public List<MediaType> getStreamingMediaTypes() {
+		return streamingMediaTypes;
+	}
+
+	public void setStreamingMediaTypes(List<MediaType> streamingMediaTypes) {
+		this.streamingMediaTypes = streamingMediaTypes;
+	}
+
+	public int getStreamingBufferSize() {
+		return streamingBufferSize;
+	}
+
+	public void setStreamingBufferSize(int streamingBufferSize) {
+		this.streamingBufferSize = streamingBufferSize;
+	}
+
 	@Override
 	public String toString() {
 		return new ToStringCreator(this).append("httpClient", httpClient)
 			.append("routes", routes)
 			.append("routesMap", routesMap)
+			.append("streamingMediaTypes", streamingMediaTypes)
+			.append("streamingBufferSize", streamingBufferSize)
 			.toString();
 	}
 

spring-cloud-gateway-server-mvc/src/main/java/org/springframework/cloud/gateway/server/mvc/handler/ClientHttpRequestFactoryProxyExchange.java

@@ -20,18 +20,28 @@
 import java.io.InputStream;
 import java.io.UncheckedIOException;
 
+import org.springframework.cloud.gateway.server.mvc.common.AbstractProxyExchange;
 import org.springframework.cloud.gateway.server.mvc.common.MvcUtils;
+import org.springframework.cloud.gateway.server.mvc.config.GatewayMvcProperties;
 import org.springframework.http.client.ClientHttpRequest;
 import org.springframework.http.client.ClientHttpRequestFactory;
 import org.springframework.http.client.ClientHttpResponse;
 import org.springframework.util.StreamUtils;
 import org.springframework.web.servlet.function.ServerResponse;
 
-public class ClientHttpRequestFactoryProxyExchange implements ProxyExchange {
+public class ClientHttpRequestFactoryProxyExchange extends AbstractProxyExchange {
 
 	private final ClientHttpRequestFactory requestFactory;
 
+	@Deprecated
 	public ClientHttpRequestFactoryProxyExchange(ClientHttpRequestFactory requestFactory) {
+		super(new GatewayMvcProperties());
+		this.requestFactory = requestFactory;
+	}
+
+	public ClientHttpRequestFactoryProxyExchange(ClientHttpRequestFactory requestFactory,
+			GatewayMvcProperties properties) {
+		super(properties);
 		this.requestFactory = requestFactory;
 	}
 
@@ -54,7 +64,8 @@ public ServerResponse exchange(Request request) {
 						InputStream inputStream = MvcUtils.getAttribute(request.getServerRequest(),
 								MvcUtils.CLIENT_RESPONSE_INPUT_STREAM_ATTR);
 						// copy body from request to clientHttpRequest
-						StreamUtils.copy(inputStream, httpServletResponse.getOutputStream());
+						ClientHttpRequestFactoryProxyExchange.this.copyResponseBody(clientHttpResponse, inputStream,
+								httpServletResponse.getOutputStream());
 					}
 					return null;
 				});

spring-cloud-gateway-server-mvc/src/main/java/org/springframework/cloud/gateway/server/mvc/handler/RestClientProxyExchange.java

@@ -21,17 +21,26 @@
 import java.io.OutputStream;
 import java.io.UncheckedIOException;
 
+import org.springframework.cloud.gateway.server.mvc.common.AbstractProxyExchange;
 import org.springframework.cloud.gateway.server.mvc.common.MvcUtils;
+import org.springframework.cloud.gateway.server.mvc.config.GatewayMvcProperties;
 import org.springframework.http.client.ClientHttpResponse;
 import org.springframework.util.StreamUtils;
 import org.springframework.web.client.RestClient;
 import org.springframework.web.servlet.function.ServerResponse;
 
-public class RestClientProxyExchange implements ProxyExchange {
+public class RestClientProxyExchange extends AbstractProxyExchange {
 
 	private final RestClient restClient;
 
+	@Deprecated
 	public RestClientProxyExchange(RestClient restClient) {
+		super(new GatewayMvcProperties());
+		this.restClient = restClient;
+	}
+
+	public RestClientProxyExchange(RestClient restClient, GatewayMvcProperties properties) {
+		super(properties);
 		this.restClient = restClient;
 	}
 
@@ -59,7 +68,7 @@ private static int copyBody(Request request, OutputStream outputStream) throws I
 		return StreamUtils.copy(request.getServerRequest().servletRequest().getInputStream(), outputStream);
 	}
 
-	private static ServerResponse doExchange(Request request, ClientHttpResponse clientResponse) throws IOException {
+	private ServerResponse doExchange(Request request, ClientHttpResponse clientResponse) throws IOException {
 		InputStream body = clientResponse.getBody();
 		// put the body input stream in a request attribute so filters can read it.
 		MvcUtils.putAttribute(request.getServerRequest(), MvcUtils.CLIENT_RESPONSE_INPUT_STREAM_ATTR, body);
@@ -71,7 +80,8 @@ private static ServerResponse doExchange(Request request, ClientHttpResponse cli
 					InputStream inputStream = MvcUtils.getAttribute(request.getServerRequest(),
 							MvcUtils.CLIENT_RESPONSE_INPUT_STREAM_ATTR);
 					// copy body from request to clientHttpRequest
-					StreamUtils.copy(inputStream, httpServletResponse.getOutputStream());
+					RestClientProxyExchange.this.copyResponseBody(clientResponse, inputStream,
+							httpServletResponse.getOutputStream());
 				}
 				return null;
 			});

spring-cloud-gateway-server-mvc/src/test/java/org/springframework/cloud/gateway/server/mvc/common/AbstractProxyExchangeTests.java

@@ -0,0 +1,99 @@
+/*
+ * Copyright 2013-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.gateway.server.mvc.common;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+
+import org.junit.jupiter.api.Test;
+
+import org.springframework.cloud.gateway.server.mvc.config.GatewayMvcProperties;
+import org.springframework.http.MediaType;
+import org.springframework.mock.http.client.MockClientHttpResponse;
+import org.springframework.web.servlet.function.ServerResponse;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/**
+ * @author Jens Mallien
+ */
+public class AbstractProxyExchangeTests {
+
+	@Test
+	public void copyResponseBodyForJson() throws IOException {
+		MockClientHttpResponse mockResponse = new MockClientHttpResponse(new byte[0], 200);
+		mockResponse.getHeaders().setContentType(MediaType.APPLICATION_JSON);
+
+		InputStream inputStream = mock(InputStream.class);
+		when(inputStream.transferTo(any())).thenReturn(3L);
+		OutputStream outputStream = mock(OutputStream.class);
+
+		int result = new TestProxyExchange().copyResponseBody(mockResponse, inputStream, outputStream);
+
+		assertThat(result).isEqualTo(3);
+		verify(outputStream, times(1)).flush();
+	}
+
+	@Test
+	public void copyResponseBodyForTextEventStream() throws IOException {
+		MockClientHttpResponse mockResponse = new MockClientHttpResponse(new byte[0], 200);
+		mockResponse.getHeaders().setContentType(MediaType.TEXT_EVENT_STREAM);
+
+		InputStream inputStream = mock(InputStream.class);
+		when(inputStream.read(any())).thenReturn(1).thenReturn(1).thenReturn(1).thenReturn(-1);
+		OutputStream outputStream = mock(OutputStream.class);
+
+		int result = new TestProxyExchange().copyResponseBody(mockResponse, inputStream, outputStream);
+
+		assertThat(result).isEqualTo(3);
+		verify(outputStream, times(4)).flush();
+	}
+
+	@Test
+	public void copyResponseBodyWithoutContentType() throws IOException {
+		MockClientHttpResponse mockResponse = new MockClientHttpResponse(new byte[0], 200);
+
+		InputStream inputStream = mock(InputStream.class);
+		when(inputStream.transferTo(any())).thenReturn(3L);
+		OutputStream outputStream = mock(OutputStream.class);
+
+		int result = new TestProxyExchange().copyResponseBody(mockResponse, inputStream, outputStream);
+
+		assertThat(result).isEqualTo(3);
+		verify(outputStream, times(1)).flush();
+	}
+
+	class TestProxyExchange extends AbstractProxyExchange {
+
+		protected TestProxyExchange() {
+			super(new GatewayMvcProperties());
+		}
+
+		@Override
+		public ServerResponse exchange(Request request) {
+			return ServerResponse.ok().build();
+		}
+
+	}
+
+}