com.subgraph.vega.internal.http.requests.RepeatableStreamingEntity.java Source code

Java tutorial

Introduction

Here is the source code for com.subgraph.vega.internal.http.requests.RepeatableStreamingEntity.java

Source

/*******************************************************************************
 * Copyright (c) 2011 Subgraph.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * which accompanies this distribution, and is available at
 * http://www.eclipse.org/legal/epl-v10.html
 * 
 * Contributors:
 *     Subgraph - initial API and implementation
 ******************************************************************************/
package com.subgraph.vega.internal.http.requests;

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

import org.apache.http.HttpEntity;
import org.apache.http.entity.AbstractHttpEntity;
import org.apache.http.entity.ByteArrayEntity;
import org.apache.http.util.ByteArrayBuffer;
import org.apache.http.util.EntityUtils;

/**
 * An entity wrapper which converts a non-repeatable streaming entity
 * into a repeatable entity by storing the data as it streams in from
 * a connection (or other streaming source).  When all streaming data has
 * been received, the entity behaves like a ByteArrayEntity. 
 */
public class RepeatableStreamingEntity extends AbstractHttpEntity {
    private final static int BUFFER_SIZE = 8192;
    private boolean isStreaming = true;
    private long length;
    private long maximumInputKilobytes = 0; // 0 = no maximum length
    private HttpEntity bufferEntity;
    private volatile InputStream input;

    /**
     * @param consumeInput Boolean indicating whether to immediately consume all input into buffer. 
     * @throws IOException 
     */
    RepeatableStreamingEntity(InputStream input, long length, boolean consumeInput, boolean isChunked,
            String contentType, String contentEncoding) throws IOException {
        setChunked(isChunked);
        setContentType(contentType);
        setContentEncoding(contentEncoding);
        setActiveInputStream(input, length);
        if (consumeInput != false) {
            consumeAllInput();
        }
    }

    RepeatableStreamingEntity(HttpEntity originalEntity) throws IOException {
        copyEntityProperties(originalEntity);
        if (originalEntity.isStreaming())
            setActiveInputStream(originalEntity.getContent(), originalEntity.getContentLength());
        else
            setActiveByteArrayEntity(EntityUtils.toByteArray(originalEntity));
    }

    void setMaximumInputKilobytes(int kb) {
        maximumInputKilobytes = kb;
    }

    private void copyEntityProperties(HttpEntity e) {
        setChunked(e.isChunked());
        if (e.getContentType() != null)
            setContentType(e.getContentType().getValue());
        if (e.getContentEncoding() != null)
            setContentEncoding(e.getContentType().getValue());
        length = e.getContentLength();
    }

    private void setActiveInputStream(InputStream input, long length) {
        if (length > Integer.MAX_VALUE)
            throw new IllegalArgumentException("HTTP entity is too large to be buffered in memory: " + length);
        this.length = length;
        final int sz = (length < 0) ? (BUFFER_SIZE) : ((int) length);
        this.input = new CachingInputStream(input, sz);
        isStreaming = true;
    }

    private void setActiveByteArrayEntity(byte[] content) {
        ByteArrayEntity entity = new ByteArrayEntity(content);
        isStreaming = false;
        length = content.length;
        bufferEntity = entity;
        input = null;
    }

    private void consumeAllInput() throws IOException {
        int rv;
        if (length < 0) {
            while ((rv = input.read()) != -1)
                ;
        } else {
            int remaining = (int) this.length;
            byte[] buffer = new byte[BUFFER_SIZE];
            while (remaining > 0) {
                int sz = (remaining < BUFFER_SIZE) ? (remaining) : (BUFFER_SIZE);
                rv = input.read(buffer, 0, sz);
                if (rv == -1)
                    break;
                remaining -= rv;
            }
        }
    }

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

    @Override
    public long getContentLength() {
        return length;
    }

    @Override
    public InputStream getContent() throws IOException {
        if (input != null)
            return input;
        else
            return bufferEntity.getContent();
    }

    @Override
    public void writeTo(OutputStream outstream) throws IOException {
        if (input == null) {
            bufferEntity.writeTo(outstream);
            return;
        }

        byte[] buffer = new byte[BUFFER_SIZE];
        int l;
        if (length < 0) {
            // consume until EOF
            while ((l = input.read(buffer)) != -1) {
                try {
                    outstream.write(buffer, 0, l);
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        } else {
            // consume no more than length
            long remaining = this.length;

            while (remaining > 0) {
                int sz = (remaining < BUFFER_SIZE) ? ((int) remaining) : (BUFFER_SIZE);
                l = input.read(buffer, 0, sz);
                if (l == -1)
                    break;
                outstream.write(buffer, 0, l);
                remaining -= l;
            }
        }
        if (input != null) {
            input.close();
            input = null;
        }
    }

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

    @Override
    public void consumeContent() throws IOException {
        if (input != null) {
            input.close();
            input = null;
        }
        if (bufferEntity != null) {
            EntityUtils.consume(bufferEntity);
        }
    }

    private class CachingInputStream extends InputStream {
        private final InputStream wrappedInput;
        private ByteArrayBuffer buffer;
        private boolean eof;

        CachingInputStream(InputStream input, int bufferSize) {
            wrappedInput = input;
            buffer = new ByteArrayBuffer(bufferSize);
        }

        @Override
        public int read() throws IOException {
            if (eof) {
                return -1;
            }

            final int b = wrappedInput.read();
            if (b == -1) {
                processEOF();
            } else {
                buffer.append(b);
            }
            checkMaximumLength();
            return b;
        }

        public int read(byte[] b, int off, int len) throws IOException {
            if (eof) {
                return -1;
            }

            final int n = wrappedInput.read(b, off, len);
            if (n == -1) {
                processEOF();
            } else {
                buffer.append(b, off, n);
            }
            checkMaximumLength();
            return n;
        }

        private void checkMaximumLength() throws IOException {
            if (maximumInputKilobytes > 0 && buffer.length() > (maximumInputKilobytes * 1024)) {
                wrappedInput.close();
                eof = true;
                buffer = null;
                setActiveByteArrayEntity(new byte[0]);
                throw new IOException(
                        "Maximum length of " + maximumInputKilobytes + " kb exceeded while streaming http entity.");
            }
        }

        public void close() throws IOException {
            wrappedInput.close();
            if (!eof) {
                setActiveByteArrayEntity(buffer.toByteArray());
            }
        }

        private void processEOF() throws IOException {
            wrappedInput.close();
            eof = true;
            setActiveByteArrayEntity(buffer.toByteArray());
        }
    }
}