Close all io streams on disconnect

Author: gthea

src/main/java/io/split/android/client/network/HttpResponseConnectionAdapter.java

@@ -1,6 +1,7 @@
 package io.split.android.client.network;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 
 import java.io.ByteArrayInputStream;
@@ -32,6 +33,8 @@ class HttpResponseConnectionAdapter extends HttpsURLConnection {
     private final URL mUrl;
     private final Certificate[] mServerCertificates;
     private OutputStream mOutputStream;
+    private InputStream mInputStream;
+    private InputStream mErrorStream;
     private boolean mDoOutput = false;
 
     /**
@@ -52,11 +55,23 @@ class HttpResponseConnectionAdapter extends HttpsURLConnection {
                                   @NonNull HttpResponse response,
                                   Certificate[] serverCertificates,
                                   @NonNull OutputStream outputStream) {
+        this(url, response, serverCertificates, outputStream, null, null);
+    }
+    
+    @VisibleForTesting
+    HttpResponseConnectionAdapter(@NonNull URL url,
+                                  @NonNull HttpResponse response,
+                                  Certificate[] serverCertificates,
+                                  @NonNull OutputStream outputStream,
+                                  @Nullable InputStream inputStream,
+                                  @Nullable InputStream errorStream) {
         super(url);
         mUrl = url;
         mResponse = response;
         mServerCertificates = serverCertificates;
         mOutputStream = outputStream;
+        mInputStream = inputStream;
+        mErrorStream = errorStream;
     }
 
     @Override
@@ -90,21 +105,27 @@ public InputStream getInputStream() throws IOException {
         if (mResponse.getHttpStatus() >= 400) {
             throw new IOException("HTTP " + mResponse.getHttpStatus());
         }
-        String data = mResponse.getData();
-        if (data == null) {
-            data = "";
+        if (mInputStream == null) {
+            String data = mResponse.getData();
+            if (data == null) {
+                data = "";
+            }
+            mInputStream = new ByteArrayInputStream(data.getBytes(StandardCharsets.UTF_8));
         }
-        return new ByteArrayInputStream(data.getBytes(StandardCharsets.UTF_8));
+        return mInputStream;
     }
 
     @Override
     public InputStream getErrorStream() {
         if (mResponse.getHttpStatus() >= 400) {
-            String data = mResponse.getData();
-            if (data == null) {
-                data = "";
+            if (mErrorStream == null) {
+                String data = mResponse.getData();
+                if (data == null) {
+                    data = "";
+                }
+                mErrorStream = new ByteArrayInputStream(data.getBytes(StandardCharsets.UTF_8));
             }
-            return new ByteArrayInputStream(data.getBytes(StandardCharsets.UTF_8));
+            return mErrorStream;
         }
         return null;
     }
@@ -121,13 +142,32 @@ public boolean usingProxy() {
 
     @Override
     public void disconnect() {
+        // Close output stream if it exists
         try {
             if (mOutputStream != null) {
                 mOutputStream.close();
             }
         } catch (IOException e) {
             // Ignore exception during disconnect
         }
+        
+        // Close input stream if it exists
+        try {
+            if (mInputStream != null) {
+                mInputStream.close();
+            }
+        } catch (IOException e) {
+            // Ignore exception during disconnect
+        }
+        
+        // Close error stream if it exists
+        try {
+            if (mErrorStream != null) {
+                mErrorStream.close();
+            }
+        } catch (IOException e) {
+            // Ignore exception during disconnect
+        }
     }
 
     // Required abstract method implementations for HTTPS connection

src/test/java/io/split/android/client/network/HttpResponseConnectionAdapterTest.java

@@ -13,6 +13,7 @@
 import org.junit.Test;
 import org.mockito.Mock;
 
+import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
@@ -374,7 +375,7 @@ public void getOutputStreamThrowsWhenNotEnabled() throws IOException {
         // Should throw exception since doOutput is not enabled
         mAdapter.getOutputStream();
     }
-    
+
     @Test
     public void setDoOutputEnablesOutput() {
         mAdapter = new HttpResponseConnectionAdapter(mTestUrl, mMockResponse, mTestCertificates);
@@ -386,15 +387,15 @@ public void setDoOutputEnablesOutput() {
         mAdapter.setDoOutput(true);
         assertEquals(true, mAdapter.getDoOutput());
     }
-    
+
     @Test
     public void getOutputStreamAfterEnablingOutput() throws IOException {
         mAdapter = new HttpResponseConnectionAdapter(mTestUrl, mMockResponse, mTestCertificates);
         mAdapter.setDoOutput(true);
 
         assertNotNull("Output stream should not be null when doOutput is enabled", mAdapter.getOutputStream());
     }
-    
+
     @Test
     public void writeToOutputStream() throws IOException {
         // Create a ByteArrayOutputStream to capture the written data
@@ -411,15 +412,15 @@ public void writeToOutputStream() throws IOException {
         // Verify that the data was written correctly
         assertEquals("Written data should match the input", testData, testOutputStream.toString(StandardCharsets.UTF_8.name()));
     }
-    
+
     @Test
     public void disconnectClosesOutputStream() throws IOException {
         // Create a custom OutputStream that tracks if it's been closed
         TestOutputStream testOutputStream = new TestOutputStream();
-        
+
         mAdapter = new HttpResponseConnectionAdapter(mTestUrl, mMockResponse, mTestCertificates, testOutputStream);
         mAdapter.setDoOutput(true);
-        
+
         // Get the output stream and write some data
         mAdapter.getOutputStream().write("Test".getBytes(StandardCharsets.UTF_8));
 
@@ -433,6 +434,37 @@ public void disconnectClosesOutputStream() throws IOException {
         assertTrue("Output stream should be closed after disconnect", testOutputStream.isClosed());
     }
 
+    @Test
+    public void disconnectClosesInputStream() throws IOException {
+        // Create a custom InputStream that tracks if it's been closed
+        TestInputStream testInputStream = new TestInputStream("Test response data".getBytes(StandardCharsets.UTF_8));
+        TestOutputStream testOutputStream = new TestOutputStream();
+        
+        // Create adapter with injected test input stream
+        when(mMockResponse.getHttpStatus()).thenReturn(200);
+        mAdapter = new HttpResponseConnectionAdapter(
+                mTestUrl, 
+                mMockResponse, 
+                mTestCertificates, 
+                testOutputStream,
+                testInputStream,
+                null);
+        
+        // Get the input stream and read some data to simulate usage
+        InputStream stream = mAdapter.getInputStream();
+        byte[] buffer = new byte[10];
+        stream.read(buffer);
+        
+        // Verify the stream is not closed yet
+        assertFalse("Input stream should not be closed before disconnect", testInputStream.isClosed());
+        
+        // Disconnect should close the input stream
+        mAdapter.disconnect();
+        
+        // Verify the stream was closed
+        assertTrue("Input stream should be closed after disconnect", testInputStream.isClosed());
+    }
+
     /**
      * Custom OutputStream implementation for testing that tracks if it's been closed.
      */
@@ -449,4 +481,48 @@ public boolean isClosed() {
             return mClosed;
         }
     }
