Repository: square/okio
Issue: DeflaterSink, InflaterSource may leave behind an empty tail

Both of these allocate a segment if necessary, but may end up not writing to that segment. We should write tests that exercise this situation and make sure we don't have a 0-byte segment still in a Buffer.


Changed files:
okio/src/main/java/okio/InflaterSource.java, okio/src/test/java/okio/InflaterSourceTest.java, okio/src/main/java/okio/DeflaterSink.java, okio/src/main/java/okio/Okio.java, okio/src/test/java/okio/DeflaterSinkTest.java

diff --git a/okio/src/main/java/okio/DeflaterSink.java b/okio/src/main/java/okio/DeflaterSink.java
index 3b23655c..a5bca15d 100644
--- a/okio/src/main/java/okio/DeflaterSink.java
+++ b/okio/src/main/java/okio/DeflaterSink.java
@@ -99,7 +99,12 @@ public final class DeflaterSink implements Sink {
         buffer.size += deflated;
         sink.emitCompleteSegments();
       } else if (deflater.needsInput()) {
-        return; // TODO(jwilson): do we have a dangling empty tail segment here?
+        if (s.pos == s.limit) {
+          // We allocated a tail segment, but didn't end up needing it. Recycle!
+          buffer.head = s.pop();
+          SegmentPool.recycle(s);
+        }
+        return;
       }
     }
   }
diff --git a/okio/src/main/java/okio/InflaterSource.java b/okio/src/main/java/okio/InflaterSource.java
index c38819dc..4011ca02 100644
--- a/okio/src/main/java/okio/InflaterSource.java
+++ b/okio/src/main/java/okio/InflaterSource.java
@@ -70,9 +70,13 @@ public final class InflaterSource implements Source {
           sink.size += bytesInflated;
           return bytesInflated;
         }
-        // TODO(jwilson): do we have an empty tail?
         if (inflater.finished() || inflater.needsDictionary()) {
           releaseInflatedBytes();
+          if (tail.pos == tail.limit) {
+            // We allocated a tail segment, but didn't end up needing it. Recycle!
+            sink.head = tail.pop();
+            SegmentPool.recycle(tail);
+          }
           return -1;
         }
         if (sourceExhausted) throw new EOFException("source exhausted prematurely");
diff --git a/okio/src/main/java/okio/Okio.java b/okio/src/main/java/okio/Okio.java
index 8df1fde8..2b68d48b 100644
--- a/okio/src/main/java/okio/Okio.java
+++ b/okio/src/main/java/okio/Okio.java
@@ -130,6 +130,7 @@ public final class Okio {
     return new Source() {
       @Override public long read(Buffer sink, long byteCount) throws IOException {
         if (byteCount < 0) throw new IllegalArgumentException("byteCount < 0: " + byteCount);
+        if (byteCount == 0) return 0;
         timeout.throwIfReached();
         Segment tail = sink.writableSegment(1);
         int maxToCopy = (int) Math.min(byteCount, Segment.SIZE - tail.limit);
diff --git a/okio/src/test/java/okio/DeflaterSinkTest.java b/okio/src/test/java/okio/DeflaterSinkTest.java
index 87d6c598..b6cafb21 100644
--- a/okio/src/test/java/okio/DeflaterSinkTest.java
+++ b/okio/src/test/java/okio/DeflaterSinkTest.java
@@ -87,6 +87,24 @@ public final class DeflaterSinkTest {
     assertEquals(repeat('a', byteCount), inflate(buffer).readUtf8(byteCount));
   }
 
+  @Test public void deflateIntoNonemptySink() throws Exception {
+    String original = "They're moving in herds. They do move in herds.";
+
+    // Exercise all possible offsets for the outgoing segment.
+    for (int i = 0; i < Segment.SIZE; i++) {
+      Buffer data = new Buffer().writeUtf8(original);
+      Buffer sink = new Buffer().writeUtf8(repeat('a', i));
+
+      DeflaterSink deflaterSink = new DeflaterSink(sink, new Deflater());
+      deflaterSink.write(data, data.size());
+      deflaterSink.close();
+
+      sink.skip(i);
+      Buffer inflated = inflate(sink);
+      assertEquals(original, inflated.readUtf8());
+    }
+  }
+
   /**
    * This test deflates a single segment of without compression because that's
    * the easiest way to force close() to emit a large amount of data to the
@@ -121,7 +139,9 @@ public final class DeflaterSinkTest {
     byte[] buffer = new byte[8192];
     while (!inflater.needsInput() || deflated.size() > 0 || deflatedIn.available() > 0) {
       int count = inflatedIn.read(buffer, 0, buffer.length);
-      result.write(buffer, 0, count);
+      if (count != -1) {
+        result.write(buffer, 0, count);
+      }
     }
     return result;
   }
diff --git a/okio/src/test/java/okio/InflaterSourceTest.java b/okio/src/test/java/okio/InflaterSourceTest.java
index e32f18f2..d09d6d33 100644
--- a/okio/src/test/java/okio/InflaterSourceTest.java
+++ b/okio/src/test/java/okio/InflaterSourceTest.java
@@ -17,6 +17,7 @@ package okio;
 
 import java.io.EOFException;
 import java.io.IOException;
+import java.util.zip.Deflater;
 import java.util.zip.DeflaterOutputStream;
 import java.util.zip.Inflater;
 import org.junit.Test;
@@ -74,6 +75,19 @@ public final class InflaterSourceTest {
     assertEquals(original, inflated.readByteString());
   }
 
+  @Test public void inflateIntoNonemptySink() throws Exception {
+    for (int i = 0; i < Segment.SIZE; i++) {
+      Buffer inflated = new Buffer().writeUtf8(repeat('a', i));
+      Buffer deflated = decodeBase64(
+          "eJxzz09RyEjNKVAoLdZRKE9VL0pVyMxTKMlIVchIzEspVshPU0jNS8/MS00tKtYDAF6CD5s=");
+      InflaterSource source = new InflaterSource(deflated, new Inflater());
+      while (source.read(inflated, Integer.MAX_VALUE) != -1) {
+      }
+      inflated.skip(i);
+      assertEquals("God help us, we're in the hands of engineers.", inflated.readUtf8());
+    }
+  }
+
   private Buffer decodeBase64(String s) {
     return new Buffer().write(ByteString.decodeBase64(s));
   }