Skip to content

Commit

Permalink
fix okhttpclient response body leak bug (#600)
Browse files Browse the repository at this point in the history
  • Loading branch information
Nicole00 authored Jun 10, 2024
1 parent 9346aeb commit 6b25e64
Showing 1 changed file with 17 additions and 10 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -16,14 +17,16 @@
import okhttp3.Response;
import okhttp3.ResponseBody;
import java.io.IOException;
import java.io.InputStream;
import java.util.HashMap;
import java.util.Map;

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();

Expand Down Expand Up @@ -75,12 +78,8 @@ public void open() {

public void close() {
try {
if (responseBody != null) {
responseBody.close();
responseBody = null;
}

requestBuffer.close();
inputStream.close();
} catch (IOException e) {
LOGGER.warn(e.getMessage());
}
Expand All @@ -92,11 +91,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.");
}
Expand All @@ -118,6 +117,7 @@ public void flush() throws TTransportException {
// Extract request and reset buffer
byte[] data = requestBuffer.toByteArray();
requestBuffer.reset();
Response response = null;
try {

// Create request object
Expand All @@ -133,15 +133,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();
}
}
}

Expand Down

0 comments on commit 6b25e64

Please sign in to comment.