Merge pull request square#517 from square/jwilson_0208_buffer_fixes

Fix some OkBuffer bugs.

okhttp-protocols/src/main/java/com/squareup/okhttp/internal/bytes/GzipSource.java

@@ -135,7 +135,7 @@ private void consumeHeader(Deadline deadline) throws IOException {
     // |...original file name, zero-terminated...| (more-->)
     // +=========================================+
     if (((flags >> FNAME) & 1) == 1) {
-      long index = seek((byte) 0, deadline);
+      long index = OkBuffers.seek(buffer, (byte) 0, source, deadline);
       if (fhcrc) updateCrc(buffer, 0, index + 1);
       buffer.skip(index + 1);
     }
@@ -145,7 +145,7 @@ private void consumeHeader(Deadline deadline) throws IOException {
     // |...file comment, zero-terminated...| (more-->)
     // +===================================+
     if (((flags >> FCOMMENT) & 1) == 1) {
-      long index = seek((byte) 0, deadline);
+      long index = OkBuffers.seek(buffer, (byte) 0, source, deadline);
       if (fhcrc) updateCrc(buffer, 0, index + 1);
       buffer.skip(index + 1);
     }
@@ -194,21 +194,10 @@ private void require(int byteCount, Deadline deadline) throws IOException {
     }
   }
 
-  /** Returns the next index of {@code b}, reading data into the buffer as necessary. */
-  private long seek(byte b, Deadline deadline) throws IOException {
-    long start = 0;
-    long index;
-    while ((index = buffer.indexOf(b, start)) == -1) {
-      start = buffer.byteCount;
-      if (source.read(buffer, Segment.SIZE, deadline) == -1) throw new EOFException();
-    }
-    return index;
-  }
-
   private void checkEqual(String name, int expected, int actual) throws IOException {
     if (actual != expected) {
       throw new IOException(String.format(
-          "%s: actual %#08x != expected %#08x", name, actual, expected));
+          "%s: actual 0x%08x != expected 0x%08x", name, actual, expected));
     }
   }
 }

okhttp-protocols/src/main/java/com/squareup/okhttp/internal/bytes/OkBuffer.java

@@ -21,6 +21,7 @@
 import java.util.Collections;
 import java.util.List;
 
