Repository: square/okio
Issue: BufferedSink.write(Source, long) hides bugs

It calls `source.read(buffer, byteCount)` which might not read `byteCount` bytes and also might return -1 (which is ignored).

Our API should make it hard to write buggy code, and this one tripped me up forcing use of the inverse API which correctly reported read count.


Changed files:
okio/src/main/java/okio/RealBufferedSink.java, okio/src/main/java/okio/Buffer.java, okio/src/test/java/okio/BufferedSinkTest.java

diff --git a/okio/src/main/java/okio/Buffer.java b/okio/src/main/java/okio/Buffer.java
index 92097b9e..44cfbf29 100644
--- a/okio/src/main/java/okio/Buffer.java
+++ b/okio/src/main/java/okio/Buffer.java
@@ -826,8 +826,10 @@ public final class Buffer implements BufferedSource, BufferedSink, Cloneable {
   }
 
   @Override public BufferedSink write(Source source, long byteCount) throws IOException {
-    if (byteCount > 0) {
-      source.read(this, byteCount);
+    while (byteCount > 0) {
+      long read = source.read(this, byteCount);
+      if (read == -1) throw new EOFException();
+      byteCount -= read;
     }
     return this;
   }
diff --git a/okio/src/main/java/okio/RealBufferedSink.java b/okio/src/main/java/okio/RealBufferedSink.java
index d5d4799d..1902fd80 100644
--- a/okio/src/main/java/okio/RealBufferedSink.java
+++ b/okio/src/main/java/okio/RealBufferedSink.java
@@ -15,6 +15,7 @@
  */
 package okio;
 
+import java.io.EOFException;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.charset.Charset;
@@ -86,8 +87,11 @@ final class RealBufferedSink implements BufferedSink {
   }
 
   @Override public BufferedSink write(Source source, long byteCount) throws IOException {
-    if (byteCount > 0) {
-      source.read(buffer, byteCount);
+    while (byteCount > 0) {
+      long read = source.read(buffer, byteCount);
+      if (read == -1) throw new EOFException();
+      byteCount -= read;
+      emitCompleteSegments();
     }
     return this;
   }
diff --git a/okio/src/test/java/okio/BufferedSinkTest.java b/okio/src/test/java/okio/BufferedSinkTest.java
index 879ecbdb..73a5cfbc 100644
--- a/okio/src/test/java/okio/BufferedSinkTest.java
+++ b/okio/src/test/java/okio/BufferedSinkTest.java
@@ -1,5 +1,6 @@
 package okio;
 
+import java.io.EOFException;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.math.BigInteger;
@@ -170,6 +171,33 @@ public class BufferedSinkTest {
     assertEquals("ef", source.readUtf8());
   }
 
+  @Test public void writeSourceReadsFully() throws Exception {
+    Source source = new ForwardingSource(new Buffer()) {
+      @Override public long read(Buffer sink, long byteCount) throws IOException {
+        sink.writeUtf8("abcd");
+        return 4;
+      }
+    };
+
+    sink.write(source, 8);
+    sink.flush();
+    assertEquals("abcdabcd", data.readUtf8());
+  }
+
+  @Test public void writeSourcePropagatesEof() throws IOException {
+    Source source = new Buffer().writeUtf8("abcd");
+
+    try {
+      sink.write(source, 8);
+      fail();
+    } catch (EOFException expected) {
+    }
+
+    // Ensure that whatever was available was correctly written.
+    sink.flush();
+    assertEquals("abcd", data.readUtf8());
+  }
+
   @Test public void writeSourceWithZeroIsNoOp() throws IOException {
     // This test ensures that a zero byte count never calls through to read the source. It may be
     // tied to something like a socket which will potentially block trying to read a segment when