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}