From c74e24495967f99e10873ff72902bf331a2655ee Mon Sep 17 00:00:00 2001
From: Looly <loolly@aliyun.com>
Date: Sat, 7 Sep 2024 19:23:12 +0800
Subject: [PATCH] fix code

---
 .../core/io/stream/LimitedInputStream.java    | 68 +++++++++++--------
 1 file changed, 38 insertions(+), 30 deletions(-)

diff --git a/hutool-core/src/main/java/org/dromara/hutool/core/io/stream/LimitedInputStream.java b/hutool-core/src/main/java/org/dromara/hutool/core/io/stream/LimitedInputStream.java
index d69f3f1ab..ef0e72955 100644
--- a/hutool-core/src/main/java/org/dromara/hutool/core/io/stream/LimitedInputStream.java
+++ b/hutool-core/src/main/java/org/dromara/hutool/core/io/stream/LimitedInputStream.java
@@ -16,66 +16,74 @@
 
 package org.dromara.hutool.core.io.stream;
 
+import org.dromara.hutool.core.lang.Assert;
+
 import java.io.FilterInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 
 /**
  * 限制读取最大长度的{@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
  */
 public class LimitedInputStream extends FilterInputStream {
 
-	private final long maxSize;
-	private long currentPos;
+	protected long limit;
+	private final boolean throwWhenReachLimit;
 
 	/**
 	 * 构造
 	 *
-	 * @param in      {@link InputStream}
-	 * @param maxSize 限制最大读取量，单位byte
+	 * @param in                  {@link InputStream}
+	 * @param limit               限制最大读取量，单位byte
+	 * @param throwWhenReachLimit 是否在达到限制时抛出异常，{@code false}则读取到限制后返回-1
 	 */
-	public LimitedInputStream(final InputStream in, final long maxSize) {
-		super(in);
-		this.maxSize = maxSize;
+	public LimitedInputStream(final InputStream in, final long limit, final boolean throwWhenReachLimit) {
+		super(Assert.notNull(in, "InputStream must not be null!"));
+		this.limit = Math.max(0L, limit);
+		this.throwWhenReachLimit = throwWhenReachLimit;
 	}
 
 	@Override
 	public int read() throws IOException {
-		final int data = super.read();
-		if (data != -1) {
-			currentPos++;
-			checkPos();
-		}
+		final int data = (limit == 0) ? -1 : super.read();
+		checkLimit(data);
+		limit = (data < 0) ? 0 : limit - 1;
 		return data;
 	}
 
-	@SuppressWarnings("NullableProblems")
 	@Override
 	public int read(final byte[] b, final int off, final int len) throws IOException {
-		final int count = super.read(b, off, len);
-		if (count > 0) {
-			currentPos += count;
-			checkPos();
-		}
-		return count;
+		final int length = (limit == 0) ? -1 : super.read(b, off, len > limit ? (int) limit : len);
+		checkLimit(length);
+		limit = (length < 0) ? 0 : limit - length;
+		return length;
 	}
 
 	@Override
-	public long skip(final long n) throws IOException {
-		final long skipped = super.skip(n);
-		if (skipped != 0) {
-			currentPos += skipped;
-			checkPos();
-		}
-		return skipped;
+	public long skip(final long len) throws IOException {
+		final long length = super.skip(Math.min(len, limit));
+		checkLimit(length);
+		limit -= length;
+		return length;
 	}
 
-	private void checkPos() {
-		if (currentPos > maxSize) {
-			throw new IllegalStateException("Read limit exceeded");
+	@Override
+	public int available() throws IOException {
+		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");
 		}
 	}
 }