This commit is contained in:
Looly 2024-09-07 19:23:12 +08:00
parent 7931966fb3
commit c74e244959

View File

@ -16,66 +16,74 @@
package org.dromara.hutool.core.io.stream; package org.dromara.hutool.core.io.stream;
import org.dromara.hutool.core.lang.Assert;
import java.io.FilterInputStream; import java.io.FilterInputStream;
import java.io.IOException; import java.io.IOException;
import java.io.InputStream; import java.io.InputStream;
/** /**
* 限制读取最大长度的{@link FilterInputStream} 实现<br> * 限制读取最大长度的{@link FilterInputStream} 实现<br>
* 来自https://github.com/skylot/jadx/blob/master/jadx-plugins/jadx-plugins-api/src/main/java/jadx/api/plugins/utils/LimitedInputStream.java
* *
* @author jadx * @author jadx
*/ */
public class LimitedInputStream extends FilterInputStream { public class LimitedInputStream extends FilterInputStream {
private final long maxSize; protected long limit;
private long currentPos; private final boolean throwWhenReachLimit;
/** /**
* 构造 * 构造
* *
* @param in {@link InputStream} * @param in {@link InputStream}
* @param maxSize 限制最大读取量单位byte * @param limit 限制最大读取量单位byte
* @param throwWhenReachLimit 是否在达到限制时抛出异常{@code false}则读取到限制后返回-1
*/ */
public LimitedInputStream(final InputStream in, final long maxSize) { public LimitedInputStream(final InputStream in, final long limit, final boolean throwWhenReachLimit) {
super(in); super(Assert.notNull(in, "InputStream must not be null!"));
this.maxSize = maxSize; this.limit = Math.max(0L, limit);
this.throwWhenReachLimit = throwWhenReachLimit;
} }
@Override @Override
public int read() throws IOException { public int read() throws IOException {
final int data = super.read(); final int data = (limit == 0) ? -1 : super.read();
if (data != -1) { checkLimit(data);
currentPos++; limit = (data < 0) ? 0 : limit - 1;
checkPos();
}
return data; return data;
} }
@SuppressWarnings("NullableProblems")
@Override @Override
public int read(final byte[] b, final int off, final int len) throws IOException { public int read(final byte[] b, final int off, final int len) throws IOException {
final int count = super.read(b, off, len); final int length = (limit == 0) ? -1 : super.read(b, off, len > limit ? (int) limit : len);
if (count > 0) { checkLimit(length);
currentPos += count; limit = (length < 0) ? 0 : limit - length;
checkPos(); return length;
}
return count;
} }
@Override @Override
public long skip(final long n) throws IOException { public long skip(final long len) throws IOException {
final long skipped = super.skip(n); final long length = super.skip(Math.min(len, limit));
if (skipped != 0) { checkLimit(length);
currentPos += skipped; limit -= length;
checkPos(); return length;
}
return skipped;
} }
private void checkPos() { @Override
if (currentPos > maxSize) { public int available() throws IOException {
throw new IllegalStateException("Read limit exceeded"); final int length = super.available();
return length > limit ? (int)limit : length;
}
/**
* 检查读取数据是否达到限制
*
* @param data 读取的数据
* @throws IOException 定义了限制并设置达到限制后抛出异常
*/
private void checkLimit(final long data) throws IOException {
if (data < 0 && limit > 0 && throwWhenReachLimit) {
throw new IOException("Read limit exceeded");
} }
} }
} }