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:
2026-08-30 18:49:55 +08:00
parent 4e9c5d242e
commit 9b1978aec8
3 changed files with 280 additions and 8 deletions
@@ -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 65transient 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();
}
}