gateway-spi/src/main/java/org/apache/knox/gateway/dispatch/CappedBufferHttpEntity.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.knox.gateway.dispatch;

import org.apache.commons.io.IOUtils;
import org.apache.http.HttpEntity;
import org.apache.http.entity.HttpEntityWrapper;

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

public class CappedBufferHttpEntity extends HttpEntityWrapper {

  public static final int DEFAULT_BUFFER_SIZE = 4096;

  private int replayWriteIndex;
  private int replayWriteLimit;
  private byte[] replayBuffer;
  private InputStream wrappedStream;

  public CappedBufferHttpEntity( final HttpEntity entity, int bufferSize ) throws IOException {
    super( entity );
    this.wrappedStream = null;
    this.replayWriteIndex = -1;
    if( !entity.isRepeatable() ) {
      this.replayBuffer = new byte[ bufferSize ];
      this.replayWriteLimit = bufferSize-1;
    } else {
      this.replayBuffer = null;
    }
  }

  public CappedBufferHttpEntity( final HttpEntity entity ) throws IOException {
    this( entity, DEFAULT_BUFFER_SIZE );
  }

  @Override
  public boolean isRepeatable() {
    return true;
  }

  @Override
  public boolean isStreaming() {
    return wrappedEntity.isStreaming();
  }

  @Override
  public boolean isChunked() {
    return wrappedEntity.isChunked();
  }

  @Override
  public long getContentLength() {
    return wrappedEntity.getContentLength();
  }

  // This will throw an IOException if an attempt is made to getContent a second time after
  // more bytes than the buffer can hold has been read on the first stream.
  @Override
  public InputStream getContent() throws IOException {
    // If the wrapped stream is repeatable return it directly.
    if( replayBuffer == null ) {
      return wrappedEntity.getContent();
    // Else if the buffer has overflowed
    } else {
      if( wrappedStream == null ) {
         wrappedStream = wrappedEntity.getContent();
      }
      return new ReplayStream();
    }
  }

  @Override
  public void writeTo( final OutputStream stream ) throws IOException {
    IOUtils.copy( getContent(), stream );
  }

  @SuppressWarnings("deprecation")
  @Override
  public void consumeContent() throws IOException {
    throw new UnsupportedOperationException();
  }

  private class ReplayStream extends InputStream {

    private int replayReadIndex = -1;

    @Override
    public int read() throws IOException {
      int b;
      // If we can read from the buffer do so.
      if( replayReadIndex < replayWriteIndex ) {
        b = replayBuffer[ ++replayReadIndex ];
      } else {
        b = wrappedStream.read();
        // If the underlying stream is not closed.
        if( b > -1 ) {
          if( replayWriteIndex < replayWriteLimit ) {
            replayBuffer[ ++replayWriteIndex ] = (byte)b;
            replayReadIndex++;
          } else {
            throw new IOException("Hit replay buffer max limit");
          }
        }
      }
      return b;
    }

    @Override
    public int read(byte[] buffer, int offset, int limit ) throws IOException {
      int count;
      // If we can read from the buffer do so.
      if( replayReadIndex < replayWriteIndex ) {
        count = replayWriteIndex - replayReadIndex;
        count = Math.min( limit, count );
        System.arraycopy( replayBuffer, replayReadIndex+1, buffer, offset, count );
        replayReadIndex += count;
      } else {
        count = wrappedStream.read( buffer, offset, limit );
        // If the underlying stream is not closed.
        if( count > -1 ) {
          if( replayWriteIndex+count < replayWriteLimit ) {
            System.arraycopy( buffer, offset, replayBuffer, replayWriteIndex+1, count );
            replayReadIndex += count;
            replayWriteIndex += count;
          } else {
            throw new IOException("Hit replay buffer max limit");
          }
        }
      }
      return count;
    }

  }

}