From aed5a74e098c3a82777007c5bf509b1359b9c487 Mon Sep 17 00:00:00 2001
From: feijermu <seedeed@qq.com>
Date: Mon, 30 Mar 2020 17:26:03 +0800
Subject: [PATCH] Release the ByteBuf when IOException occurs in
 AbstractMemoryHttpData. (#10133)

Motivation:

An IOException may be thrown from InputStream.read or checkSize method, and cause the ByteBuf leak.

Modification:

Release the ByteBuf when IOException occurs.

Result:
Avoid ByteBuf leak.
---
 .../multipart/AbstractMemoryHttpData.java     | 20 ++++++++++------
 .../multipart/AbstractMemoryHttpDataTest.java | 23 +++++++++++++++++++
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/codec-http/src/main/java/io/netty/handler/codec/http/multipart/AbstractMemoryHttpData.java b/codec-http/src/main/java/io/netty/handler/codec/http/multipart/AbstractMemoryHttpData.java
index a08cdad70a..bc925c2466 100644
--- a/codec-http/src/main/java/io/netty/handler/codec/http/multipart/AbstractMemoryHttpData.java
+++ b/codec-http/src/main/java/io/netty/handler/codec/http/multipart/AbstractMemoryHttpData.java
@@ -66,18 +66,24 @@ public abstract class AbstractMemoryHttpData extends AbstractHttpData {
     public void setContent(InputStream inputStream) throws IOException {
         ObjectUtil.checkNotNull(inputStream, "inputStream");
 
-        ByteBuf buffer = buffer();
         byte[] bytes = new byte[4096 * 4];
-        int read = inputStream.read(bytes);
+        ByteBuf buffer = buffer();
         int written = 0;
-        while (read > 0) {
-            buffer.writeBytes(bytes, 0, read);
-            written += read;
-            checkSize(written);
-            read = inputStream.read(bytes);
+        try {
+            int read = inputStream.read(bytes);
+            while (read > 0) {
+                buffer.writeBytes(bytes, 0, read);
+                written += read;
+                checkSize(written);
+                read = inputStream.read(bytes);
+            }
+        } catch (IOException e) {
+            buffer.release();
+            throw e;
         }
         size = written;
         if (definedSize > 0 && definedSize < size) {
+            buffer.release();
             throw new IOException("Out of size: " + size + " > " + definedSize);
         }
         if (byteBuf != null) {
diff --git a/codec-http/src/test/java/io/netty/handler/codec/http/multipart/AbstractMemoryHttpDataTest.java b/codec-http/src/test/java/io/netty/handler/codec/http/multipart/AbstractMemoryHttpDataTest.java
index df4a425de5..5e6d0a23a5 100644
--- a/codec-http/src/test/java/io/netty/handler/codec/http/multipart/AbstractMemoryHttpDataTest.java
+++ b/codec-http/src/test/java/io/netty/handler/codec/http/multipart/AbstractMemoryHttpDataTest.java
@@ -16,9 +16,15 @@
 package io.netty.handler.codec.http.multipart;
 
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufInputStream;
+import io.netty.buffer.ByteBufUtil;
+import io.netty.buffer.Unpooled;
+import io.netty.buffer.UnpooledByteBufAllocator;
+
 import org.junit.Test;
 
 import java.io.ByteArrayInputStream;
+import java.io.IOException;
 import java.nio.charset.Charset;
 import java.security.SecureRandom;
 import java.util.Arrays;
@@ -36,6 +42,22 @@ public class AbstractMemoryHttpDataTest {
      */
     @Test
     public void testSetContentFromStream() throws Exception {
+        // definedSize=0
+        TestHttpData test = new TestHttpData("test", UTF_8, 0);
+        String contentStr = "foo_test";
+        ByteBuf buf = Unpooled.wrappedBuffer(contentStr.getBytes(UTF_8));
+        buf.markReaderIndex();
+        ByteBufInputStream is = new ByteBufInputStream(buf);
+        try {
+            test.setContent(is);
+            assertFalse(buf.isReadable());
+            assertEquals(test.getString(UTF_8), contentStr);
+            buf.resetReaderIndex();
+            assertTrue(ByteBufUtil.equals(buf, test.getByteBuf()));
+        } finally {
+            is.close();
+        }
+
         Random random = new SecureRandom();
 
         for (int i = 0; i < 20; i++) {
@@ -56,6 +78,7 @@ public class AbstractMemoryHttpDataTest {
             assertEquals(0, buffer.readerIndex());
             assertEquals(bytes.length, buffer.writerIndex());
             assertArrayEquals(bytes, Arrays.copyOf(buffer.array(), bytes.length));
+            assertArrayEquals(bytes, data.get());
         }
     }