From 855c114f4f6fce3e0e1ae71479f89c9d5ad066d6 Mon Sep 17 00:00:00 2001 From: jwilson Date: Wed, 29 Jan 2014 00:08:27 -0500 Subject: [PATCH] Read/write access for byte, short and int. This is big-endian only because that's all we use in OkHttp. --- .../okhttp/internal/bytes/OkBuffer.java | 151 ++++++++++++++++-- .../okhttp/internal/bytes/OkBuffers.java | 2 +- .../okhttp/internal/bytes/OkBufferTest.java | 100 ++++++++++++ 3 files changed, 239 insertions(+), 14 deletions(-) diff --git a/okhttp-protocols/src/main/java/com/squareup/okhttp/internal/bytes/OkBuffer.java b/okhttp-protocols/src/main/java/com/squareup/okhttp/internal/bytes/OkBuffer.java index 8a646337b..fc6c25f7d 100644 --- a/okhttp-protocols/src/main/java/com/squareup/okhttp/internal/bytes/OkBuffer.java +++ b/okhttp-protocols/src/main/java/com/squareup/okhttp/internal/bytes/OkBuffer.java @@ -51,6 +51,91 @@ public final class OkBuffer implements Source, Sink { return byteCount; } + /** Reads a byte from the front of this buffer and returns it. */ + public byte readByte() { + if (byteCount < 1) throw new IllegalArgumentException("byteCount < 1: " + byteCount); + + Segment segment = head; + int pos = segment.pos; + int limit = segment.limit; + + byte[] data = segment.data; + byte b = data[pos++]; + byteCount -= 1; + + if (pos == limit) { + head = segment.pop(); + SegmentPool.INSTANCE.recycle(segment); + } else { + segment.pos = pos; + } + + return b; + } + + /** Reads a Big-Endian short from the front of this buffer and returns it. */ + public short readShort() { + if (byteCount < 2) throw new IllegalArgumentException("byteCount < 2: " + byteCount); + + Segment segment = head; + int pos = segment.pos; + int limit = segment.limit; + + // If the short is split across multiple segments, delegate to readByte(). + if (limit - pos < 2) { + int s = (readByte() & 0xff) << 8 + | (readByte() & 0xff); + return (short) s; + } + + byte[] data = segment.data; + int s = (data[pos++] & 0xff) << 8 + | (data[pos++] & 0xff); + byteCount -= 2; + + if (pos == limit) { + head = segment.pop(); + SegmentPool.INSTANCE.recycle(segment); + } else { + segment.pos = pos; + } + + return (short) s; + } + + /** Reads a Big-Endian int from the front of this buffer and returns it. */ + public int readInt() { + if (byteCount < 4) throw new IllegalArgumentException("byteCount < 4: " + byteCount); + + Segment segment = head; + int pos = segment.pos; + int limit = segment.limit; + + // If the int is split across multiple segments, delegate to readByte(). + if (limit - pos < 4) { + return (readByte() & 0xff) << 24 + | (readByte() & 0xff) << 16 + | (readByte() & 0xff) << 8 + | (readByte() & 0xff); + } + + byte[] data = segment.data; + int i = (data[pos++] & 0xff) << 24 + | (data[pos++] & 0xff) << 16 + | (data[pos++] & 0xff) << 8 + | (data[pos++] & 0xff); + byteCount -= 4; + + if (pos == limit) { + head = segment.pop(); + SegmentPool.INSTANCE.recycle(segment); + } else { + segment.pos = pos; + } + + return i; + } + /** Removes {@code byteCount} bytes from this and returns them as a byte string. */ public ByteString readByteString(int byteCount) { return new ByteString(readBytes(byteCount)); @@ -87,20 +172,21 @@ public final class OkBuffer implements Source, Sink { /** Appends {@code byteString} to this. */ public void write(ByteString byteString) { - write(byteString.data); + write(byteString.data, 0, byteString.data.length); } /** Encodes {@code string} as UTF-8 and appends the bytes to this. */ public void writeUtf8(String string) { - write(string.getBytes(Util.UTF_8)); + byte[] data = string.getBytes(Util.UTF_8); + write(data, 0, data.length); } - private void write(byte[] data) { - int offset = 0; - while (offset < data.length) { - Segment tail = writableSegment(); + void write(byte[] data, int offset, int byteCount) { + int limit = offset + byteCount; + while (offset < limit) { + Segment tail = writableSegment(1); - int toCopy = Math.min(data.length - offset, Segment.SIZE - tail.limit); + int toCopy = Math.min(limit - offset, Segment.SIZE - tail.limit); System.arraycopy(data, offset, tail.data, tail.limit, toCopy); offset += toCopy; @@ -110,15 +196,54 @@ public final class OkBuffer implements Source, Sink { this.byteCount += data.length; } - /** Returns a tail segment that we can write bytes to, creating it if necessary. */ - Segment writableSegment() { + /** Appends a Big-Endian byte to the end of this buffer. */ + public OkBuffer writeByte(int b) { + Segment tail = writableSegment(1); + tail.data[tail.limit++] = (byte) b; + byteCount += 1; + return this; + } + + /** Appends a Big-Endian short to the end of this buffer. */ + public OkBuffer writeShort(int s) { + Segment tail = writableSegment(2); + byte[] data = tail.data; + int limit = tail.limit; + data[limit++] = (byte) ((s >> 8) & 0xff); + data[limit++] = (byte) (s & 0xff); + tail.limit = limit; + byteCount += 2; + return this; + } + + /** Appends a Big-Endian int to the end of this buffer. */ + public OkBuffer writeInt(int i) { + Segment tail = writableSegment(4); + byte[] data = tail.data; + int limit = tail.limit; + data[limit++] = (byte) ((i >> 24) & 0xff); + data[limit++] = (byte) ((i >> 16) & 0xff); + data[limit++] = (byte) ((i >> 8) & 0xff); + data[limit++] = (byte) (i & 0xff); + tail.limit = limit; + byteCount += 4; + return this; + } + + /** + * Returns a tail segment that we can write at least {@code minimumCapacity} + * bytes to, creating it if necessary. + */ + Segment writableSegment(int minimumCapacity) { + if (minimumCapacity < 1 || minimumCapacity > Segment.SIZE) throw new IllegalArgumentException(); + if (head == null) { head = SegmentPool.INSTANCE.take(); // Acquire a first segment. return head.next = head.prev = head; } Segment tail = head.prev; - if (tail.limit == Segment.SIZE) { + if (tail.limit + minimumCapacity > Segment.SIZE) { tail = tail.push(SegmentPool.INSTANCE.take()); // Append a new empty segment to fill up. } return tail; @@ -264,14 +389,14 @@ public final class OkBuffer implements Source, Sink { */ @Override public String toString() { if (byteCount > 0x100000) return super.toString(); - char[] result = new char[(int) (byteCount * 2)]; + int charCount = (int) (byteCount * 2); + char[] result = new char[charCount]; int offset = 0; - for (Segment s = head; offset < byteCount; s = s.next) { + for (Segment s = head; offset < charCount; s = s.next) { for (int i = s.pos; i < s.limit; i++) { result[offset++] = HEX_DIGITS[(s.data[i] >> 4) & 0xf]; result[offset++] = HEX_DIGITS[s.data[i] & 0xf]; } - offset += s.limit - s.pos; } return new String(result); } diff --git a/okhttp-protocols/src/main/java/com/squareup/okhttp/internal/bytes/OkBuffers.java b/okhttp-protocols/src/main/java/com/squareup/okhttp/internal/bytes/OkBuffers.java index 0f8583953..ec7c60a3e 100644 --- a/okhttp-protocols/src/main/java/com/squareup/okhttp/internal/bytes/OkBuffers.java +++ b/okhttp-protocols/src/main/java/com/squareup/okhttp/internal/bytes/OkBuffers.java @@ -67,7 +67,7 @@ public final class OkBuffers { OkBuffer sink, long byteCount, Deadline deadline) throws IOException { if (byteCount < 0) throw new IllegalArgumentException("byteCount < 0: " + byteCount); deadline.throwIfReached(); - Segment tail = sink.writableSegment(); + Segment tail = sink.writableSegment(1); int maxToCopy = (int) Math.min(byteCount, Segment.SIZE - tail.limit); int bytesRead = in.read(tail.data, tail.limit, maxToCopy); if (bytesRead == -1) return -1; diff --git a/okhttp-protocols/src/test/java/com/squareup/okhttp/internal/bytes/OkBufferTest.java b/okhttp-protocols/src/test/java/com/squareup/okhttp/internal/bytes/OkBufferTest.java index f91c180d0..ed52cb689 100644 --- a/okhttp-protocols/src/test/java/com/squareup/okhttp/internal/bytes/OkBufferTest.java +++ b/okhttp-protocols/src/test/java/com/squareup/okhttp/internal/bytes/OkBufferTest.java @@ -316,6 +316,106 @@ public final class OkBufferTest { assertEquals(-1, source.read(sink, 1, Deadline.NONE)); } + @Test public void writeBytes() throws Exception { + OkBuffer data = new OkBuffer(); + data.writeByte(0xab); + data.writeByte(0xcd); + assertEquals("abcd", data.toString()); + } + + @Test public void writeLastByteInSegment() throws Exception { + OkBuffer data = new OkBuffer(); + data.writeUtf8(repeat('a', Segment.SIZE - 1)); + data.writeByte(0x20); + data.writeByte(0x21); + assertEquals(asList(Segment.SIZE, 1), data.segmentSizes()); + assertEquals(repeat('a', Segment.SIZE - 1), data.readUtf8(Segment.SIZE - 1)); + assertEquals("2021", data.toString()); + } + + @Test public void writeShort() throws Exception { + OkBuffer data = new OkBuffer(); + data.writeShort(0xabcd); + data.writeShort(0x4321); + assertEquals("abcd4321", data.toString()); + } + + @Test public void writeInt() throws Exception { + OkBuffer data = new OkBuffer(); + data.writeInt(0xabcdef01); + data.writeInt(0x87654321); + assertEquals("abcdef0187654321", data.toString()); + } + + @Test public void writeLastIntegerInSegment() throws Exception { + OkBuffer data = new OkBuffer(); + data.writeUtf8(repeat('a', Segment.SIZE - 4)); + data.writeInt(0xabcdef01); + data.writeInt(0x87654321); + assertEquals(asList(Segment.SIZE, 4), data.segmentSizes()); + assertEquals(repeat('a', Segment.SIZE - 4), data.readUtf8(Segment.SIZE - 4)); + assertEquals("abcdef0187654321", data.toString()); + } + + @Test public void writeIntegerDoesntQuiteFitInSegment() throws Exception { + OkBuffer data = new OkBuffer(); + data.writeUtf8(repeat('a', Segment.SIZE - 3)); + data.writeInt(0xabcdef01); + data.writeInt(0x87654321); + assertEquals(asList(Segment.SIZE - 3, 8), data.segmentSizes()); + assertEquals(repeat('a', Segment.SIZE - 3), data.readUtf8(Segment.SIZE - 3)); + assertEquals("abcdef0187654321", data.toString()); + } + + @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(0, data.byteCount()); + } + + @Test public void readShort() throws Exception { + OkBuffer data = new OkBuffer(); + data.write(new ByteString(new byte[] { + (byte) 0xab, (byte) 0xcd, (byte) 0xef, (byte) 0x01 + })); + assertEquals((short) 0xabcd, data.readShort()); + assertEquals((short) 0xef01, data.readShort()); + assertEquals(0, data.byteCount()); + } + + @Test public void readShortSplitAcrossMultipleSegments() throws Exception { + OkBuffer data = new OkBuffer(); + data.writeUtf8(repeat('a', Segment.SIZE - 1)); + data.write(new ByteString(new byte[] { (byte) 0xab, (byte) 0xcd })); + data.readUtf8(Segment.SIZE - 1); + assertEquals((short) 0xabcd, data.readShort()); + assertEquals(0, data.byteCount()); + } + + @Test public void readInt() throws Exception { + OkBuffer data = new OkBuffer(); + data.write(new ByteString(new byte[] { + (byte) 0xab, (byte) 0xcd, (byte) 0xef, (byte) 0x01, + (byte) 0x87, (byte) 0x65, (byte) 0x43, (byte) 0x21 + })); + assertEquals(0xabcdef01, data.readInt()); + assertEquals(0x87654321, data.readInt()); + assertEquals(0, data.byteCount()); + } + + @Test public void readIntSplitAcrossMultipleSegments() throws Exception { + OkBuffer data = new OkBuffer(); + data.writeUtf8(repeat('a', Segment.SIZE - 3)); + data.write(new ByteString(new byte[] { + (byte) 0xab, (byte) 0xcd, (byte) 0xef, (byte) 0x01 + })); + data.readUtf8(Segment.SIZE - 3); + assertEquals(0xabcdef01, data.readInt()); + assertEquals(0, data.byteCount()); + } + private String repeat(char c, int count) { char[] array = new char[count]; Arrays.fill(array, c);