com.groupbyinc.common.util.stream.RestrictedInputStream Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of common-test Show documentation

${project.name}

The newest version!

package com.groupbyinc.common.util.stream;

import java.io.IOException;
import java.io.InputStream;

public class RestrictedInputStream extends InputStream {

  public static final String EXCEEDED_MAX_BYTES = "InputStream exceeded maximum byte size";

  private final InputStream original;

  private final long maxBytes;

  private long totalBytesRead;

  public RestrictedInputStream(InputStream original, long maxBytes) {
    this.original = original;
    this.maxBytes = maxBytes;
  }

  @Override
  public int read() throws IOException {
    int b = original.read();
    if (b >= 0) {
      incrementByteCounter(1);
    }
    return b;
  }

  @Override
  public int read(byte b[]) throws IOException {
    return read(b, 0, b.length);
  }

  @Override
  public int read(byte b[], int off, int len) throws IOException {
    int bytesRead = original.read(b, off, len);
    if (bytesRead >= 0) {
      incrementByteCounter(bytesRead);
    }
    return bytesRead;
  }

  private void incrementByteCounter(int bytesRead) throws IOException {
    totalBytesRead += bytesRead;
    if (totalBytesRead > maxBytes) {
      throw new IOException(EXCEEDED_MAX_BYTES + ": " + maxBytes);
    }
  }
}