mirror of
https://gitee.com/mateos/mateclaw.git
synced 2026-09-13 03:13:41 +08:00
feat(workspace): memory snapshot export and import with whitelist
This commit is contained in:
matevip
parent
81c6488a3c
commit
725fdd1e0f
@ -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.
|
||||
* <p>
|
||||
* 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.
|
||||
* <p>
|
||||
* Three operations:
|
||||
* <ul>
|
||||
* <li>{@link #export} — Build a ZIP of the agent's whitelisted memory
|
||||
* files plus a {@code manifest.json}. {@code enabled} / {@code sortOrder}
|
||||
* are deliberately NOT serialised — those are UI preferences and an
|
||||
* import must NOT clobber the local toggle state.</li>
|
||||
* <li>{@link #previewImport} — Parse the ZIP without writing anything.
|
||||
* Classifies every entry as create / update (with old/new size + hash)
|
||||
* / skip (with a reason). Required so the UI can show a diff before the
|
||||
* user confirms.</li>
|
||||
* <li>{@link #apply} — Same classification, but actually persist; runs in
|
||||
* a single Spring transaction so an OutOfMemory on entry N-1 doesn't
|
||||
* leave the agent's memory half-rewritten.</li>
|
||||
* </ul>
|
||||
* <p>
|
||||
* Defences:
|
||||
* <ul>
|
||||
* <li>White-list match per entry name (regex + literal set); no path-traversal,
|
||||
* no absolute paths, no arbitrary files. Anything outside lands in the
|
||||
* skip list, never in the file table.</li>
|
||||
* <li>ZIP-bomb caps: ≤ {@link #MAX_ENTRIES} entries, ≤ {@link #MAX_ENTRY_BYTES}
|
||||
* per entry (post-decompression), ≤ {@link #MAX_TOTAL_BYTES} total.</li>
|
||||
* <li>Ownership gate ({@link #assertOwnership}) — the caller-supplied
|
||||
* {@code workspaceId} must match the target agent's row. Controller-level
|
||||
* {@code @RequireWorkspaceRole} guards the HTTP route; this is the
|
||||
* service-layer second line of defence so a unit test or alternative
|
||||
* caller can't bypass it.</li>
|
||||
* </ul>
|
||||
*
|
||||
* @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<String> 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<WorkspaceFileEntity> 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<WorkspaceFileEntity> 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<String, Object> 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<String, byte[]> 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<String, byte[]> 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<String, byte[]> readAndValidateZip(byte[] zipBytes) {
|
||||
if (zipBytes == null || zipBytes.length == 0) {
|
||||
throw new MateClawException(400, "Empty archive");
|
||||
}
|
||||
Map<String, byte[]> 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<String, byte[]> entries,
|
||||
boolean applyWrites, int[] counter) {
|
||||
List<String> willCreate = new ArrayList<>();
|
||||
List<FileDiff> willUpdate = new ArrayList<>();
|
||||
List<SkipEntry> willSkip = new ArrayList<>();
|
||||
|
||||
for (Map.Entry<String, byte[]> 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<String> willCreate,
|
||||
List<FileDiff> willUpdate,
|
||||
List<SkipEntry> 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) {}
|
||||
}
|
||||
@ -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<byte[]> 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<WorkspaceMemoryArchiveService.ImportPreview> 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<WorkspaceMemoryArchiveService.ImportResult> 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;
|
||||
|
||||
@ -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.
|
||||
* <p>
|
||||
* 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<String, byte[]> 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<String, Object> manifest = (Map<String, Object>) 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<String> nameCap = ArgumentCaptor.forClass(String.class);
|
||||
ArgumentCaptor<String> 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<String, String> 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<String, String> entries) {
|
||||
ByteArrayOutputStream baos = new ByteArrayOutputStream();
|
||||
try (ZipOutputStream zip = new ZipOutputStream(baos)) {
|
||||
for (Map.Entry<String, String> 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<String, byte[]> readZip(byte[] data) throws Exception {
|
||||
Map<String, byte[]> 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;
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user