Merge "Updated the AdbUtils class to better manage opening and closing the resources."
diff --git a/tests/DumpRenderTree2/src/com/android/dumprendertree2/FsUtils.java b/tests/DumpRenderTree2/src/com/android/dumprendertree2/FsUtils.java
index e2e07ad..4438811 100644
--- a/tests/DumpRenderTree2/src/com/android/dumprendertree2/FsUtils.java
+++ b/tests/DumpRenderTree2/src/com/android/dumprendertree2/FsUtils.java
@@ -228,4 +228,24 @@
 
         return new LinkedList<String>();
     }
+
+    public static void closeInputStream(InputStream inputStream) {
+        try {
+            if (inputStream != null) {
+                inputStream.close();
+            }
+        } catch (IOException e) {
+            Log.e(LOG_TAG, "Couldn't close stream!", e);
+        }
+    }
+
+    public static void closeOutputStream(OutputStream outputStream) {
+        try {
+            if (outputStream != null) {
+                outputStream.close();
+            }
+        } catch (IOException e) {
+            Log.e(LOG_TAG, "Couldn't close stream!", e);
+        }
+    }
 }
\ No newline at end of file
diff --git a/tests/DumpRenderTree2/src/com/android/dumprendertree2/forwarder/AdbUtils.java b/tests/DumpRenderTree2/src/com/android/dumprendertree2/forwarder/AdbUtils.java
index d165a1a..086ff59 100644
--- a/tests/DumpRenderTree2/src/com/android/dumprendertree2/forwarder/AdbUtils.java
+++ b/tests/DumpRenderTree2/src/com/android/dumprendertree2/forwarder/AdbUtils.java
@@ -36,52 +36,45 @@
     private static final int ADB_RESPONSE_SIZE = 4;
 
     /**
-     * Send an ADB command using existing socket connection
+     * Creates a new socket that can be configured to serve as a transparent proxy to a
+     * remote machine. This can be achieved by calling configureSocket()
      *
-     * The streams provided must be from a socket connected to adb already
+     * @return a socket that can be configured to link to remote machine
+     */
+    public static Socket createSocket() {
+        Socket socket = null;
+        try {
+            socket = new Socket(ADB_HOST, ADB_PORT);
+        } catch (IOException e) {
+            Log.e(LOG_TAG, "Creation failed.", e);
+        }
+        return socket;
+    }
+
+    /**
+     * Configures the connection to serve as a transparent proxy to a remote machine.
+     * The given streams must belong to a socket created by createSocket().
      *
-     * @param is input stream of the socket connection
-     * @param os output stream of the socket
-     * @param cmd the adb command to send
-     * @return if adb gave a success response
+     * @param inputStream inputStream of the socket we want to configure
+     * @param outputStream outputStream of the socket we want to configure
+     * @param remoteAddress address of the remote machine (as you would type in a browser
+     *      in a machine that the device is connected to via adb)
+     * @param remotePort port on which to connect
+     * @return if the configuration suceeded
      * @throws IOException
      */
-    private static boolean sendAdbCmd(InputStream is, OutputStream os, String cmd)
-            throws IOException {
-        byte[] buf = new byte[ADB_RESPONSE_SIZE];
-
+    public static boolean configureConnection(InputStream inputStream, OutputStream outputStream,
+            String remoteAddress, int remotePort) throws IOException {
+        String cmd = "tcp:" + remotePort + ":" + remoteAddress;
         cmd = String.format("%04X", cmd.length()) + cmd;
-        os.write(cmd.getBytes());
-        int read = is.read(buf);
+
+        byte[] buf = new byte[ADB_RESPONSE_SIZE];
+        outputStream.write(cmd.getBytes());
+        int read = inputStream.read(buf);
         if (read != ADB_RESPONSE_SIZE || !ADB_OK.equals(new String(buf))) {
             Log.w(LOG_TAG, "adb cmd faild.");
             return false;
         }
         return true;
     }
-
-    /**
-     * Get a tcp socket connection to specified IP address and port proxied by adb
-     *
-     * The proxying is transparent, e.g. if a socket is returned, then it can be written to and
-     * read from as if it is directly connected to the target
-     *
-     * @param remoteAddress IP address of the host to connect to
-     * @param remotePort port of the host to connect to
-     * @return a valid Socket instance if successful, null otherwise
-     */
-    public static Socket getSocketToRemoteMachine(String remoteAddress, int remotePort) {
-        try {
-            Socket socket = new Socket(ADB_HOST, ADB_PORT);
-            String cmd = "tcp:" + remotePort + ":" + remoteAddress;
-            if (!sendAdbCmd(socket.getInputStream(), socket.getOutputStream(), cmd)) {
-                socket.close();
-                return null;
-            }
-            return socket;
-        } catch (IOException ioe) {
-            Log.w(LOG_TAG, "error creating adb socket", ioe);
-            return null;
-        }
-    }
 }