+import static com.squareup.okhttp.internal.Util.UTF_8;
 import static com.squareup.okhttp.internal.Util.checkOffsetAndCount;
 
 /**
@@ -165,7 +166,25 @@ public ByteString readByteString(int byteCount) {
 
   /** Removes {@code byteCount} bytes from this, decodes them as UTF-8 and returns the string. */
   public String readUtf8(int byteCount) {
-    return new String(readBytes(byteCount), Util.UTF_8);
+    checkOffsetAndCount(this.byteCount, 0, byteCount);
+    if (byteCount == 0) return "";
+
+    Segment head = this.head;
+    if (head.pos + byteCount > head.limit) {
+      // If the string spans multiple segments, delegate to readBytes().
+      return new String(readBytes(byteCount), Util.UTF_8);
+    }
+
+    String result = new String(head.data, head.pos, byteCount, UTF_8);
+    head.pos += byteCount;
+    this.byteCount -= byteCount;
+
+    if (head.pos == head.limit) {
+      this.head = head.pop();
+      SegmentPool.INSTANCE.recycle(head);
+    }
+
+    return result;
   }
 
   private byte[] readBytes(int byteCount) {
@@ -354,8 +373,8 @@ Segment writableSegment(int minimumCapacity) {
     while (byteCount > 0) {
       // Is a prefix of the source's head segment all that we need to move?
       if (byteCount < (source.head.limit - source.head.pos)) {
-        Segment tail = head.prev;
-        if (byteCount + (tail.limit - tail.pos) > Segment.SIZE) {
+        Segment tail = head != null ? head.prev : null;
+        if (tail == null || byteCount + (tail.limit - tail.pos) > Segment.SIZE) {
           // We're going to need another segment. Split the source's head
           // segment in two, then move the first of those two to this buffer.
           source.head = source.head.split((int) byteCount);

okhttp-protocols/src/main/java/com/squareup/okhttp/internal/bytes/OkBuffers.java

@@ -15,6 +15,7 @@
  */
 package com.squareup.okhttp.internal.bytes;
 
+import java.io.EOFException;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
@@ -25,6 +26,22 @@ public final class OkBuffers {
   private OkBuffers() {
   }
 
+  /**
+   * Returns the index of {@code b} in {@code buffer}, refilling it if necessary
+   * until it is found. This reads an unbounded number of bytes into {@code
+   * buffer}.
+   */
+  public static long seek(OkBuffer buffer, byte b, Source source, Deadline deadline)
+      throws IOException {
+    long start = 0;
+    long index;
+    while ((index = buffer.indexOf(b, start)) == -1) {
+      start = buffer.byteCount;
+      if (source.read(buffer, Segment.SIZE, deadline) == -1) throw new EOFException();
+    }
+    return index;
+  }
+
   /** Returns a sink that writes to {@code out}. */
   public static Sink sink(final OutputStream out) {
     return new Sink() {
@@ -148,7 +165,7 @@ public static InputStream inputStream(final Source source) {
           long count = source.read(buffer, Segment.SIZE, Deadline.NONE);
           if (count == -1) return -1;
         }
-        return buffer.readByte();
+        return buffer.readByte() & 0xff;
       }
 
       @Override public int read(byte[] data, int offset, int byteCount) throws IOException {
@@ -179,7 +196,7 @@ public static InputStream inputStream(final Source source) {
       }
 
       @Override public void close() throws IOException {
-        super.close();
+        source.close(Deadline.NONE);
       }
 
       @Override public String toString() {

okhttp-protocols/src/test/java/com/squareup/okhttp/internal/bytes/GzipSourceTest.java

@@ -115,7 +115,7 @@ private void assertGzipped(OkBuffer gzipped) throws IOException {
       gunzip(gzipped);
       fail();
     } catch (IOException e) {
-      assertEquals("FHCRC: actual 0x00261d != expected 0x000000", e.getMessage());
+      assertEquals("FHCRC: actual 0x0000261d != expected 0x00000000", e.getMessage());
     }
   }
 
@@ -130,7 +130,7 @@ private void assertGzipped(OkBuffer gzipped) throws IOException {
       gunzip(gzipped);
       fail();
     } catch (IOException e) {
-      assertEquals("CRC: actual 0x37ad8f8d != expected 0x1234567", e.getMessage());
+      assertEquals("CRC: actual 0x37ad8f8d != expected 0x01234567", e.getMessage());
     }
   }
 
@@ -145,7 +145,7 @@ private void assertGzipped(OkBuffer gzipped) throws IOException {
       gunzip(gzipped);
       fail();
     } catch (IOException e) {
-      assertEquals("ISIZE: actual 0x000020 != expected 0x123456", e.getMessage());
+      assertEquals("ISIZE: actual 0x00000020 != expected 0x00123456", e.getMessage());
     }
   }
 

okhttp-protocols/src/test/java/com/squareup/okhttp/internal/bytes/OkBufferTest.java

@@ -47,6 +47,19 @@ public final class OkBufferTest {
     }
   }
 
+  @Test public void readUtf8SpansSegments() throws Exception {
+    OkBuffer buffer = new OkBuffer();
+    buffer.writeUtf8(repeat('a', Segment.SIZE * 2));
+    buffer.readUtf8(Segment.SIZE - 1);
+    assertEquals("aa", buffer.readUtf8(2));
+  }
+
+  @Test public void readUtf8EntireBuffer() throws Exception {
+    OkBuffer buffer = new OkBuffer();
+    buffer.writeUtf8(repeat('a', Segment.SIZE));
+    assertEquals(repeat('a', Segment.SIZE), buffer.readUtf8(Segment.SIZE));
+  }
+
   @Test public void bufferToString() throws Exception {
     OkBuffer buffer = new OkBuffer();
     buffer.writeUtf8("\u0000\u0001\u0002\u007f");
@@ -465,8 +478,8 @@ private List<Integer> moveBytesBetweenBuffers(String... contents) {
   @Test public void readByte() throws Exception {
     OkBuffer data = new OkBuffer();
     data.write(new ByteString(new byte[] { (byte) 0xab, (byte) 0xcd }));
-    assertEquals((byte) 0xab, data.readByte());
-    assertEquals((byte) 0xcd, data.readByte());
+    assertEquals(0xab, data.readByte() & 0xff);
+    assertEquals(0xcd, data.readByte() & 0xff);
     assertEquals(0, data.byteCount());
   }
 
@@ -538,13 +551,21 @@ private List<Integer> moveBytesBetweenBuffers(String... contents) {
     buffer.writeUtf8(repeat('b', Segment.SIZE));
     buffer.writeUtf8("c");
     buffer.skip(1);
-    assertEquals('b', buffer.readByte());
+    assertEquals('b', buffer.readByte() & 0xff);
     buffer.skip(Segment.SIZE - 2);
-    assertEquals('b', buffer.readByte());
+    assertEquals('b', buffer.readByte() & 0xff);
     buffer.skip(1);
     assertEquals(0, buffer.byteCount());
   }
 
+  @Test public void testWritePrefixToEmptyBuffer() {
+    OkBuffer sink = new OkBuffer();
+    OkBuffer source = new OkBuffer();
+    source.writeUtf8("abcd");
+    sink.write(source, 2, Deadline.NONE);
+    assertEquals("ab", sink.readUtf8(2));
+  }
+
   private String repeat(char c, int count) {
     char[] array = new char[count];
     Arrays.fill(array, c);