OutputStream capabilities in connection adapter

Author: gthea

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

@@ -1,8 +1,10 @@
 package io.split.android.client.network;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.VisibleForTesting;
 
 import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
@@ -22,13 +24,15 @@
 /**
  * Adapter that wraps an HttpResponse as an HttpURLConnection.
  * <p>
- * This is only used to adapt the response from the CONNECT method.
+ * This is only used to adapt the response from request through the TLS tunnel.
  */
 class HttpResponseConnectionAdapter extends HttpsURLConnection {
 
     private final HttpResponse mResponse;
     private final URL mUrl;
     private final Certificate[] mServerCertificates;
+    private OutputStream mOutputStream;
+    private boolean mDoOutput = false;
 
     /**
      * Creates an adapter that wraps an HttpResponse as an HttpURLConnection.
@@ -38,12 +42,21 @@ class HttpResponseConnectionAdapter extends HttpsURLConnection {
      * @param serverCertificates The server certificates from the SSL connection
      */
     HttpResponseConnectionAdapter(@NonNull URL url,
-                                         @NonNull HttpResponse response,
-                                         Certificate[] serverCertificates) {
+                                  @NonNull HttpResponse response,
+                                  Certificate[] serverCertificates) {
+        this(url, response, serverCertificates, new ByteArrayOutputStream());
+    }
+
+    @VisibleForTesting
+    HttpResponseConnectionAdapter(@NonNull URL url,
+                                  @NonNull HttpResponse response,
+                                  Certificate[] serverCertificates,
+                                  @NonNull OutputStream outputStream) {
         super(url);
         mUrl = url;
         mResponse = response;
         mServerCertificates = serverCertificates;
+        mOutputStream = outputStream;
     }
 
     @Override
@@ -108,6 +121,13 @@ public boolean usingProxy() {
 
     @Override
     public void disconnect() {
+        try {
+            if (mOutputStream != null) {
+                mOutputStream.close();
+            }
+        } catch (IOException e) {
+            // Ignore exception during disconnect
+        }
     }
 
     // Required abstract method implementations for HTTPS connection
@@ -148,11 +168,12 @@ public boolean getInstanceFollowRedirects() {
 
     @Override
     public void setDoOutput(boolean doOutput) {
+        mDoOutput = doOutput;
     }
 
     @Override
     public boolean getDoOutput() {
-        return false;
+        return mDoOutput;
     }
 
     @Override
@@ -350,7 +371,10 @@ public Permission getPermission() throws IOException {
 
     @Override
     public OutputStream getOutputStream() throws IOException {
-        throw new IOException("Output not supported for SSL proxy responses");
+        if (!mDoOutput) {
+            throw new IOException("Output not enabled for this connection. Call setDoOutput(true) first.");
+        }
+        return mOutputStream;
     }
 
     @Override

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

@@ -1,6 +1,7 @@
 package io.split.android.client.network;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
@@ -12,6 +13,7 @@
 import org.junit.Test;
 import org.mockito.Mock;
 
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.MalformedURLException;
@@ -367,8 +369,84 @@ public void urlCanBeRetrieved() {
     }
 
     @Test(expected = IOException.class)
-    public void getOutputStreamThrows() throws IOException {
+    public void getOutputStreamThrowsWhenNotEnabled() throws IOException {
         mAdapter = new HttpResponseConnectionAdapter(mTestUrl, mMockResponse, mTestCertificates);
+        // Should throw exception since doOutput is not enabled
         mAdapter.getOutputStream();
     }
+    
+    @Test
+    public void setDoOutputEnablesOutput() {
+        mAdapter = new HttpResponseConnectionAdapter(mTestUrl, mMockResponse, mTestCertificates);
+        
+        // Initially doOutput should be false
+        assertEquals(false, mAdapter.getDoOutput());
+        
+        // After setting doOutput to true, getDoOutput should return true
+        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
+        ByteArrayOutputStream testOutputStream = new ByteArrayOutputStream();
+        
+        // Use the constructor that accepts a custom OutputStream
+        mAdapter = new HttpResponseConnectionAdapter(mTestUrl, mMockResponse, mTestCertificates, testOutputStream);
+        mAdapter.setDoOutput(true);
+        
+        // Write test data to the output stream
+        String testData = "Test output data";
+        mAdapter.getOutputStream().write(testData.getBytes(StandardCharsets.UTF_8));
+        
+        // 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));
+        
+        // Verify the stream is not closed yet
+        assertFalse("Output stream should not be closed before disconnect", testOutputStream.isClosed());
+        
+        // Disconnect should close the output stream
+        mAdapter.disconnect();
+        
+        // Verify the stream was closed
+        assertTrue("Output stream should be closed after disconnect", testOutputStream.isClosed());
+    }
+    
+    /**
+     * Custom OutputStream implementation for testing that tracks if it's been closed.
+     */
+    private static class TestOutputStream extends ByteArrayOutputStream {
+        private boolean mClosed = false;
+        
+        @Override
+        public void close() throws IOException {
+            super.close();
+            mClosed = true;
+        }
+        
+        public boolean isClosed() {
+            return mClosed;
+        }
+    }
 }