task-65: transient payload 读取流式解压 + 解压后字节上限
- 新增 aiimage.transient-storage.max-decompressed-payload-bytes 配置 (默认 100MB,即 2x 存储上限),覆盖 gzip 二进制与 gzip64 兼容两条路径 - decodeGzipStream 按 8KB 缓冲流式解压,累计输出超过上限立即中止并报错, 防止压缩炸弹(zip bomb)在读取时无界膨胀内存 - 超限错误保留可识别消息(exceeds configured limit),由 resolvePayload 统一包装
This commit is contained in:
@@ -37,6 +37,11 @@ public class TransientStorageProperties {
|
|||||||
private long warnPayloadBytes = 5L * 1024 * 1024;
|
private long warnPayloadBytes = 5L * 1024 * 1024;
|
||||||
private long maxPayloadBytes = 50L * 1024 * 1024;
|
private long maxPayloadBytes = 50L * 1024 * 1024;
|
||||||
private long maxStoredPayloadBytes = 50L * 1024 * 1024;
|
private long maxStoredPayloadBytes = 50L * 1024 * 1024;
|
||||||
|
/**
|
||||||
|
* 读取端解压后字节上限:防止压缩炸弹(zip bomb)在流式解压时无界膨胀内存。
|
||||||
|
* 默认 2x 存储上限,覆盖 gzip 二进制路径与 gzip64 兼容路径。
|
||||||
|
*/
|
||||||
|
private long maxDecompressedPayloadBytes = 100L * 1024 * 1024;
|
||||||
private boolean fallbackToLocalOnOversize = true;
|
private boolean fallbackToLocalOnOversize = true;
|
||||||
private boolean deleteRetryEnabled = true;
|
private boolean deleteRetryEnabled = true;
|
||||||
private String deleteRetryCron = "0 */5 * * * *";
|
private String deleteRetryCron = "0 */5 * * * *";
|
||||||
|
|||||||
+35
-8
@@ -467,19 +467,15 @@ public class TransientPayloadStorageService {
|
|||||||
return "";
|
return "";
|
||||||
}
|
}
|
||||||
if (isGzip(storedContent)) {
|
if (isGzip(storedContent)) {
|
||||||
try (GZIPInputStream gzip = new GZIPInputStream(new ByteArrayInputStream(storedContent))) {
|
return decodeGzipStream(storedContent);
|
||||||
return new String(gzip.readAllBytes(), StandardCharsets.UTF_8);
|
|
||||||
} catch (Exception ex) {
|
|
||||||
throw new IllegalStateException("failed to decode transient payload", ex);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
String legacyText = new String(storedContent, StandardCharsets.UTF_8);
|
String legacyText = new String(storedContent, StandardCharsets.UTF_8);
|
||||||
if (legacyText.startsWith("gzip64:")) {
|
if (legacyText.startsWith("gzip64:")) {
|
||||||
try {
|
try {
|
||||||
byte[] compressed = Base64.getDecoder().decode(legacyText.substring("gzip64:".length()));
|
byte[] compressed = Base64.getDecoder().decode(legacyText.substring("gzip64:".length()));
|
||||||
try (GZIPInputStream gzip = new GZIPInputStream(new ByteArrayInputStream(compressed))) {
|
return decodeGzipStream(compressed);
|
||||||
return new String(gzip.readAllBytes(), StandardCharsets.UTF_8);
|
} catch (IllegalStateException ex) {
|
||||||
}
|
throw ex;
|
||||||
} catch (Exception ex) {
|
} catch (Exception ex) {
|
||||||
throw new IllegalStateException("failed to decode transient payload", ex);
|
throw new IllegalStateException("failed to decode transient payload", ex);
|
||||||
}
|
}
|
||||||
@@ -488,6 +484,37 @@ public class TransientPayloadStorageService {
|
|||||||
return legacyText;
|
return legacyText;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 流式解压 gzip 字节流,解压输出累计超过 maxDecompressedPayloadBytes 立即中止,
|
||||||
|
* 防止压缩炸弹在内存中无界膨胀。上限未配置(<=0)时不限流。
|
||||||
|
*/
|
||||||
|
private String decodeGzipStream(byte[] compressed) {
|
||||||
|
long maxBytes = properties.getMaxDecompressedPayloadBytes();
|
||||||
|
try {
|
||||||
|
ByteArrayOutputStream out = new ByteArrayOutputStream();
|
||||||
|
try (GZIPInputStream gzip = new GZIPInputStream(new ByteArrayInputStream(compressed))) {
|
||||||
|
byte[] buffer = new byte[8192];
|
||||||
|
int read;
|
||||||
|
while ((read = gzip.read(buffer)) >= 0) {
|
||||||
|
if (read == 0) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if (isPositiveLimit(maxBytes) && (long) out.size() + read > maxBytes) {
|
||||||
|
throw new IllegalStateException(
|
||||||
|
"transient payload decompressed size exceeds configured limit "
|
||||||
|
+ maxBytes + " bytes");
|
||||||
|
}
|
||||||
|
out.write(buffer, 0, read);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return out.toString(StandardCharsets.UTF_8);
|
||||||
|
} catch (IllegalStateException ex) {
|
||||||
|
throw ex;
|
||||||
|
} catch (Exception ex) {
|
||||||
|
throw new IllegalStateException("failed to decode transient payload", ex);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
private boolean isGzip(byte[] content) {
|
private boolean isGzip(byte[] content) {
|
||||||
return content.length >= 2
|
return content.length >= 2
|
||||||
&& (content[0] & 0xFF) == 0x1F
|
&& (content[0] & 0xFF) == 0x1F
|
||||||
|
|||||||
+240
@@ -0,0 +1,240 @@
|
|||||||
|
package com.nanri.aiimage.modules.task.service;
|
||||||
|
|
||||||
|
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||||
|
import com.nanri.aiimage.config.InstanceMetadata;
|
||||||
|
import com.nanri.aiimage.config.StorageProperties;
|
||||||
|
import com.nanri.aiimage.config.TransientStorageProperties;
|
||||||
|
import com.nanri.aiimage.modules.file.service.object.RustfsObjectStorageService;
|
||||||
|
import com.nanri.aiimage.modules.file.service.oss.OssStorageService;
|
||||||
|
import com.nanri.aiimage.modules.task.mapper.TaskChunkMapper;
|
||||||
|
import com.nanri.aiimage.modules.task.mapper.TaskScopeStateMapper;
|
||||||
|
import org.junit.jupiter.api.Test;
|
||||||
|
import org.junit.jupiter.api.io.TempDir;
|
||||||
|
|
||||||
|
import java.io.ByteArrayOutputStream;
|
||||||
|
import java.nio.charset.StandardCharsets;
|
||||||
|
import java.nio.file.Path;
|
||||||
|
import java.util.ArrayList;
|
||||||
|
import java.util.List;
|
||||||
|
import java.util.zip.GZIPOutputStream;
|
||||||
|
|
||||||
|
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||||
|
import static org.junit.jupiter.api.Assertions.assertThrows;
|
||||||
|
import static org.junit.jupiter.api.Assertions.assertTrue;
|
||||||
|
import static org.mockito.ArgumentMatchers.any;
|
||||||
|
import static org.mockito.ArgumentMatchers.anyBoolean;
|
||||||
|
import static org.mockito.ArgumentMatchers.anyString;
|
||||||
|
import static org.mockito.Mockito.doReturn;
|
||||||
|
import static org.mockito.Mockito.mock;
|
||||||
|
import static org.mockito.Mockito.times;
|
||||||
|
import static org.mockito.Mockito.verify;
|
||||||
|
import static org.mockito.Mockito.when;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Task 65:transient payload 读取增加流式解压和解压后字节上限。
|
||||||
|
* 读取路径(resolvePayload)从对象存储拿到 gzip 二进制流后流式解压,
|
||||||
|
* 解压输出字节数超过 maxDecompressedPayloadBytes 立即中止并报错,
|
||||||
|
* 防止压缩炸弹(zip bomb)导致无界内存增长;旧的 gzip64 文本兼容分支同样受限。
|
||||||
|
*/
|
||||||
|
class TransientPayloadStorageDecompressionTest {
|
||||||
|
|
||||||
|
@TempDir
|
||||||
|
Path tempDir;
|
||||||
|
|
||||||
|
private static final long BIG_LIMIT = 100_000L;
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void test_task_065_payload_normal_default_path() throws Exception {
|
||||||
|
RustfsObjectStorageService rustfs = mock(RustfsObjectStorageService.class);
|
||||||
|
when(rustfs.isConfigured()).thenReturn(true);
|
||||||
|
when(rustfs.uploadBytes(anyString(), any(byte[].class), anyBoolean()))
|
||||||
|
.thenAnswer(invocation -> invocation.getArgument(0));
|
||||||
|
TransientPayloadStorageService service = newService(rustfs, propertiesWithDecompressionLimit(BIG_LIMIT));
|
||||||
|
String content = compressiblePayload(3000);
|
||||||
|
|
||||||
|
String pointer = service.storeParsedPayloadFast("TEST", 1L, "scope", content, false);
|
||||||
|
when(rustfs.readObjectBytes(anyString())).thenReturn(gzipBytes(content));
|
||||||
|
|
||||||
|
assertEquals(content, service.resolvePayload(pointer, "read failed"));
|
||||||
|
verify(rustfs).readObjectBytes(anyString());
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void test_task_065_payload_normal_multiple_items() throws Exception {
|
||||||
|
RustfsObjectStorageService rustfs = mock(RustfsObjectStorageService.class);
|
||||||
|
when(rustfs.isConfigured()).thenReturn(true);
|
||||||
|
when(rustfs.uploadBytes(anyString(), any(byte[].class), anyBoolean()))
|
||||||
|
.thenAnswer(invocation -> invocation.getArgument(0));
|
||||||
|
TransientPayloadStorageService service = newService(rustfs, propertiesWithDecompressionLimit(BIG_LIMIT));
|
||||||
|
List<String> contents = List.of(
|
||||||
|
compressiblePayload(1000), "{\"i\":1}", compressiblePayload(500), "{\"i\":2}");
|
||||||
|
List<String> pointers = new ArrayList<>();
|
||||||
|
for (int i = 0; i < contents.size(); i++) {
|
||||||
|
pointers.add(service.storeChunkPayload("TEST", 10L + i, "scope-" + i, i, contents.get(i)));
|
||||||
|
}
|
||||||
|
when(rustfs.readObjectBytes(anyString())).thenAnswer(invocation -> {
|
||||||
|
String objectKey = invocation.getArgument(0);
|
||||||
|
int index = Integer.parseInt(objectKey.substring(objectKey.indexOf("chunk-") + "chunk-".length(),
|
||||||
|
objectKey.indexOf(".json")));
|
||||||
|
return gzipBytes(contents.get(index));
|
||||||
|
});
|
||||||
|
|
||||||
|
for (int i = 0; i < contents.size(); i++) {
|
||||||
|
assertEquals(contents.get(i), service.resolvePayload(pointers.get(i), "read failed"));
|
||||||
|
}
|
||||||
|
verify(rustfs, times(contents.size())).readObjectBytes(anyString());
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void test_task_065_payload_normal_repeated_operation_is_idempotent() throws Exception {
|
||||||
|
RustfsObjectStorageService rustfs = mock(RustfsObjectStorageService.class);
|
||||||
|
when(rustfs.isConfigured()).thenReturn(true);
|
||||||
|
when(rustfs.uploadBytes(anyString(), any(byte[].class), anyBoolean()))
|
||||||
|
.thenAnswer(invocation -> invocation.getArgument(0));
|
||||||
|
TransientPayloadStorageService service = newService(rustfs, propertiesWithDecompressionLimit(BIG_LIMIT));
|
||||||
|
String content = compressiblePayload(1500);
|
||||||
|
String pointer = service.storeParsedPayloadFast("TEST", 1L, "scope", content, false);
|
||||||
|
when(rustfs.readObjectBytes(anyString())).thenReturn(gzipBytes(content));
|
||||||
|
|
||||||
|
assertEquals(content, service.resolvePayload(pointer, "read failed"));
|
||||||
|
assertEquals(content, service.resolvePayload(pointer, "read failed"));
|
||||||
|
|
||||||
|
verify(rustfs, times(2)).readObjectBytes(anyString());
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void test_task_065_payload_boundary_empty_input() throws Exception {
|
||||||
|
RustfsObjectStorageService rustfs = mock(RustfsObjectStorageService.class);
|
||||||
|
when(rustfs.isConfigured()).thenReturn(true);
|
||||||
|
when(rustfs.uploadBytes(anyString(), any(byte[].class), anyBoolean()))
|
||||||
|
.thenAnswer(invocation -> invocation.getArgument(0));
|
||||||
|
TransientPayloadStorageService service = newService(rustfs, propertiesWithDecompressionLimit(BIG_LIMIT));
|
||||||
|
|
||||||
|
String pointer = service.storeParsedPayloadFast("TEST", 1L, "scope", "", false);
|
||||||
|
when(rustfs.readObjectBytes(anyString())).thenReturn(gzipBytes(""));
|
||||||
|
|
||||||
|
assertEquals("", service.resolvePayload(pointer, "read failed"));
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void test_task_065_payload_boundary_single_item() throws Exception {
|
||||||
|
RustfsObjectStorageService rustfs = mock(RustfsObjectStorageService.class);
|
||||||
|
when(rustfs.isConfigured()).thenReturn(true);
|
||||||
|
when(rustfs.uploadBytes(anyString(), any(byte[].class), anyBoolean()))
|
||||||
|
.thenAnswer(invocation -> invocation.getArgument(0));
|
||||||
|
TransientPayloadStorageService service = newService(rustfs, propertiesWithDecompressionLimit(BIG_LIMIT));
|
||||||
|
String content = "{\"row\":\"single\"}";
|
||||||
|
|
||||||
|
String pointer = service.storeResultItemPayload("TEST", 7L, "scope", "row-1", content);
|
||||||
|
when(rustfs.readObjectBytes(anyString())).thenReturn(gzipBytes(content));
|
||||||
|
|
||||||
|
assertEquals(content, service.resolvePayload(pointer, "read failed"));
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void test_task_065_payload_boundary_limit_and_overflow() throws Exception {
|
||||||
|
// 上限/超限:解压后字节数超过 maxDecompressedPayloadBytes 时中止并报错(防压缩炸弹),
|
||||||
|
// 等于上限时恰好成功;不产生无界内存增长。
|
||||||
|
RustfsObjectStorageService rustfs = mock(RustfsObjectStorageService.class);
|
||||||
|
when(rustfs.isConfigured()).thenReturn(true);
|
||||||
|
when(rustfs.uploadBytes(anyString(), any(byte[].class), anyBoolean()))
|
||||||
|
.thenAnswer(invocation -> invocation.getArgument(0));
|
||||||
|
String content = compressiblePayload(2000);
|
||||||
|
int exactLimit = content.getBytes(StandardCharsets.UTF_8).length;
|
||||||
|
byte[] compressed = gzipBytes(content);
|
||||||
|
|
||||||
|
TransientPayloadStorageService exactService =
|
||||||
|
newService(rustfs, propertiesWithDecompressionLimit(exactLimit));
|
||||||
|
String exactPointer = exactService.storeParsedPayloadFast("TEST", 1L, "scope", content, false);
|
||||||
|
when(rustfs.readObjectBytes(anyString())).thenReturn(compressed);
|
||||||
|
assertEquals(content, exactService.resolvePayload(exactPointer, "read failed"));
|
||||||
|
|
||||||
|
TransientPayloadStorageService strictService =
|
||||||
|
newService(rustfs, propertiesWithDecompressionLimit(exactLimit - 1L));
|
||||||
|
String strictPointer = strictService.storeParsedPayloadFast("TEST", 1L, "scope", content, false);
|
||||||
|
IllegalStateException ex = assertThrows(IllegalStateException.class,
|
||||||
|
() -> strictService.resolvePayload(strictPointer, "read failed"));
|
||||||
|
assertTrue(ex.getCause() != null && ex.getCause().getMessage() != null
|
||||||
|
&& ex.getCause().getMessage().contains("exceeds"), "超限错误消息应可识别");
|
||||||
|
assertTrue(ex.getMessage().contains("read failed"), "调用方上下文应保留在顶层消息");
|
||||||
|
|
||||||
|
when(rustfs.readObjectBytes(anyString())).thenReturn(compressed);
|
||||||
|
assertThrows(IllegalStateException.class,
|
||||||
|
() -> strictService.resolvePayload(strictPointer, "read failed"));
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void test_task_065_payload_invalid_input_rejected() throws Exception {
|
||||||
|
// 非法输入:损坏的 gzip 流(magic 头 + 垃圾字节)读回时抛 IllegalStateException 且错误消息可识别。
|
||||||
|
RustfsObjectStorageService rustfs = mock(RustfsObjectStorageService.class);
|
||||||
|
when(rustfs.isConfigured()).thenReturn(true);
|
||||||
|
when(rustfs.uploadBytes(anyString(), any(byte[].class), anyBoolean()))
|
||||||
|
.thenAnswer(invocation -> invocation.getArgument(0));
|
||||||
|
TransientPayloadStorageService service = newService(rustfs, propertiesWithDecompressionLimit(BIG_LIMIT));
|
||||||
|
String pointer = service.storeParsedPayloadFast("TEST", 1L, "scope", "{\"x\":1}", false);
|
||||||
|
when(rustfs.readObjectBytes(anyString()))
|
||||||
|
.thenReturn(new byte[]{0x1f, (byte) 0x8b, 0x00, (byte) 0xFF, 0x00, 0x00});
|
||||||
|
|
||||||
|
IllegalStateException ex = assertThrows(IllegalStateException.class,
|
||||||
|
() -> service.resolvePayload(pointer, "read failed"));
|
||||||
|
assertTrue(ex.getMessage().contains("read failed"), "错误消息应包含调用方上下文");
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void test_task_065_payload_dependency_failure_releases_resources() throws Exception {
|
||||||
|
// 依赖失败:对象存储读取失败传播且不残留;恢复后再次读取成功。
|
||||||
|
RustfsObjectStorageService rustfs = mock(RustfsObjectStorageService.class);
|
||||||
|
when(rustfs.isConfigured()).thenReturn(true);
|
||||||
|
when(rustfs.uploadBytes(anyString(), any(byte[].class), anyBoolean()))
|
||||||
|
.thenAnswer(invocation -> invocation.getArgument(0));
|
||||||
|
TransientPayloadStorageService service = newService(rustfs, propertiesWithDecompressionLimit(BIG_LIMIT));
|
||||||
|
String content = compressiblePayload(800);
|
||||||
|
String pointer = service.storeParsedPayloadFast("TEST", 1L, "scope", content, false);
|
||||||
|
|
||||||
|
when(rustfs.readObjectBytes(anyString())).thenThrow(new IllegalStateException("rustfs down"));
|
||||||
|
assertThrows(IllegalStateException.class, () -> service.resolvePayload(pointer, "read failed"));
|
||||||
|
|
||||||
|
// thenThrow 后必须用 doReturn 重新打桩:when() 内的方法调用会立即触发旧异常桩。
|
||||||
|
doReturn(gzipBytes(content)).when(rustfs).readObjectBytes(anyString());
|
||||||
|
assertEquals(content, service.resolvePayload(pointer, "read failed"));
|
||||||
|
verify(rustfs, times(2)).readObjectBytes(anyString());
|
||||||
|
}
|
||||||
|
|
||||||
|
private TransientPayloadStorageService newService(RustfsObjectStorageService rustfs,
|
||||||
|
TransientStorageProperties transientProperties) {
|
||||||
|
StorageProperties storageProperties = new StorageProperties();
|
||||||
|
storageProperties.setLocalTempDir(tempDir.toString());
|
||||||
|
return new TransientPayloadStorageService(
|
||||||
|
transientProperties,
|
||||||
|
storageProperties,
|
||||||
|
rustfs,
|
||||||
|
mock(OssStorageService.class),
|
||||||
|
new ObjectMapper(),
|
||||||
|
new InstanceMetadata("test-instance"),
|
||||||
|
mock(TaskChunkMapper.class),
|
||||||
|
mock(TaskScopeStateMapper.class));
|
||||||
|
}
|
||||||
|
|
||||||
|
private TransientStorageProperties propertiesWithDecompressionLimit(long maxDecompressedPayloadBytes) {
|
||||||
|
TransientStorageProperties properties = new TransientStorageProperties();
|
||||||
|
properties.setEnabled(true);
|
||||||
|
properties.setMaxPayloadBytes(100_000L);
|
||||||
|
properties.setMaxStoredPayloadBytes(100_000L);
|
||||||
|
properties.setWarnPayloadBytes(1);
|
||||||
|
properties.setMaxDecompressedPayloadBytes(maxDecompressedPayloadBytes);
|
||||||
|
properties.setFallbackToLocalOnOversize(true);
|
||||||
|
return properties;
|
||||||
|
}
|
||||||
|
|
||||||
|
private String compressiblePayload(int repeatCount) {
|
||||||
|
return "{\"payload\":\"c".repeat(repeatCount) + "\"}";
|
||||||
|
}
|
||||||
|
|
||||||
|
private byte[] gzipBytes(String content) throws Exception {
|
||||||
|
ByteArrayOutputStream baos = new ByteArrayOutputStream();
|
||||||
|
try (GZIPOutputStream gzip = new GZIPOutputStream(baos)) {
|
||||||
|
gzip.write(content.getBytes(StandardCharsets.UTF_8));
|
||||||
|
}
|
||||||
|
return baos.toByteArray();
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user