Fix memory leaks in ProtobufDecoder

Issue: SPR-17418

Author: poutsma

spring-core/src/test/java/org/springframework/core/io/buffer/AbstractDataBufferAllocatingTestCase.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 the original author or authors.
+ * Copyright 2002-2018 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.
@@ -66,9 +66,12 @@ protected DataBuffer createDataBuffer(int capacity) {
 	}
 
 	protected DataBuffer stringBuffer(String value) {
-		byte[] bytes = value.getBytes(StandardCharsets.UTF_8);
-		DataBuffer buffer = this.bufferFactory.allocateBuffer(bytes.length);
-		buffer.write(bytes);
+		return byteBuffer(value.getBytes(StandardCharsets.UTF_8));
+	}
+
+	protected DataBuffer byteBuffer(byte[] value) {
+		DataBuffer buffer = this.bufferFactory.allocateBuffer(value.length);
+		buffer.write(value);
 		return buffer;
 	}
 

spring-web/src/main/java/org/springframework/http/codec/protobuf/ProtobufDecoder.java

@@ -216,12 +216,18 @@ public Iterable<? extends Message> apply(DataBuffer input) {
 				} while (remainingBytesToRead > 0);
 				return messages;
 			}
+			catch (DecodingException ex) {
+				throw ex;
+			}
 			catch (IOException ex) {
 				throw new DecodingException("I/O error while parsing input stream", ex);
 			}
 			catch (Exception ex) {
 				throw new DecodingException("Could not read Protobuf message: " + ex.getMessage(), ex);
 			}
+			finally {
+				DataBufferUtils.release(input);
+			}
 		}
 	}
 

spring-web/src/test/java/org/springframework/http/codec/protobuf/ProtobufDecoderTests.java

@@ -36,8 +36,7 @@
 import org.springframework.util.MimeType;
 
 import static java.util.Collections.emptyMap;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.*;
 import static org.springframework.core.ResolvableType.forClass;
 
 /**
@@ -81,7 +80,7 @@ public void canDecode() {
 
 	@Test
 	public void decodeToMono() {
-		DataBuffer data = this.bufferFactory.wrap(testMsg.toByteArray());
+		DataBuffer data = byteBuffer(testMsg.toByteArray());
 		ResolvableType elementType = forClass(Msg.class);
 
 		Mono<Message> mono = this.decoder.decodeToMono(Flux.just(data), elementType, null, emptyMap());
@@ -106,12 +105,12 @@ public void decodeToMonoWithLargerDataBuffer() {
 
 	@Test
 	public void decodeChunksToMono() {
-		DataBuffer buffer = this.bufferFactory.wrap(testMsg.toByteArray());
+		DataBuffer buffer = byteBuffer(testMsg.toByteArray());
 		Flux<DataBuffer> chunks = Flux.just(
-				buffer.slice(0, 4),
-				buffer.slice(4, buffer.readableByteCount() - 4));
-		DataBufferUtils.retain(buffer);
+				DataBufferUtils.retain(buffer.slice(0, 4)),
+				DataBufferUtils.retain(buffer.slice(4, buffer.readableByteCount() - 4)));
 		ResolvableType elementType = forClass(Msg.class);
+		release(buffer);
 
 		Mono<Message> mono = this.decoder.decodeToMono(chunks, elementType, null,
 				emptyMap());
@@ -123,10 +122,11 @@ public void decodeChunksToMono() {
 
 	@Test
 	public void decode() throws IOException {
-		DataBuffer buffer = bufferFactory.allocateBuffer();
+		DataBuffer buffer = this.bufferFactory.allocateBuffer();
 		testMsg.writeDelimitedTo(buffer.asOutputStream());
-		DataBuffer buffer2 = bufferFactory.allocateBuffer();
+		DataBuffer buffer2 = this.bufferFactory.allocateBuffer();
 		testMsg2.writeDelimitedTo(buffer2.asOutputStream());
+
 		Flux<DataBuffer> source = Flux.just(buffer, buffer2);
 		ResolvableType elementType = forClass(Msg.class);
 
@@ -136,22 +136,22 @@ public void decode() throws IOException {
 				.expectNext(testMsg)
 				.expectNext(testMsg2)
 				.verifyComplete();
-
-		DataBufferUtils.release(buffer);
-		DataBufferUtils.release(buffer2);
 	}
 
 	@Test
 	public void decodeSplitChunks() throws IOException {
-		DataBuffer buffer = bufferFactory.allocateBuffer();
+		DataBuffer buffer = this.bufferFactory.allocateBuffer();
 		testMsg.writeDelimitedTo(buffer.asOutputStream());
-		DataBuffer buffer2 = bufferFactory.allocateBuffer();
+		DataBuffer buffer2 = this.bufferFactory.allocateBuffer();
 		testMsg2.writeDelimitedTo(buffer2.asOutputStream());
+
 		Flux<DataBuffer> chunks = Flux.just(
-				buffer.slice(0, 4),
-				buffer.slice(4, buffer.readableByteCount() - 4),
-				buffer2.slice(0, 2),
-				buffer2.slice(2, buffer2.readableByteCount() - 2));
+				DataBufferUtils.retain(buffer.slice(0, 4)),
+				DataBufferUtils.retain(buffer.slice(4, buffer.readableByteCount() - 4)),
+				DataBufferUtils.retain(buffer2.slice(0, 2)),
+				DataBufferUtils.retain(buffer2
+						.slice(2, buffer2.readableByteCount() - 2)));
+		release(buffer, buffer2);
 
 		ResolvableType elementType = forClass(Msg.class);
 		Flux<Message> messages = this.decoder.decode(chunks, elementType, null, emptyMap());
@@ -160,9 +160,6 @@ public void decodeSplitChunks() throws IOException {
 				.expectNext(testMsg)
 				.expectNext(testMsg2)
 				.verifyComplete();
-
-		DataBufferUtils.release(buffer);
-		DataBufferUtils.release(buffer2);
 	}
 
 	@Test
@@ -178,15 +175,12 @@ public void decodeMergedChunks() throws IOException {
 				.expectNext(testMsg)
 				.expectNext(testMsg)
 				.verifyComplete();
-
-		DataBufferUtils.release(buffer);
 	}
 
 	@Test
 	public void exceedMaxSize() {
 		this.decoder.setMaxMessageSize(1);
-		byte[] body = testMsg.toByteArray();
-		Flux<DataBuffer> source = Flux.just(this.bufferFactory.wrap(body));
+		Flux<DataBuffer> source = Flux.just(byteBuffer(testMsg.toByteArray()));
 		ResolvableType elementType = forClass(Msg.class);
 		Flux<Message> messages = this.decoder.decode(source, elementType, null,
 				emptyMap());