\ No newline at end of file
diff --git a/tests/DumpRenderTree2/src/com/android/dumprendertree2/forwarder/ConnectionHandler.java b/tests/DumpRenderTree2/src/com/android/dumprendertree2/forwarder/ConnectionHandler.java
index 5e9f24e..4f01dae 100644
--- a/tests/DumpRenderTree2/src/com/android/dumprendertree2/forwarder/ConnectionHandler.java
+++ b/tests/DumpRenderTree2/src/com/android/dumprendertree2/forwarder/ConnectionHandler.java
@@ -18,6 +18,8 @@
 
 import android.util.Log;
 
+import com.android.dumprendertree2.FsUtils;
+
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
@@ -38,36 +40,25 @@
 
     private class SocketPipeThread extends Thread {
 
-        private Socket mInSocket, mOutSocket;
+        private InputStream mInputStream;
+        private OutputStream mOutputStream;
 
-        public SocketPipeThread(Socket inSocket, Socket outSocket) {
-            mInSocket = inSocket;
-            mOutSocket = outSocket;
+        public SocketPipeThread(InputStream inputStream, OutputStream outputStream) {
+            mInputStream = inputStream;
+            mOutputStream = outputStream;
+            setName("SocketPipeThread: " + getName());
         }
 
         @Override
         public void run() {
-            InputStream is;
-            OutputStream os;
-            try {
-                synchronized (this) {
-                    is = mInSocket.getInputStream();
-                    os = mOutSocket.getOutputStream();
-                }
-            } catch (IOException e) {
-                Log.w(LOG_TAG, this.toString(), e);
-                finish();
-                return;
-            }
-
             byte[] buffer = new byte[4096];
             int length;
             while (true) {
                 try {
-                    if ((length = is.read(buffer)) <= 0) {
+                    if ((length = mInputStream.read(buffer)) < 0) {
                         break;
                     }
-                    os.write(buffer, 0, length);
+                    mOutputStream.write(buffer, 0, length);
                 } catch (IOException e) {
                     /** This exception means one of the streams is closed */
                     Log.v(LOG_TAG, this.toString(), e);
@@ -75,10 +66,6 @@
                 }
             }
 
-            finish();
-        }
-
-        private void finish() {
             synchronized (mThreadsRunning) {
                 mThreadsRunning--;
                 if (mThreadsRunning == 0) {
@@ -90,7 +77,7 @@
 
         @Override
         public String toString() {
-            return "SocketPipeThread:\n" + mInSocket + "\n=>\n" + mOutSocket;
+            return getName();
         }
     }
 
@@ -98,20 +85,51 @@
 
     private Socket mFromSocket, mToSocket;
     private SocketPipeThread mFromToPipe, mToFromPipe;
+    private InputStream mFromSocketInputStream, mToSocketInputStream;
+    private OutputStream mFromSocketOutputStream, mToSocketOutputStream;
+
+    private int mPort;
+    private String mRemoteMachineIpAddress;
 
     private OnFinishedCallback mOnFinishedCallback;
 
-    public ConnectionHandler(Socket fromSocket, Socket toSocket) {
+    public ConnectionHandler(String remoteMachineIp, int port, Socket fromSocket, Socket toSocket) {
+        mRemoteMachineIpAddress = remoteMachineIp;
+        mPort = port;
+
         mFromSocket = fromSocket;
         mToSocket = toSocket;
-        mFromToPipe = new SocketPipeThread(mFromSocket, mToSocket);
-        mToFromPipe = new SocketPipeThread(mToSocket, mFromSocket);
+
+        try {
+            mFromSocketInputStream = mFromSocket.getInputStream();
+            mToSocketInputStream = mToSocket.getInputStream();
+            mFromSocketOutputStream = mFromSocket.getOutputStream();
+            mToSocketOutputStream = mToSocket.getOutputStream();
+            if (!AdbUtils.configureConnection(mToSocketInputStream, mToSocketOutputStream,
+                    mRemoteMachineIpAddress, mPort)) {
+                throw new IOException("Configuring socket failed!");
+            }
+        } catch (IOException e) {
+            Log.e(LOG_TAG, "Unable to start ConnectionHandler", e);
+            closeStreams();
+            return;
+        }
+
+        mFromToPipe = new SocketPipeThread(mFromSocketInputStream, mToSocketOutputStream);
+        mToFromPipe = new SocketPipeThread(mToSocketInputStream, mFromSocketOutputStream);
     }
 
     public void registerOnConnectionHandlerFinishedCallback(OnFinishedCallback callback) {
         mOnFinishedCallback = callback;
     }
 
+    private void closeStreams() {
+        FsUtils.closeInputStream(mFromSocketInputStream);
+        FsUtils.closeInputStream(mToSocketInputStream);
+        FsUtils.closeOutputStream(mFromSocketOutputStream);
+        FsUtils.closeOutputStream(mToSocketOutputStream);
+    }
+
     public void start() {
         /** We have 2 threads running, one for each pipe, that we start here. */
         mThreadsRunning = 2;
diff --git a/tests/DumpRenderTree2/src/com/android/dumprendertree2/forwarder/Forwarder.java b/tests/DumpRenderTree2/src/com/android/dumprendertree2/forwarder/Forwarder.java
index 31cd8ea..b361a89 100644
--- a/tests/DumpRenderTree2/src/com/android/dumprendertree2/forwarder/Forwarder.java
+++ b/tests/DumpRenderTree2/src/com/android/dumprendertree2/forwarder/Forwarder.java
@@ -60,19 +60,18 @@
     @Override
     public void run() {
         while (true) {
-            /** These sockets will be closed when Forwarder.stop() is called */
             Socket localSocket;
             Socket remoteSocket;
             try {
                 localSocket = mServerSocket.accept();
-                remoteSocket = AdbUtils.getSocketToRemoteMachine(mRemoteMachineIpAddress,
-                        mPort);
             } catch (IOException e) {
                 /** This most likely means that mServerSocket is already closed */
                 Log.w(LOG_TAG, "mPort=" + mPort, e);
                 break;
             }
 
+            remoteSocket = AdbUtils.createSocket();
+
             if (remoteSocket == null) {
                 try {
                     localSocket.close();
@@ -86,7 +85,8 @@
             }
 
             final ConnectionHandler connectionHandler =
-                    new ConnectionHandler(localSocket, remoteSocket);
+                    new ConnectionHandler(mRemoteMachineIpAddress, mPort, localSocket,
+                            remoteSocket);
 
             /**
              * We have to close the sockets after the ConnectionHandler finishes, so we
@@ -98,9 +98,7 @@
                 @Override
                 public void onFinished() {
                     synchronized (this) {
-                        if (mConnectionHandlers.remove(connectionHandler)) {
-                            Log.d(LOG_TAG, "removeConnectionHandler(): removed");
-                        } else {
+                        if (!mConnectionHandlers.remove(connectionHandler)) {
                             assert false : "removeConnectionHandler(): not in the collection";
                         }
                     }