From 4c8a1d605188ea4152fb168fe8a3cf9d7ea89f06 Mon Sep 17 00:00:00 2001
From: Violeta Georgieva <violetag@vmware.com>
Date: Mon, 2 Nov 2020 19:05:54 +0200
Subject: [PATCH] Fix memory leak in NettyWriteResponseFilter.wrap (#2019)

NettyWriteResponseFilter.wrap never releases the original poolable
buffer in case of DefaultDataBufferFactory.

Related to https://github.com/reactor/reactor-netty/issues/1359
---
 .../filter/NettyWriteResponseFilter.java      | 19 +++---
 .../filter/NettyWriteResponseFilterTests.java | 65 +++++++++++++++++++
 2 files changed, 76 insertions(+), 8 deletions(-)
 create mode 100644 spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/filter/NettyWriteResponseFilterTests.java

diff --git a/spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/filter/NettyWriteResponseFilter.java b/spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/filter/NettyWriteResponseFilter.java
index e6af21da..6dffb401 100644
--- a/spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/filter/NettyWriteResponseFilter.java
+++ b/spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/filter/NettyWriteResponseFilter.java
@@ -27,6 +27,7 @@ import reactor.netty.Connection;
 
 import org.springframework.core.Ordered;
 import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.DataBufferFactory;
 import org.springframework.core.io.buffer.DefaultDataBufferFactory;
 import org.springframework.core.io.buffer.NettyDataBufferFactory;
 import org.springframework.http.MediaType;
@@ -103,19 +104,21 @@ public class NettyWriteResponseFilter implements GlobalFilter, Ordered {
 	}
 
 	protected DataBuffer wrap(ByteBuf byteBuf, ServerHttpResponse response) {
-		if (response.bufferFactory() instanceof NettyDataBufferFactory) {
-			NettyDataBufferFactory factory = (NettyDataBufferFactory) response
-					.bufferFactory();
+		DataBufferFactory bufferFactory = response.bufferFactory();
+		if (bufferFactory instanceof NettyDataBufferFactory) {
+			NettyDataBufferFactory factory = (NettyDataBufferFactory) bufferFactory;
 			return factory.wrap(byteBuf);
 		}
 		// MockServerHttpResponse creates these
-		else if (response.bufferFactory() instanceof DefaultDataBufferFactory) {
-			DefaultDataBufferFactory factory = (DefaultDataBufferFactory) response
-					.bufferFactory();
-			return factory.wrap(byteBuf.nioBuffer());
+		else if (bufferFactory instanceof DefaultDataBufferFactory) {
+			DataBuffer buffer = ((DefaultDataBufferFactory) bufferFactory)
+					.allocateBuffer(byteBuf.readableBytes());
+			buffer.write(byteBuf.nioBuffer());
+			byteBuf.release();
+			return buffer;
 		}
 		throw new IllegalArgumentException(
-				"Unkown DataBufferFactory type " + response.bufferFactory().getClass());
+				"Unkown DataBufferFactory type " + bufferFactory.getClass());
 	}
 
 	private void cleanup(ServerWebExchange exchange) {
diff --git a/spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/filter/NettyWriteResponseFilterTests.java b/spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/filter/NettyWriteResponseFilterTests.java
new file mode 100644
index 00000000..2e758f1b
--- /dev/null
+++ b/spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/filter/NettyWriteResponseFilterTests.java
@@ -0,0 +1,65 @@
+/*
+ * Copyright 2013-2020 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.filter;
+
+import java.nio.charset.Charset;
+import java.util.ArrayList;
+
+import io.netty.buffer.ByteBuf;
+import org.junit.Test;
+
+import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.NettyDataBufferFactory;
+import org.springframework.core.io.buffer.PooledDataBuffer;
+import org.springframework.mock.http.server.reactive.MockServerHttpResponse;
+
+import static io.netty.buffer.PooledByteBufAllocator.DEFAULT;
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * @author Violeta Georgieva
+ */
+public class NettyWriteResponseFilterTests {
+
+	@Test
+	public void testWrap_NettyDataBufferFactory() {
+		doTestWrap(new MockServerHttpResponse(new NettyDataBufferFactory(DEFAULT)));
+	}
+
+	@Test
+	public void testWrap_DefaultDataBufferFactory() {
+		doTestWrap(new MockServerHttpResponse());
+	}
+
+	private void doTestWrap(MockServerHttpResponse response) {
+		NettyWriteResponseFilter filter = new NettyWriteResponseFilter(new ArrayList<>());
+
+		ByteBuf buffer = DEFAULT.buffer();
+		buffer.writeCharSequence("test", Charset.defaultCharset());
+
+		DataBuffer result = filter.wrap(buffer, response);
+
+		assertThat(result.toString(Charset.defaultCharset())).isEqualTo("test");
+
+		if (result instanceof PooledDataBuffer) {
+			((PooledDataBuffer) result).release();
+		}
+
+		assertThat(buffer.refCnt()).isEqualTo(0);
+	}
+
+}