Merge korg/donut into korg/master

src/com/google/common/io/protocol/BoundInputStream.java

+// Copyright 2008 Google Inc. All Rights Reserved.
+
+package com.google.common.io.protocol;
+
+import java.io.*;
+
+/**
+ * An input stream backed by another input stream, where reading from the
+ * underlying input stream is limited to a fixed number of bytes. Also does
+ * some buffering.
+ *
+ */
+public class BoundInputStream extends InputStream {
+
+  /** Buffer size */
+  static final int BUF_SIZE = 4096;
+
+  /** Number of bytes that may still be read from the underlying stream */
+  private int remaining;
+
+  /** Small buffer to avoid making OS calls for each byte read. */
+  private byte[] buf;
+
+  /** Current position in the buffer */
+  private int bufPos;
+
+  /** Filled size of the buffer */
+  private int bufSize;
+
+  /** Underlying stream to read from */
+  private InputStream base;
+
+  public BoundInputStream(InputStream base, int len) {
+    this.base = base;
+    this.remaining = len;
+
+    buf = new byte[Math.min(len, BUF_SIZE)];
+  }
+
+  /**
+   * Make sure there is at least one byte in the buffer. If not possible,
+   * return false.
+   */
+  private boolean checkBuf() throws IOException {
+    if (remaining <= 0) {
+      return false;
+    }
+
+    if (bufPos >= bufSize) {
+      bufSize = base.read(buf, 0, Math.min(remaining, buf.length));
+      if (bufSize <= 0) {
+        remaining = 0;
+        return false;
+      }
+      bufPos = 0;
+    }
+    return true;
+  }
+
+  public int available() {
+    return bufSize - bufPos;
+  }
+
+  public int read() throws IOException {
+    if (!checkBuf()) {
+      return -1;
+    }
+    remaining--;
+    return buf[bufPos++] & 255;
+  }
+
+  public int read(byte[] data, int start, int count) throws IOException {
+    if (!checkBuf()) {
+      return -1;
+    }
+    count = Math.min(count, bufSize - bufPos);
+    System.arraycopy(buf, bufPos, data, start, count);
+    bufPos += count;
+    remaining -= count;
+    return count;
+  }
+
+  /**
+   * How many bytes are remaining, based on the length provided to the
+   * constructor. The underlying stream may terminate earlier. Provided mainly
+   * for testing purposes.
+   */
+  public int getRemaining() {
+    return remaining;
+  }
+}

src/com/google/common/io/protocol/ProtoBufUtil.java

@@ -110,6 +110,32 @@ public final class ProtoBufUtil {
     }
   }
 
+  /**
+   * Returns an input stream for reading protocol buffer
+   * responses. This method reads a 32-bit signed integer from the
+   * stream, which determines the data size and compression. If the
+   * integer is negative, indicating a GZipped input stream which we
+   * do not support, an exception is thrown. Otherwise, just a
+   * BoundInputStream is returned. The input stream returned is always
+   * limited to the data available.
+   *
+   * @param dataInput the data input to read from
+   * @return an input stream, limited to the data size read from the stream
+   * @throws IOException if the incoming stream is gzipped.
+   */
+  public static InputStream getInputStreamForProtoBufResponse(
+      DataInput dataInput) throws IOException {
+
+    int size = dataInput.readInt();
+    InputStream is = new BoundInputStream((InputStream) dataInput,
+        Math.abs(size));
+
+    if (size < 0) {
+        throw new IOException("Cannot read gzipped streams");
+    }
+    return is;
+  }
+
   /**
    * Reads a single protocol buffer from the given input stream. This method is
    * provided where the client needs incremental access to the contents of a
@@ -143,6 +169,28 @@ public final class ProtoBufUtil {
     return tag;
   }
 
+  /**
+   * Reads a size int and a protocol buffer from a DataInput. If the size
+   * is negative, this is interpreted as an indicator that the protocol buffer
+   * is packed with GZIP. In this case, -size bytes are read, and the data is
+   * unpacked with GZIP before constructing the protocol buffer.
+   *
+   * @param protoBufType the protocol buffer type to read
+   * @param dataInput the data input to read from
+   * @return a protocol buffer of the given type
+   * @throws IOException
+   */
+  public static ProtoBuf readProtoBufResponse(ProtoBufType protoBufType,
+      DataInput dataInput) throws IOException {
+    ProtoBuf response = new ProtoBuf(protoBufType);
+    InputStream is = getInputStreamForProtoBufResponse(dataInput);
+    response.parse(is);
+    if (is.read() != -1) {
+      throw new IOException();
+    }
+    return response;
+  }
+
   /**
    * A wrapper for <code> getProtoValueOrNegativeOne </code> that drills into
    * a sub message returning the long value if it exists, returning -1 if it