task-65: transient payload 读取流式解压 + 解压后字节上限

- 新增 aiimage.transient-storage.max-decompressed-payload-bytes 配置（默认 100MB，即 2x 存储上限），覆盖 gzip 二进制与 gzip64 兼容两条路径 - decodeGzipStream 按 8KB 缓冲流式解压，累计输出超过上限立即中止并报错，防止压缩炸弹（zip bomb）在读取时无界膨胀内存 - 超限错误保留可识别消息（exceeds configured limit），由 resolvePayload 统一包装
2026-08-30 18:49:55 +08:00
parent 4e9c5d242e
commit 9b1978aec8
3 changed files with 280 additions and 8 deletions
@@ -37,6 +37,11 @@ public class TransientStorageProperties {
    private long warnPayloadBytes = 5L * 1024 * 1024;
    private long maxPayloadBytes = 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 deleteRetryEnabled = true;
    private String deleteRetryCron = "0 */5 * * * *";
@@ -467,19 +467,15 @@ public class TransientPayloadStorageService {
            return "";
        }
        if (isGzip(storedContent)) {
-            try (GZIPInputStream gzip = new GZIPInputStream(new ByteArrayInputStream(storedContent))) {
-                return new String(gzip.readAllBytes(), StandardCharsets.UTF_8);
-            } catch (Exception ex) {
-                throw new IllegalStateException("failed to decode transient payload", ex);
-            }
+            return decodeGzipStream(storedContent);
        }
        String legacyText = new String(storedContent, StandardCharsets.UTF_8);
        if (legacyText.startsWith("gzip64:")) {
            try {
                byte[] compressed = Base64.getDecoder().decode(legacyText.substring("gzip64:".length()));
-                try (GZIPInputStream gzip = new GZIPInputStream(new ByteArrayInputStream(compressed))) {
-                    return new String(gzip.readAllBytes(), StandardCharsets.UTF_8);
-                }
+                return decodeGzipStream(compressed);
+            } catch (IllegalStateException ex) {
+                throw ex;
            } catch (Exception ex) {
                throw new IllegalStateException("failed to decode transient payload", ex);
            }
@@ -488,6 +484,37 @@ public class TransientPayloadStorageService {
        return legacyText;
    }

+    /**
+     * 流式解压 gzip 字节流，解压输出累计超过 maxDecompressedPayloadBytes 立即中止，
+     * 防止压缩炸弹在内存中无界膨胀。上限未配置（&lt;=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) {
        return content.length >= 2
                && (content[0] & 0xFF) == 0x1F
@@ -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();
+    }
+}