001package org.jsoup.internal;
002
003import org.jsoup.helper.Validate;
004
005import java.io.BufferedInputStream;
006import java.io.ByteArrayOutputStream;
007import java.io.IOException;
008import java.io.InputStream;
009import java.net.SocketTimeoutException;
010import java.nio.ByteBuffer;
011
012/**
013 * A jsoup internal class (so don't use it as there is no contract API) that enables constraints on an Input Stream,
014 * namely a maximum read size, and the ability to Thread.interrupt() the read.
015 */
016public final class ConstrainableInputStream extends BufferedInputStream {
017    private static final int DefaultSize = 1024 * 32;
018
019    private final boolean capped;
020    private final int maxSize;
021    private long startTime;
022    private long timeout = -1; // optional max time of request
023    private int remaining;
024    private boolean interrupted;
025
026    private ConstrainableInputStream(InputStream in, int bufferSize, int maxSize) {
027        super(in, bufferSize);
028        Validate.isTrue(maxSize >= 0);
029        this.maxSize = maxSize;
030        remaining = maxSize;
031        capped = maxSize != 0;
032        startTime = System.nanoTime();
033    }
034
035    /**
036     * If this InputStream is not already a ConstrainableInputStream, let it be one.
037     * @param in the input stream to (maybe) wrap
038     * @param bufferSize the buffer size to use when reading
039     * @param maxSize the maximum size to allow to be read. 0 == infinite.
040     * @return a constrainable input stream
041     */
042    public static ConstrainableInputStream wrap(InputStream in, int bufferSize, int maxSize) {
043        return in instanceof ConstrainableInputStream
044            ? (ConstrainableInputStream) in
045            : new ConstrainableInputStream(in, bufferSize, maxSize);
046    }
047
048    @Override
049    public int read(byte[] b, int off, int len) throws IOException {
050        if (interrupted || capped && remaining <= 0)
051            return -1;
052        if (Thread.interrupted()) {
053            // interrupted latches, because parse() may call twice (and we still want the thread interupt to clear)
054            interrupted = true;
055            return -1;
056        }
057        if (expired())
058            throw new SocketTimeoutException("Read timeout");
059
060        if (capped && len > remaining)
061            len = remaining; // don't read more than desired, even if available
062
063        try {
064            final int read = super.read(b, off, len);
065            remaining -= read;
066            return read;
067        } catch (SocketTimeoutException e) {
068            return 0;
069        }
070    }
071
072    /**
073     * Reads this inputstream to a ByteBuffer. The supplied max may be less than the inputstream's max, to support
074     * reading just the first bytes.
075     */
076    public ByteBuffer readToByteBuffer(int max) throws IOException {
077        Validate.isTrue(max >= 0, "maxSize must be 0 (unlimited) or larger");
078        final boolean localCapped = max > 0; // still possibly capped in total stream
079        final int bufferSize = localCapped && max < DefaultSize ? max : DefaultSize;
080        final byte[] readBuffer = new byte[bufferSize];
081        final ByteArrayOutputStream outStream = new ByteArrayOutputStream(bufferSize);
082
083        int read;
084        int remaining = max;
085
086        while (true) {
087            read = read(readBuffer);
088            if (read == -1) break;
089            if (localCapped) { // this local byteBuffer cap may be smaller than the overall maxSize (like when reading first bytes)
090                if (read >= remaining) {
091                    outStream.write(readBuffer, 0, remaining);
092                    break;
093                }
094                remaining -= read;
095            }
096            outStream.write(readBuffer, 0, read);
097        }
098        return ByteBuffer.wrap(outStream.toByteArray());
099    }
100
101    @Override
102    public void reset() throws IOException {
103        super.reset();
104        remaining = maxSize - markpos;
105    }
106
107    public ConstrainableInputStream timeout(long startTimeNanos, long timeoutMillis) {
108        this.startTime = startTimeNanos;
109        this.timeout = timeoutMillis * 1000000;
110        return this;
111    }
112
113    private boolean expired() {
114        if (timeout == -1)
115            return false;
116
117        final long now = System.nanoTime();
118        final long dur = now - startTime;
119        return (dur > timeout);
120    }
121}