+
+    private static class TestInputStream extends ByteArrayInputStream {
+        private boolean mClosed = false;
+        
+        public TestInputStream(byte[] data) {
+            super(data);
+        }
+        
+        @Override
+        public void close() throws IOException {
+            super.close();
+            mClosed = true;
+        }
+        
+        public boolean isClosed() {
+            return mClosed;
+        }
+    }
+    
+    @Test
+    public void disconnectClosesErrorStream() throws IOException {
+        TestInputStream testErrorStream = new TestInputStream("Error data".getBytes(StandardCharsets.UTF_8));
+        TestOutputStream testOutputStream = new TestOutputStream();
+        
+        when(mMockResponse.getHttpStatus()).thenReturn(404); // Error status
+        mAdapter = new HttpResponseConnectionAdapter(
+                mTestUrl, 
+                mMockResponse, 
+                mTestCertificates, 
+                testOutputStream,
+                null,
+                testErrorStream);
+        
+        // Get the error stream and read some data to simulate usage
+        InputStream stream = mAdapter.getErrorStream();
+        byte[] buffer = new byte[10];
+        stream.read(buffer);
+        
+        assertFalse("Error stream should not be closed before disconnect", testErrorStream.isClosed());
+        
+        mAdapter.disconnect();
+        
+        assertTrue("Error stream should be closed after disconnect", testErrorStream.isClosed());
+    }
 }