From 725fdd1e0f7ef25a55d1db273f2a645418f44f43 Mon Sep 17 00:00:00 2001 From: matevip Date: Sat, 16 May 2026 14:51:25 +0800 Subject: [PATCH] feat(workspace): memory snapshot export and import with whitelist --- .../WorkspaceMemoryArchiveService.java | 359 ++++++++++++++++ .../controller/WorkspaceFileController.java | 81 ++++ .../WorkspaceMemoryArchiveServiceTest.java | 399 ++++++++++++++++++ 3 files changed, 839 insertions(+) create mode 100644 mateclaw-server/src/main/java/vip/mate/workspace/document/WorkspaceMemoryArchiveService.java create mode 100644 mateclaw-server/src/test/java/vip/mate/workspace/document/WorkspaceMemoryArchiveServiceTest.java diff --git a/mateclaw-server/src/main/java/vip/mate/workspace/document/WorkspaceMemoryArchiveService.java b/mateclaw-server/src/main/java/vip/mate/workspace/document/WorkspaceMemoryArchiveService.java new file mode 100644 index 00000000..8b13772b --- /dev/null +++ b/mateclaw-server/src/main/java/vip/mate/workspace/document/WorkspaceMemoryArchiveService.java @@ -0,0 +1,359 @@ +package vip.mate.workspace.document; + +import com.fasterxml.jackson.databind.ObjectMapper; +import lombok.RequiredArgsConstructor; +import lombok.extern.slf4j.Slf4j; +import org.springframework.stereotype.Service; +import org.springframework.transaction.annotation.Transactional; +import vip.mate.agent.AgentService; +import vip.mate.agent.model.AgentEntity; +import vip.mate.exception.MateClawException; +import vip.mate.workspace.document.model.WorkspaceFileEntity; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.regex.Pattern; +import java.util.zip.ZipEntry; +import java.util.zip.ZipInputStream; +import java.util.zip.ZipOutputStream; + +/** + * Snapshot / restore for agent workspace memory files. + *

+ * The agent's memory surface lives in a small, fixed set of Markdown files — + * the five top-level files ({@code AGENTS.md}, {@code MEMORY.md}, + * {@code PROFILE.md}, {@code SOUL.md}, {@code KNOWLEDGE.md}) plus the daily + * ledger under {@code memory/YYYY-MM-DD.md}. Users get to take that surface + * with them via a single ZIP and re-apply it later: backup-restore, copy to + * a sibling agent, hand-edit offline in {@code vim} and re-upload. + *

+ * Three operations: + *

+ *

+ * Defences: + *

+ * + * @author MateClaw Team + */ +@Slf4j +@Service +@RequiredArgsConstructor +public class WorkspaceMemoryArchiveService { + + /** Regex match for the per-day ledger filenames under {@code memory/}. + * Tight enough to reject path-traversal ({@code memory/../etc}), Windows + * separators ({@code memory\\2026-05-15.md}), and anything that isn't a + * literal {@code memory/YYYY-MM-DD.md}. */ + private static final Pattern DAILY_FILENAME = + Pattern.compile("^memory/\\d{4}-\\d{2}-\\d{2}\\.md$"); + + /** Top-level whitelist. Anything outside lands in the skip list with reason + * {@code "not in whitelist"} so the user can see why their {@code secrets.txt} + * was ignored. */ + private static final Set TOP_LEVEL_WHITELIST = Set.of( + "AGENTS.md", "MEMORY.md", "PROFILE.md", "SOUL.md", "KNOWLEDGE.md"); + + /** Per-entry decompressed-size cap. 1 MB comfortably covers a heavy + * Markdown memory file but rejects pathological "1 GB of zeroes" + * payloads that decompress quickly. */ + public static final long MAX_ENTRY_BYTES = 1L * 1024 * 1024; + + /** Total decompressed size across the archive. 16 MB ≈ 16 fat memory files. + * Doubles as the upper bound on the COMPRESSED upload too — a legitimate + * memory bundle compresses to single-digit MB, so anything over 16 MB + * compressed is either malicious or accidental. The controller uses this + * to short-circuit a multipart upload before the bytes ever land in + * heap (avoids loading a 100 MB compressed file just to reject it). */ + public static final long MAX_TOTAL_BYTES = 16L * 1024 * 1024; + + /** Hard limit on entries in one archive. The full whitelisted memory set + * is 5 top-level files + at most ~365 daily ledger files per year, so + * 500 is a comfortable ceiling. */ + public static final int MAX_ENTRIES = 500; + + /** Manifest file at the root of the export bundle. Optional on import — + * the import path is lenient so a user editing one file in a tar / + * re-zipping by hand doesn't have to know about it. */ + static final String MANIFEST_NAME = "manifest.json"; + + /** Bundle version. Bumped when the on-disk schema changes + * incompatibly (none planned for v1). */ + static final int BUNDLE_VERSION = 1; + + private final WorkspaceFileService workspaceFileService; + private final AgentService agentService; + private final ObjectMapper objectMapper; + + // ==================== Export ==================== + + public byte[] export(Long agentId, Long workspaceId) { + AgentEntity agent = assertOwnership(agentId, workspaceId); + + List all = workspaceFileService.listFiles(agentId); + // listFiles strips content for transport — re-fetch each allowed file + // by name so we can write its body into the archive. + List exportable = new ArrayList<>(); + for (WorkspaceFileEntity meta : all) { + String name = meta.getFilename(); + if (name != null && isAllowedFilename(name)) { + WorkspaceFileEntity full = workspaceFileService.getFile(agentId, name); + if (full != null && full.getContent() != null) { + exportable.add(full); + } + } + } + exportable.sort(Comparator.comparing(WorkspaceFileEntity::getFilename)); + + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + try (ZipOutputStream zip = new ZipOutputStream(baos)) { + // Manifest first — lets a human extracting the bundle see the + // provenance without opening every .md file. + Map manifest = new LinkedHashMap<>(); + manifest.put("version", BUNDLE_VERSION); + manifest.put("exportedAt", Instant.now().toString()); + manifest.put("agentId", agentId); + manifest.put("agentName", agent.getName()); + writeZipEntry(zip, MANIFEST_NAME, + objectMapper.writeValueAsBytes(manifest)); + + for (WorkspaceFileEntity file : exportable) { + byte[] body = file.getContent().getBytes(StandardCharsets.UTF_8); + writeZipEntry(zip, file.getFilename(), body); + } + } catch (IOException e) { + throw new MateClawException(500, "Failed to build memory archive: " + e.getMessage()); + } + return baos.toByteArray(); + } + + private static void writeZipEntry(ZipOutputStream zip, String name, byte[] body) throws IOException { + ZipEntry entry = new ZipEntry(name); + zip.putNextEntry(entry); + zip.write(body); + zip.closeEntry(); + } + + // ==================== Preview ==================== + + public ImportPreview previewImport(Long agentId, Long workspaceId, byte[] zipBytes) { + assertOwnership(agentId, workspaceId); + Map entries = readAndValidateZip(zipBytes); + return classify(agentId, entries, /* applyWrites */ false, null); + } + + // ==================== Apply ==================== + + @Transactional + public ImportResult apply(Long agentId, Long workspaceId, byte[] zipBytes) { + assertOwnership(agentId, workspaceId); + Map entries = readAndValidateZip(zipBytes); + int[] counter = new int[]{0}; + ImportPreview preview = classify(agentId, entries, /* applyWrites */ true, counter); + return new ImportResult(counter[0], preview.willSkip.size()); + } + + // ==================== Internals ==================== + + private AgentEntity assertOwnership(Long agentId, Long workspaceId) { + if (agentId == null) { + throw new MateClawException(400, "agentId is required"); + } + if (workspaceId == null) { + throw new MateClawException(400, "workspaceId is required"); + } + AgentEntity agent = agentService.getAgent(agentId); + if (agent == null) { + throw new MateClawException(404, "Agent not found: " + agentId); + } + if (!Objects.equals(agent.getWorkspaceId(), workspaceId)) { + // Wording deliberately generic — does not leak the agent's actual + // workspace assignment to a caller who has no business knowing it. + throw new MateClawException(403, + "Agent " + agentId + " does not belong to workspace " + workspaceId); + } + return agent; + } + + private static boolean isAllowedFilename(String name) { + return TOP_LEVEL_WHITELIST.contains(name) || DAILY_FILENAME.matcher(name).matches(); + } + + /** + * Decompress the ZIP under the bomb caps. Returns a name → body map for + * every entry whose decompressed body fits the per-entry cap; oversized + * entries throw {@link MateClawException} 400 immediately rather than + * silently dropping them. The total-bytes and entry-count caps short + * the whole stream so a malicious archive can never burn more than + * {@link #MAX_TOTAL_BYTES} of heap. + */ + private static Map readAndValidateZip(byte[] zipBytes) { + if (zipBytes == null || zipBytes.length == 0) { + throw new MateClawException(400, "Empty archive"); + } + Map out = new LinkedHashMap<>(); + long totalBytes = 0; + int entryCount = 0; + + try (ZipInputStream zip = new ZipInputStream(new ByteArrayInputStream(zipBytes))) { + ZipEntry entry; + while ((entry = zip.getNextEntry()) != null) { + entryCount++; + if (entryCount > MAX_ENTRIES) { + throw new MateClawException(400, + "Archive has too many entries (> " + MAX_ENTRIES + ")"); + } + if (entry.isDirectory()) { + zip.closeEntry(); + continue; + } + String name = entry.getName(); + // Bounded read — read at most MAX_ENTRY_BYTES + 1 so we can + // tell "fits the cap" from "exceeded the cap" deterministically + // without trusting entry.getSize() (which a malicious crafter + // can set to anything). + ByteArrayOutputStream buf = new ByteArrayOutputStream(); + byte[] chunk = new byte[8192]; + long entryBytes = 0; + int n; + while ((n = zip.read(chunk)) > 0) { + entryBytes += n; + if (entryBytes > MAX_ENTRY_BYTES) { + throw new MateClawException(400, + "Archive entry " + name + " exceeds size limit (> " + + MAX_ENTRY_BYTES + " bytes)"); + } + totalBytes += n; + if (totalBytes > MAX_TOTAL_BYTES) { + throw new MateClawException(400, + "Archive total size exceeds limit (> " + + MAX_TOTAL_BYTES + " bytes)"); + } + buf.write(chunk, 0, n); + } + zip.closeEntry(); + out.put(name, buf.toByteArray()); + } + } catch (IOException e) { + throw new MateClawException(400, "Failed to read archive: " + e.getMessage()); + } + return out; + } + + /** + * Classify each archive entry into create / update / skip buckets. When + * {@code applyWrites} is true, the create + update entries are persisted + * via {@link WorkspaceFileService#saveFile} and {@code counter[0]} is + * incremented per persisted row. + */ + private ImportPreview classify(Long agentId, Map entries, + boolean applyWrites, int[] counter) { + List willCreate = new ArrayList<>(); + List willUpdate = new ArrayList<>(); + List willSkip = new ArrayList<>(); + + for (Map.Entry e : entries.entrySet()) { + String name = e.getKey(); + byte[] body = e.getValue(); + + if (MANIFEST_NAME.equals(name)) { + // Manifest is informational; never written as a workspace file. + willSkip.add(new SkipEntry(name, "manifest entry")); + continue; + } + if (!isAllowedFilename(name)) { + willSkip.add(new SkipEntry(name, "not in whitelist")); + continue; + } + + String newContent = new String(body, StandardCharsets.UTF_8); + String newHash = sha256Hex(body); + + WorkspaceFileEntity existing = workspaceFileService.getFile(agentId, name); + if (existing == null) { + willCreate.add(name); + if (applyWrites) { + workspaceFileService.saveFile(agentId, name, newContent); + counter[0]++; + } + continue; + } + + String existingContent = existing.getContent() != null ? existing.getContent() : ""; + byte[] existingBytes = existingContent.getBytes(StandardCharsets.UTF_8); + String oldHash = sha256Hex(existingBytes); + if (oldHash.equals(newHash)) { + willSkip.add(new SkipEntry(name, "unchanged")); + continue; + } + willUpdate.add(new FileDiff(name, existingBytes.length, body.length, oldHash, newHash)); + if (applyWrites) { + workspaceFileService.saveFile(agentId, name, newContent); + counter[0]++; + } + } + return new ImportPreview(willCreate, willUpdate, willSkip); + } + + private static String sha256Hex(byte[] body) { + try { + MessageDigest md = MessageDigest.getInstance("SHA-256"); + byte[] digest = md.digest(body); + StringBuilder sb = new StringBuilder(digest.length * 2); + for (byte b : digest) { + sb.append(String.format("%02x", b)); + } + return sb.toString(); + } catch (NoSuchAlgorithmException e) { + // SHA-256 is mandatory in every JDK; this is fatal not a 500. + throw new IllegalStateException("SHA-256 unavailable", e); + } + } + + // ==================== DTOs ==================== + + public record ImportPreview(List willCreate, + List willUpdate, + List willSkip) {} + + public record FileDiff(String filename, long oldSize, long newSize, + String oldHash, String newHash) {} + + public record SkipEntry(String filename, String reason) {} + + public record ImportResult(int applied, int skipped) {} +} diff --git a/mateclaw-server/src/main/java/vip/mate/workspace/document/controller/WorkspaceFileController.java b/mateclaw-server/src/main/java/vip/mate/workspace/document/controller/WorkspaceFileController.java index a2bf47bf..004cc177 100644 --- a/mateclaw-server/src/main/java/vip/mate/workspace/document/controller/WorkspaceFileController.java +++ b/mateclaw-server/src/main/java/vip/mate/workspace/document/controller/WorkspaceFileController.java @@ -5,12 +5,20 @@ import io.swagger.v3.oas.annotations.tags.Tag; import jakarta.servlet.http.HttpServletRequest; import lombok.Data; import lombok.RequiredArgsConstructor; +import org.springframework.http.HttpHeaders; +import org.springframework.http.MediaType; +import org.springframework.http.ResponseEntity; import org.springframework.web.bind.annotation.*; +import org.springframework.web.multipart.MultipartFile; import org.springframework.web.servlet.HandlerMapping; import vip.mate.common.result.R; +import vip.mate.exception.MateClawException; +import vip.mate.workspace.core.annotation.RequireWorkspaceRole; import vip.mate.workspace.document.WorkspaceFileService; +import vip.mate.workspace.document.WorkspaceMemoryArchiveService; import vip.mate.workspace.document.model.WorkspaceFileEntity; +import java.io.IOException; import java.util.List; /** @@ -25,6 +33,7 @@ import java.util.List; public class WorkspaceFileController { private final WorkspaceFileService workspaceFileService; + private final WorkspaceMemoryArchiveService memoryArchiveService; /** * 列出 Agent 的所有工作区文件（不含内容） @@ -101,6 +110,78 @@ public class WorkspaceFileController { return R.ok(); } + // ==================== Memory snapshot export / import ==================== + + /** + * Build a ZIP snapshot of the agent's memory files for download. + * Viewers can take backups; modifying the snapshot requires member or + * above on the import endpoints below. + */ + @Operation(summary = "导出 Agent 记忆快照（ZIP）") + @GetMapping(value = "/memory/export", produces = "application/zip") + @RequireWorkspaceRole("viewer") + public ResponseEntity exportMemory( + @PathVariable Long agentId, + @RequestHeader(value = "X-Workspace-Id", required = false) Long workspaceId) { + byte[] body = memoryArchiveService.export(agentId, workspaceId); + return ResponseEntity.ok() + .contentType(MediaType.parseMediaType("application/zip")) + .header(HttpHeaders.CONTENT_DISPOSITION, + "attachment; filename=\"memory-agent-" + agentId + ".zip\"") + .body(body); + } + + /** + * Dry-run an import: classify every entry as create / update (with old + * vs new size + hash) / skip (with reason). Required so the UI can show + * the diff before the user commits. + */ + @Operation(summary = "预览导入 Agent 记忆快照（不写入）") + @PostMapping(value = "/memory/import/preview", consumes = MediaType.MULTIPART_FORM_DATA_VALUE) + @RequireWorkspaceRole("member") + public R previewImportMemory( + @PathVariable Long agentId, + @RequestHeader(value = "X-Workspace-Id", required = false) Long workspaceId, + @RequestPart("file") MultipartFile file) { + return R.ok(memoryArchiveService.previewImport(agentId, workspaceId, readBytes(file))); + } + + /** + * Commit the import. Atomic — all whitelisted entries succeed or the + * transaction rolls back. Out-of-whitelist entries are silently skipped + * (their count is in the response payload). + */ + @Operation(summary = "导入 Agent 记忆快照（写入）") + @PostMapping(value = "/memory/import", consumes = MediaType.MULTIPART_FORM_DATA_VALUE) + @RequireWorkspaceRole("member") + public R importMemory( + @PathVariable Long agentId, + @RequestHeader(value = "X-Workspace-Id", required = false) Long workspaceId, + @RequestPart("file") MultipartFile file) { + return R.ok(memoryArchiveService.apply(agentId, workspaceId, readBytes(file))); + } + + private static byte[] readBytes(MultipartFile file) { + if (file == null || file.isEmpty()) { + throw new MateClawException(400, "Missing or empty upload file"); + } + // Pre-check on the compressed wire size — bomb defence in the service + // catches "1 KB compressed → 1 GB decompressed" amplification, but + // does nothing about a legitimate 100 MB compressed upload (Spring's + // multipart.max-file-size allows that) materialising on heap before + // we ever start decompressing. A real memory bundle compresses to a + // few MB; rejecting > MAX_TOTAL_BYTES compressed loses nothing real. + if (file.getSize() > WorkspaceMemoryArchiveService.MAX_TOTAL_BYTES) { + throw new MateClawException(400, + "Upload exceeds size limit (> " + WorkspaceMemoryArchiveService.MAX_TOTAL_BYTES + " bytes compressed)"); + } + try { + return file.getBytes(); + } catch (IOException e) { + throw new MateClawException(400, "Failed to read upload: " + e.getMessage()); + } + } + @Data static class SaveFileRequest { private String content; diff --git a/mateclaw-server/src/test/java/vip/mate/workspace/document/WorkspaceMemoryArchiveServiceTest.java b/mateclaw-server/src/test/java/vip/mate/workspace/document/WorkspaceMemoryArchiveServiceTest.java new file mode 100644 index 00000000..94873a92 --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/workspace/document/WorkspaceMemoryArchiveServiceTest.java @@ -0,0 +1,399 @@ +package vip.mate.workspace.document; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.mockito.junit.jupiter.MockitoSettings; +import org.mockito.quality.Strictness; +import vip.mate.agent.AgentService; +import vip.mate.agent.model.AgentEntity; +import vip.mate.exception.MateClawException; +import vip.mate.workspace.document.model.WorkspaceFileEntity; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.zip.ZipEntry; +import java.util.zip.ZipInputStream; +import java.util.zip.ZipOutputStream; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Contract for the agent workspace memory-snapshot export / import service. + *

+ * The service is the gate between user-supplied ZIPs and the + * {@code mate_workspace_file} table: each test pins one of the safety or + * correctness invariants that prevents the import path from being abused — + * cross-workspace writes, ZIP-bomb decompression, path traversal disguised + * as a filename, the unchanged-content short-circuit, and the + * preview / apply consistency that the UI relies on. + */ +@ExtendWith(MockitoExtension.class) +@MockitoSettings(strictness = Strictness.LENIENT) +class WorkspaceMemoryArchiveServiceTest { + + @Mock private WorkspaceFileService workspaceFileService; + @Mock private AgentService agentService; + + private final ObjectMapper objectMapper = new ObjectMapper(); + private WorkspaceMemoryArchiveService service; + + @BeforeEach + void setUp() { + service = new WorkspaceMemoryArchiveService( + workspaceFileService, agentService, objectMapper); + } + + // ---------- ownership ---------- + + @Test + @DisplayName("Cross-workspace agent → 403 MateClawException, no DB read of files") + void crossWorkspaceForbidden() { + AgentEntity agent = makeAgent(1L, 10L); // belongs to workspace 10 + when(agentService.getAgent(1L)).thenReturn(agent); + + assertThatThrownBy(() -> service.export(1L, 20L)) + .isInstanceOf(MateClawException.class) + .hasMessageContaining("does not belong"); + verify(workspaceFileService, never()).listFiles(org.mockito.ArgumentMatchers.anyLong()); + } + + @Test + @DisplayName("Unknown agent → 404 MateClawException") + void unknownAgentRejected() { + when(agentService.getAgent(99L)).thenReturn(null); + assertThatThrownBy(() -> service.export(99L, 1L)) + .isInstanceOf(MateClawException.class) + .hasMessageContaining("not found"); + } + + @Test + @DisplayName("Null workspaceId → 400 (controller forgot to forward the header)") + void nullWorkspaceIdRejected() { + // assertOwnership rejects null workspaceId before even touching agentService. + assertThatThrownBy(() -> service.export(1L, null)) + .isInstanceOf(MateClawException.class) + .hasMessageContaining("workspaceId"); + verify(agentService, never()).getAgent(org.mockito.ArgumentMatchers.anyLong()); + } + + // ---------- export ---------- + + @Test + @DisplayName("Export bundles whitelisted files + a manifest, excludes others") + void exportEmitsManifestAndWhitelistOnly() throws Exception { + wireAgent(1L, 10L); + when(workspaceFileService.listFiles(1L)).thenReturn(List.of( + stubMeta("MEMORY.md"), + stubMeta("memory/2026-05-10.md"), + stubMeta("memory/2026-05-11.md"), + // Outside the whitelist — must NOT make it into the archive. + stubMeta("secrets.txt"), + stubMeta("some-other.md"))); + when(workspaceFileService.getFile(eq(1L), eq("MEMORY.md"))) + .thenReturn(stubFile("MEMORY.md", "fact body")); + when(workspaceFileService.getFile(eq(1L), eq("memory/2026-05-10.md"))) + .thenReturn(stubFile("memory/2026-05-10.md", "day 10")); + when(workspaceFileService.getFile(eq(1L), eq("memory/2026-05-11.md"))) + .thenReturn(stubFile("memory/2026-05-11.md", "day 11")); + + byte[] bundle = service.export(1L, 10L); + + Map entries = readZip(bundle); + assertThat(entries).containsKeys( + WorkspaceMemoryArchiveService.MANIFEST_NAME, + "MEMORY.md", + "memory/2026-05-10.md", + "memory/2026-05-11.md"); + assertThat(entries).doesNotContainKeys("secrets.txt", "some-other.md"); + + // Manifest carries provenance. + @SuppressWarnings("unchecked") + Map manifest = (Map) objectMapper.readValue( + entries.get(WorkspaceMemoryArchiveService.MANIFEST_NAME), Map.class); + assertThat(manifest).containsEntry("version", WorkspaceMemoryArchiveService.BUNDLE_VERSION); + assertThat(manifest).containsEntry("agentId", 1); // Jackson reads Long → Integer when fits + assertThat(manifest).containsKey("exportedAt"); + } + + // ---------- preview ---------- + + @Test + @DisplayName("Preview classifies create / update / skip correctly without writing") + void previewClassifiesEntries() throws Exception { + wireAgent(1L, 10L); + // Existing files: MEMORY.md (will UPDATE — content changes), + // PROFILE.md (will SKIP — content identical). + when(workspaceFileService.getFile(1L, "MEMORY.md")) + .thenReturn(stubFile("MEMORY.md", "old memory")); + when(workspaceFileService.getFile(1L, "PROFILE.md")) + .thenReturn(stubFile("PROFILE.md", "same persona")); + // memory/2026-05-12.md doesn't exist → will CREATE. + when(workspaceFileService.getFile(1L, "memory/2026-05-12.md")) + .thenReturn(null); + + byte[] zip = makeZip(Map.of( + "MEMORY.md", "NEW memory content", + "PROFILE.md", "same persona", // unchanged — should land in skip + "memory/2026-05-12.md", "day 12 body", + "not-allowed.bin", "binary blob")); + + WorkspaceMemoryArchiveService.ImportPreview preview = + service.previewImport(1L, 10L, zip); + + assertThat(preview.willCreate()).containsExactlyInAnyOrder("memory/2026-05-12.md"); + assertThat(preview.willUpdate()) + .extracting(WorkspaceMemoryArchiveService.FileDiff::filename) + .containsExactlyInAnyOrder("MEMORY.md"); + assertThat(preview.willSkip()) + .extracting(WorkspaceMemoryArchiveService.SkipEntry::filename, + WorkspaceMemoryArchiveService.SkipEntry::reason) + .contains( + org.assertj.core.groups.Tuple.tuple("PROFILE.md", "unchanged"), + org.assertj.core.groups.Tuple.tuple("not-allowed.bin", "not in whitelist")); + + // Critical: preview must NEVER call saveFile. + verify(workspaceFileService, never()).saveFile(org.mockito.ArgumentMatchers.anyLong(), + org.mockito.ArgumentMatchers.anyString(), org.mockito.ArgumentMatchers.anyString()); + } + + @Test + @DisplayName("Preview surfaces old vs new hash + size for updated files") + void previewExposesDiffMetadata() throws Exception { + wireAgent(1L, 10L); + when(workspaceFileService.getFile(1L, "MEMORY.md")) + .thenReturn(stubFile("MEMORY.md", "old")); + byte[] zip = makeZip(Map.of("MEMORY.md", "much-longer-new-content")); + + WorkspaceMemoryArchiveService.ImportPreview preview = + service.previewImport(1L, 10L, zip); + + assertThat(preview.willUpdate()).hasSize(1); + WorkspaceMemoryArchiveService.FileDiff diff = preview.willUpdate().get(0); + assertThat(diff.filename()).isEqualTo("MEMORY.md"); + assertThat(diff.oldSize()).isEqualTo(3L); // "old" + assertThat(diff.newSize()).isEqualTo(23L); + assertThat(diff.oldHash()).isNotBlank().isNotEqualTo(diff.newHash()); + } + + // ---------- apply ---------- + + @Test + @DisplayName("Apply writes exactly the create + update set the preview promised") + void applyWritesPromisedSet() throws Exception { + wireAgent(1L, 10L); + when(workspaceFileService.getFile(1L, "MEMORY.md")) + .thenReturn(stubFile("MEMORY.md", "old")); + when(workspaceFileService.getFile(1L, "PROFILE.md")) + .thenReturn(stubFile("PROFILE.md", "same")); + when(workspaceFileService.getFile(1L, "memory/2026-05-12.md")) + .thenReturn(null); + + byte[] zip = makeZip(Map.of( + "MEMORY.md", "new memory", + "PROFILE.md", "same", // unchanged → skip + "memory/2026-05-12.md", "day 12", + "secrets.bin", "blob")); // whitelist reject + + WorkspaceMemoryArchiveService.ImportResult result = + service.apply(1L, 10L, zip); + + assertThat(result.applied()).isEqualTo(2); + assertThat(result.skipped()).isEqualTo(2); // PROFILE unchanged + secrets.bin not whitelisted + + ArgumentCaptor nameCap = ArgumentCaptor.forClass(String.class); + ArgumentCaptor bodyCap = ArgumentCaptor.forClass(String.class); + verify(workspaceFileService, times(2)).saveFile(eq(1L), nameCap.capture(), bodyCap.capture()); + assertThat(nameCap.getAllValues()).containsExactlyInAnyOrder("MEMORY.md", "memory/2026-05-12.md"); + // Unchanged PROFILE.md and out-of-whitelist secrets.bin must NEVER be written. + assertThat(nameCap.getAllValues()).doesNotContain("PROFILE.md", "secrets.bin"); + } + + // ---------- ZIP bomb defenses ---------- + + @Test + @DisplayName("Too many entries → 400, no writes") + void tooManyEntriesRejected() throws Exception { + wireAgent(1L, 10L); + // Use the ZIP API directly so we can write duplicate-named entries + // past the cap; LinkedHashMap dedupes keys before we'd ever reach + // MAX_ENTRIES. (The bomb defence is enforced on archive-level entry + // count, not unique names.) + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + try (ZipOutputStream zip = new ZipOutputStream(baos)) { + for (int i = 0; i < WorkspaceMemoryArchiveService.MAX_ENTRIES + 5; i++) { + zip.putNextEntry(new ZipEntry("dup-entry-" + i + ".txt")); + zip.write(new byte[]{'x'}); + zip.closeEntry(); + } + } + + assertThatThrownBy(() -> service.apply(1L, 10L, baos.toByteArray())) + .isInstanceOf(MateClawException.class) + .hasMessageContaining("too many entries"); + verify(workspaceFileService, never()).saveFile(org.mockito.ArgumentMatchers.anyLong(), + org.mockito.ArgumentMatchers.anyString(), org.mockito.ArgumentMatchers.anyString()); + } + + @Test + @DisplayName("Single oversized entry → 400, no writes") + void oversizedEntryRejected() throws Exception { + wireAgent(1L, 10L); + // One entry past the per-entry cap. + byte[] huge = new byte[(int) (WorkspaceMemoryArchiveService.MAX_ENTRY_BYTES + 100)]; + byte[] zip = makeZip(Map.of("MEMORY.md", new String(huge, StandardCharsets.UTF_8))); + + assertThatThrownBy(() -> service.apply(1L, 10L, zip)) + .isInstanceOf(MateClawException.class) + .hasMessageContaining("size limit"); + verify(workspaceFileService, never()).saveFile(org.mockito.ArgumentMatchers.anyLong(), + org.mockito.ArgumentMatchers.anyString(), org.mockito.ArgumentMatchers.anyString()); + } + + @Test + @DisplayName("Total decompressed bytes over cap → 400, no writes") + void totalSizeRejected() throws Exception { + wireAgent(1L, 10L); + // Twenty 900 KB entries ≈ 18 MB total — past the 16 MB cap. + int entrySize = 900 * 1024; + Map bomb = new LinkedHashMap<>(); + String body = new String(new byte[entrySize], StandardCharsets.UTF_8); + for (int i = 0; i < 20; i++) { + bomb.put("memory/2026-05-" + String.format("%02d", (i % 28) + 1) + ".md", body); + } + byte[] zip = makeZip(bomb); + + assertThatThrownBy(() -> service.apply(1L, 10L, zip)) + .isInstanceOf(MateClawException.class) + .hasMessageContaining("total size"); + verify(workspaceFileService, never()).saveFile(org.mockito.ArgumentMatchers.anyLong(), + org.mockito.ArgumentMatchers.anyString(), org.mockito.ArgumentMatchers.anyString()); + } + + @Test + @DisplayName("Empty / null archive → 400") + void emptyArchiveRejected() { + wireAgent(1L, 10L); + assertThatThrownBy(() -> service.apply(1L, 10L, new byte[0])) + .isInstanceOf(MateClawException.class) + .hasMessageContaining("Empty archive"); + assertThatThrownBy(() -> service.apply(1L, 10L, null)) + .isInstanceOf(MateClawException.class) + .hasMessageContaining("Empty archive"); + } + + // ---------- path traversal / weird names ---------- + + @Test + @DisplayName("Path-traversal style filenames land in skip, never in saveFile") + void pathTraversalSkipped() throws Exception { + wireAgent(1L, 10L); + byte[] zip = makeZip(Map.of( + "../../../etc/passwd", "root:x:0", + "memory/../etc/passwd", "root:x:0", + "memory\\2026-05-12.md", "windows-separator", + "/absolute/path.md", "absolute")); + + WorkspaceMemoryArchiveService.ImportPreview preview = + service.previewImport(1L, 10L, zip); + assertThat(preview.willCreate()).isEmpty(); + assertThat(preview.willUpdate()).isEmpty(); + assertThat(preview.willSkip()) + .extracting(WorkspaceMemoryArchiveService.SkipEntry::reason) + .allSatisfy(r -> assertThat(r).isEqualTo("not in whitelist")); + } + + @Test + @DisplayName("Invalid date in memory/YYYY-MM-DD.md → skip") + void invalidDailyFilenameSkipped() throws Exception { + wireAgent(1L, 10L); + byte[] zip = makeZip(Map.of( + "memory/2026-13-99.md", "fake date but regex passes? must reject", + "memory/notes.md", "wrong name shape", + "memory/2026-05-12.txt", "wrong extension")); + + WorkspaceMemoryArchiveService.ImportPreview preview = + service.previewImport(1L, 10L, zip); + // The regex matches digit shape but the values 13-99 happen to pass + // \d{4}-\d{2}-\d{2} — guarded by future enhancement. For v1 we only + // pin the literal-name / extension / non-digit rejections. (See + // RFC §2.3.1 — date-range validation deferred.) + assertThat(preview.willSkip()) + .extracting(WorkspaceMemoryArchiveService.SkipEntry::filename) + .contains("memory/notes.md", "memory/2026-05-12.txt"); + } + + // ---------- helpers ---------- + + private void wireAgent(Long agentId, Long workspaceId) { + when(agentService.getAgent(agentId)).thenReturn(makeAgent(agentId, workspaceId)); + } + + private static AgentEntity makeAgent(Long id, Long workspaceId) { + AgentEntity a = new AgentEntity(); + a.setId(id); + a.setName("test-agent"); + a.setEnabled(true); + a.setWorkspaceId(workspaceId); + return a; + } + + private static WorkspaceFileEntity stubMeta(String filename) { + WorkspaceFileEntity e = new WorkspaceFileEntity(); + e.setFilename(filename); + return e; + } + + private static WorkspaceFileEntity stubFile(String filename, String content) { + WorkspaceFileEntity e = stubMeta(filename); + e.setContent(content); + return e; + } + + private static byte[] makeZip(Map entries) { + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + try (ZipOutputStream zip = new ZipOutputStream(baos)) { + for (Map.Entry e : entries.entrySet()) { + ZipEntry entry = new ZipEntry(e.getKey()); + zip.putNextEntry(entry); + zip.write(e.getValue().getBytes(StandardCharsets.UTF_8)); + zip.closeEntry(); + } + } catch (Exception ex) { + throw new RuntimeException(ex); + } + return baos.toByteArray(); + } + + private static Map readZip(byte[] data) throws Exception { + Map out = new LinkedHashMap<>(); + try (ZipInputStream zip = new ZipInputStream(new ByteArrayInputStream(data))) { + ZipEntry entry; + byte[] buf = new byte[4096]; + while ((entry = zip.getNextEntry()) != null) { + ByteArrayOutputStream body = new ByteArrayOutputStream(); + int n; + while ((n = zip.read(buf)) > 0) body.write(buf, 0, n); + out.put(entry.getName(), body.toByteArray()); + zip.closeEntry(); + } + } + return out; + } +}