From 734b58f0f49261ed325857cefd187bf48bd35bd5 Mon Sep 17 00:00:00 2001
From: Marcus Linke <linke@newsaktuell.de>
Date: Tue, 7 Jul 2015 22:58:18 +0200
Subject: [PATCH] Fix FrameReader

---
 .../dockerjava/core/command/FrameReader.java  | 65 ++++++++++++-------
 .../command/AttachContainerCmdImplTest.java   |  4 +-
 2 files changed, 44 insertions(+), 25 deletions(-)

diff --git a/src/main/java/com/github/dockerjava/core/command/FrameReader.java b/src/main/java/com/github/dockerjava/core/command/FrameReader.java
index d8f0cd7f6..a6db922db 100644
--- a/src/main/java/com/github/dockerjava/core/command/FrameReader.java
+++ b/src/main/java/com/github/dockerjava/core/command/FrameReader.java
@@ -14,13 +14,13 @@
  */
 public class FrameReader implements AutoCloseable {
 
-    private static final int BUFFER_SIZE = 100;
-
     private static final int HEADER_SIZE = 8;
 
     private final InputStream inputStream;
 
-    private boolean rawDetected = false;
+    private boolean rawStreamDetected = false;
+
+    private final byte[] rawBuffer = new byte[1000];
 
     public FrameReader(InputStream inputStream) {
         this.inputStream = inputStream;
@@ -35,7 +35,7 @@ private static StreamType streamType(byte streamType) {
         case 2:
             return StreamType.STDERR;
         default:
-            throw new IllegalArgumentException("invalid streamType");
+            return StreamType.RAW;
         }
     }
 
@@ -44,34 +44,55 @@ private static StreamType streamType(byte streamType) {
      */
     public Frame readFrame() throws IOException {
 
-        byte[] buffer = new byte[BUFFER_SIZE];
+        if (rawStreamDetected) {
 
-        int readBytes = inputStream.read(buffer);
 
-        if (readBytes == -1) {
-            return null;
-        }
+            int read = inputStream.read(rawBuffer);
 
-        if (rawDetected || readBytes != HEADER_SIZE) {
-            rawDetected = true;
+            return new Frame(StreamType.RAW, Arrays.copyOf(rawBuffer, read));
 
-            byte[] read = Arrays.copyOfRange(buffer, 0, readBytes);
-
-            return new Frame(StreamType.RAW, read);
         } else {
 
-            int payloadSize = ((buffer[4] & 0xff) << 24) + ((buffer[5] & 0xff) << 16) + ((buffer[6] & 0xff) << 8)
-                    + (buffer[7] & 0xff);
+            byte[] header = new byte[HEADER_SIZE];
 
-            byte[] payload = new byte[payloadSize];
-            int actualPayloadSize = inputStream.read(payload);
-            if (actualPayloadSize != payloadSize) {
-                throw new IOException(String.format("payload must be %d bytes long, but was %d", payloadSize,
-                        actualPayloadSize));
+            int actualHeaderSize = 0;
+
+            do {
+                int headerCount = inputStream.read(header, actualHeaderSize, HEADER_SIZE - actualHeaderSize);
+
+                if (headerCount == -1) {
+                    return null;
+                }
+                actualHeaderSize += headerCount;
+            } while (actualHeaderSize < HEADER_SIZE);
+
+            StreamType streamType = streamType(header[0]);
+
+            if(streamType.equals(StreamType.RAW)) {
+                rawStreamDetected = true;
+                return new Frame(StreamType.RAW, Arrays.copyOf(header, HEADER_SIZE));
             }
 
-            return new Frame(streamType(buffer[0]), payload);
+            int payloadSize = ((header[4] & 0xff) << 24) + ((header[5] & 0xff) << 16) + ((header[6] & 0xff) << 8)
+                    + (header[7] & 0xff);
 
+            byte[] payload = new byte[payloadSize];
+            int actualPayloadSize = 0;
+
+            do {
+                int count = inputStream.read(payload, actualPayloadSize, payloadSize - actualPayloadSize);
+
+                if (count == -1) {
+                    if (actualPayloadSize != payloadSize) {
+                        throw new IOException(String.format("payload must be %d bytes long, but was %d", payloadSize,
+                                actualPayloadSize));
+                    }
+                    break;
+                }
+                actualPayloadSize += count;
+            } while (actualPayloadSize < payloadSize);
+
+            return new Frame(streamType, payload);
         }
     }
 
diff --git a/src/test/java/com/github/dockerjava/core/command/AttachContainerCmdImplTest.java b/src/test/java/com/github/dockerjava/core/command/AttachContainerCmdImplTest.java
index fa5ec44ae..9c5c5c691 100644
--- a/src/test/java/com/github/dockerjava/core/command/AttachContainerCmdImplTest.java
+++ b/src/test/java/com/github/dockerjava/core/command/AttachContainerCmdImplTest.java
@@ -2,10 +2,8 @@
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.containsString;
-import static org.hamcrest.Matchers.endsWith;
 import static org.hamcrest.Matchers.isEmptyString;
 import static org.hamcrest.Matchers.not;
-import static org.hamcrest.Matchers.containsString;
 
 import java.io.File;
 import java.io.InputStream;
@@ -66,7 +64,7 @@ public void attachContainerWithoutTTY() throws Exception {
         CollectFramesCallback collectFramesCallback = new CollectFramesCallback() {
             @Override
             public void onNext(Frame frame) {
-                assertEquals(frame.getStreamType(), StreamType.RAW);
+                assertEquals(frame.getStreamType(), StreamType.STDOUT);
                 super.onNext(frame);
             };
         };