From 3259fe8727cbd828020f23a3b3c0603e2855a237 Mon Sep 17 00:00:00 2001
From: Anqi <anqi.wang@vesoft.com>
Date: Fri, 7 Jun 2024 16:34:16 +0800
Subject: [PATCH] fix okhttpclient response body leak bug

---
 .../thrift/transport/THttp2Client.java        | 26 ++++++++++++-------
 1 file changed, 16 insertions(+), 10 deletions(-)

diff --git a/client/src/main/fbthrift/com/facebook/thrift/transport/THttp2Client.java b/client/src/main/fbthrift/com/facebook/thrift/transport/THttp2Client.java
index 334be15d1..ec0003105 100644
--- a/client/src/main/fbthrift/com/facebook/thrift/transport/THttp2Client.java
+++ b/client/src/main/fbthrift/com/facebook/thrift/transport/THttp2Client.java
@@ -6,6 +6,7 @@
 package com.facebook.thrift.transport;
 
 import com.facebook.thrift.utils.Logger;
+import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import javax.net.ssl.SSLSocketFactory;
 import javax.net.ssl.TrustManager;
@@ -16,6 +17,7 @@
 import okhttp3.Response;
 import okhttp3.ResponseBody;
 import java.io.IOException;
+import java.io.InputStream;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -23,7 +25,8 @@ public class THttp2Client extends TTransport {
     private static final Logger LOGGER = Logger.getLogger(THttp2Client.class.getName());
 
     private final ByteArrayOutputStream requestBuffer = new ByteArrayOutputStream();
-    private ResponseBody responseBody = null;
+
+    private InputStream inputStream = null;
     private Map<String, String> customHeaders = null;
     private static final Map<String, String> defaultHeaders = getDefaultHeaders();
 
@@ -75,11 +78,6 @@ public void open() {
 
     public void close() {
         try {
-            if (responseBody != null) {
-                responseBody.close();
-                responseBody = null;
-            }
-
             requestBuffer.close();
         } catch (IOException e) {
             LOGGER.warn(e.getMessage());
@@ -92,11 +90,11 @@ public boolean isOpen() {
     }
 
     public int read(byte[] buf, int off, int len) throws TTransportException {
-        if (responseBody == null) {
+        if (inputStream == null) {
             throw new TTransportException("Response buffer is empty, no request.");
         }
         try {
-            int ret = responseBody.byteStream().read(buf, off, len);
+            int ret = inputStream.read(buf, off, len);
             if (ret == -1) {
                 throw new TTransportException("No more data available.");
             }
@@ -118,6 +116,7 @@ public void flush() throws TTransportException {
         // Extract request and reset buffer
         byte[] data = requestBuffer.toByteArray();
         requestBuffer.reset();
+        Response response = null;
         try {
 
             // Create request object
@@ -133,15 +132,22 @@ public void flush() throws TTransportException {
             Request request = requestBuilder.build();
 
             // Make the request
-            Response response = client.newCall(request).execute();
+            response = client.newCall(request).execute();
             if (!response.isSuccessful()) {
                 throw new TTransportException("HTTP Response code: " + response.code());
             }
 
+            if (response.body() == null) {
+                throw new TTransportException("response body is null");
+            }
             // Read the response
-            responseBody = response.body();
+            inputStream = new ByteArrayInputStream(response.body().bytes());
         } catch (IOException iox) {
             throw new TTransportException(iox);
+        } finally {
+            if (response != null) {
+                response.close();
+            }
         }
     }