From 7ccbb0e2e4cdb888987c0fbbe37739fe47f894f0 Mon Sep 17 00:00:00 2001 From: Nicolas Meienberger Date: Sun, 10 May 2026 22:25:31 +0200 Subject: [PATCH] feat: move restore execution to agent RPC --- app/server/modules/agents/agents-manager.ts | 34 +++- .../modules/agents/controller/server.ts | 36 +++- .../modules/agents/controller/session.ts | 74 +++++-- app/server/modules/backups/backup.helpers.ts | 2 +- .../__tests__/repositories.service.test.ts | 137 +++++++++++-- .../modules/repositories/repositories.dto.ts | 1 + .../repositories/repositories.service.ts | 113 +++++++---- .../commands/{ => __tests__}/volume.test.ts | 8 +- apps/agent/src/commands/backup-run.ts | 2 +- .../__tests__/backup.helpers.test.ts | 0 .../commands/{ => helpers}/backup.helpers.ts | 0 apps/agent/src/commands/index.ts | 4 + apps/agent/src/commands/restore.ts | 73 +++++++ packages/contracts/src/agent-protocol.ts | 49 +++++ .../restic/commands/__tests__/restore.test.ts | 187 +++++++++--------- packages/core/src/restic/commands/restore.ts | 68 ++++--- packages/core/src/restic/index.ts | 1 + 17 files changed, 582 insertions(+), 207 deletions(-) rename apps/agent/src/commands/{ => __tests__}/volume.test.ts (93%) rename apps/agent/src/commands/{ => helpers}/__tests__/backup.helpers.test.ts (100%) rename apps/agent/src/commands/{ => helpers}/backup.helpers.ts (100%) create mode 100644 apps/agent/src/commands/restore.ts diff --git a/app/server/modules/agents/agents-manager.ts b/app/server/modules/agents/agents-manager.ts index 7a5b9c90..f5b6a8a2 100644 --- a/app/server/modules/agents/agents-manager.ts +++ b/app/server/modules/agents/agents-manager.ts @@ -1,7 +1,13 @@ import { logger } from "@zerobyte/core/node"; -import type { BackupRunPayload, VolumeCommand, VolumeCommandResult } from "@zerobyte/contracts/agent-protocol"; +import type { + BackupRunPayload, + RestoreCommandPayload, + VolumeCommand, + VolumeCommandResult, +} from "@zerobyte/contracts/agent-protocol"; import { Effect } from "effect"; import { config } from "../../core/config"; +import { serverEvents } from "../../core/events"; import { createAgentManagerRuntime, type AgentManagerEvent } from "./controller/server"; import { LOCAL_AGENT_ID } from "./constants"; import { spawnLocalAgentProcess, stopLocalAgentProcess } from "./local/process"; @@ -212,6 +218,15 @@ const handleAgentManagerEvent = (event: AgentManagerEvent) => { }); break; } + case "restore.progress": { + serverEvents.emit("restore:progress", { + organizationId: event.payload.organizationId, + repositoryId: event.payload.repositoryId, + snapshotId: event.payload.snapshotId, + ...event.payload.progress, + }); + break; + } } }; @@ -307,6 +322,23 @@ export const agentManager = { return response.command; }, + runRestoreCommand: async (agentId: string, payload: Omit) => { + const runtime = getAgentManagerRuntime(); + if (!runtime) { + throw new Error(`Restore agent ${agentId} is not connected`); + } + + const response = await Effect.runPromise(runtime.runRestoreCommand(agentId, payload)); + if (!response) { + throw new Error(`Failed to send restore command ${payload.snapshotId} to agent ${agentId}`); + } + + if (response.status === "error") { + throw new Error(response.error); + } + + return response.result; + }, }; export const startLocalAgent = async () => { diff --git a/app/server/modules/agents/controller/server.ts b/app/server/modules/agents/controller/server.ts index b3de99f5..f84a9d89 100644 --- a/app/server/modules/agents/controller/server.ts +++ b/app/server/modules/agents/controller/server.ts @@ -5,6 +5,8 @@ import type { AgentMessage, BackupCancelPayload, BackupRunPayload, + RestoreCommandPayload, + RestoreCommandResponsePayload, VolumeCommand, VolumeCommandResponsePayload, } from "@zerobyte/contracts/agent-protocol"; @@ -22,16 +24,9 @@ type AgentEventContext = { agentName: string; }; -type AgentBackupMessage = Extract< - AgentMessage, - { - type: "backup.started" | "backup.progress" | "backup.completed" | "backup.failed" | "backup.cancelled"; - } ->; - export type AgentManagerEvent = | (AgentEventContext & { type: "agent.disconnected" }) - | (AgentEventContext & AgentBackupMessage); + | (AgentEventContext & AgentMessage); type ControllerAgentSessionHandle = { agentId: string; @@ -384,6 +379,31 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) => yield* logger.effect.info(`Completed volume command ${command.name} on agent ${agentId}`); return result; }), + runRestoreCommand: ( + agentId: string, + payload: Omit, + ): Effect.Effect => + Effect.gen(function* () { + const session = getSession(agentId); + + if (!session) { + yield* logger.effect.warn( + `Cannot send restore command ${payload.snapshotId}. Agent ${agentId} is not connected.`, + ); + return null; + } + + if (!(yield* session.isReady())) { + yield* logger.effect.warn( + `Cannot send restore command ${payload.snapshotId}. Agent ${agentId} is not ready.`, + ); + return null; + } + + const result = yield* session.runRestoreCommand(payload); + yield* logger.effect.info(`Completed restore command ${payload.snapshotId} on agent ${agentId}`); + return result; + }), stop, }; } diff --git a/app/server/modules/agents/controller/session.ts b/app/server/modules/agents/controller/session.ts index 5a68402e..2956465c 100644 --- a/app/server/modules/agents/controller/session.ts +++ b/app/server/modules/agents/controller/session.ts @@ -7,6 +7,8 @@ import { type BackupCancelPayload, type BackupRunPayload, type ControllerWireMessage, + type RestoreCommandPayload, + type RestoreCommandResponsePayload, type VolumeCommand, type VolumeCommandResponsePayload, } from "@zerobyte/contracts/agent-protocol"; @@ -29,16 +31,11 @@ type SessionState = { lastPongAt: number | null; }; -type PendingCommand = { - deferred: Deferred.Deferred; - description: string; -}; +type PendingCommand = { deferred: Deferred.Deferred; description: string }; export type ControllerAgentSessionEvent = - | Exclude - | { - type: "agent.disconnected"; - }; + | Exclude + | { type: "agent.disconnected" }; export type ControllerAgentSession = { readonly connectionId: string; @@ -46,6 +43,9 @@ export type ControllerAgentSession = { sendBackup: (payload: BackupRunPayload) => Effect.Effect; sendBackupCancel: (payload: BackupCancelPayload) => Effect.Effect; runVolumeCommand: (command: VolumeCommand) => Effect.Effect; + runRestoreCommand: ( + payload: Omit, + ) => Effect.Effect; isReady: () => Effect.Effect; run: Effect.Effect; }; @@ -68,7 +68,9 @@ export const createControllerAgentSession = ( Queue.offer(outboundQueue, message).pipe( Effect.catchAllCause((cause) => Effect.sync(() => { - logger.error(`Failed to queue outbound message for agent ${socket.data.agentId}: ${toMessage(cause)}`); + logger.error( + `Failed to queue outbound message for agent ${socket.data.agentId}: ${toMessage(cause)}`, + ); return false; }), ), @@ -177,6 +179,17 @@ export const createControllerAgentSession = ( yield* Deferred.succeed(pending.deferred, payload); }); + const handleRestoreCommandResult = (payload: RestoreCommandResponsePayload) => + Effect.gen(function* () { + const pending = yield* removePendingCommand(payload.commandId); + if (!pending) { + yield* logger.effect.warn(`Received response for unknown restore command ${payload.commandId}`); + return; + } + + yield* Deferred.succeed(pending.deferred, payload); + }); + const handleAgentMessage = (message: AgentMessage) => Effect.gen(function* () { switch (message.type) { @@ -189,7 +202,11 @@ export const createControllerAgentSession = ( } case "heartbeat.pong": { const seenAt = Date.now(); - yield* updateState((current) => ({ ...current, lastSeenAt: seenAt, lastPongAt: message.payload.sentAt })); + yield* updateState((current) => ({ + ...current, + lastSeenAt: seenAt, + lastPongAt: message.payload.sentAt, + })); yield* onEvent(message); break; } @@ -197,6 +214,10 @@ export const createControllerAgentSession = ( yield* handleVolumeCommandResult(message.payload); break; } + case "restore.commandResult": { + yield* handleRestoreCommandResult(message.payload); + break; + } default: { yield* onEvent(message); break; @@ -216,7 +237,9 @@ export const createControllerAgentSession = ( } if (!parsed.success) { - yield* logger.effect.warn(`Invalid agent message from ${socket.data.agentId}: ${parsed.error.message}`); + yield* logger.effect.warn( + `Invalid agent message from ${socket.data.agentId}: ${parsed.error.message}`, + ); return; } @@ -232,19 +255,44 @@ export const createControllerAgentSession = ( const deferred = yield* Deferred.make(); yield* setPendingCommand(commandId, { deferred, description }); - const queued = yield* offerOutbound(createControllerMessage("volume.command", { commandId, command })); + const queued = yield* offerOutbound( + createControllerMessage("volume.command", { commandId, command }), + ); if (!queued) { yield* removePendingCommand(commandId); return yield* Effect.fail(new Error(`Failed to queue volume command ${command.name}`)); } - return yield* Deferred.await(deferred).pipe( + return (yield* Deferred.await(deferred).pipe( Effect.timeoutFail({ duration: "60 seconds", onTimeout: () => new Error(`Volume command ${command.name} timed out`), }), Effect.ensuring(removePendingCommand(commandId)), + )) as VolumeCommandResponsePayload; + }), + runRestoreCommand: (payload) => + Effect.gen(function* () { + const commandId = Bun.randomUUIDv7(); + const description = `restore command ${payload.snapshotId}`; + const deferred = yield* Deferred.make(); + yield* setPendingCommand(commandId, { deferred, description }); + + const queued = yield* offerOutbound( + createControllerMessage("restore.command", { commandId, ...payload }), ); + if (!queued) { + yield* removePendingCommand(commandId); + return yield* Effect.fail(new Error(`Failed to queue restore command ${payload.snapshotId}`)); + } + + return (yield* Deferred.await(deferred).pipe( + Effect.timeoutFail({ + duration: "1 hour", + onTimeout: () => new Error(`Restore command ${payload.snapshotId} timed out`), + }), + Effect.ensuring(removePendingCommand(commandId)), + )) as RestoreCommandResponsePayload; }), isReady: () => Ref.get(state).pipe(Effect.map((current) => current.isReady)), run, diff --git a/app/server/modules/backups/backup.helpers.ts b/app/server/modules/backups/backup.helpers.ts index dcafb27e..095f831d 100644 --- a/app/server/modules/backups/backup.helpers.ts +++ b/app/server/modules/backups/backup.helpers.ts @@ -1,5 +1,5 @@ import { CronExpressionParser } from "cron-parser"; -import { createBackupOptions as createAgentBackupOptions } from "../../../../apps/agent/src/commands/backup.helpers"; +import { createBackupOptions as createAgentBackupOptions } from "../../../../apps/agent/src/commands/helpers/backup.helpers"; import type { BackupSchedule } from "~/server/db/schema"; import { toMessage } from "~/server/utils/errors"; import { logger } from "@zerobyte/core/node"; diff --git a/app/server/modules/repositories/__tests__/repositories.service.test.ts b/app/server/modules/repositories/__tests__/repositories.service.test.ts index 14ab2450..4184dc14 100644 --- a/app/server/modules/repositories/__tests__/repositories.service.test.ts +++ b/app/server/modules/repositories/__tests__/repositories.service.test.ts @@ -4,22 +4,24 @@ import * as os from "node:os"; import nodePath from "node:path"; import { randomUUID } from "node:crypto"; import { Readable } from "node:stream"; +import { Effect } from "effect"; import { afterEach, beforeAll, beforeEach, describe, expect, test, vi } from "vitest"; import type { RepositoryConfig } from "@zerobyte/core/restic"; import { REPOSITORY_BASE } from "~/server/core/constants"; +import { config } from "~/server/core/config"; import { withContext } from "~/server/core/request-context"; import { db } from "~/server/db/db"; -import { repositoriesTable } from "~/server/db/schema"; +import { repositoriesTable, type RepositoryInsert } from "~/server/db/schema"; import { generateShortId } from "~/server/utils/id"; import { restic } from "~/server/core/restic"; +import { agentManager } from "~/server/modules/agents/agents-manager"; import { createTestSession } from "~/test/helpers/auth"; import { createTestBackupSchedule } from "~/test/helpers/backup"; import { cache, cacheKeys } from "~/server/utils/cache"; import { ResticError } from "@zerobyte/core/restic/server"; import { repositoriesService } from "../repositories.service"; -import { Effect } from "effect"; -const createTestRepository = async (organizationId: string) => { +const createTestRepository = async (organizationId: string, overrides: Partial = {}) => { const id = randomUUID(); const shortId = generateShortId(); const [repository] = await db @@ -33,6 +35,7 @@ const createTestRepository = async (organizationId: string) => { compressionMode: "auto", status: "healthy", organizationId, + ...overrides, }) .returning(); return repository; @@ -404,7 +407,15 @@ describe("repositoriesService.dumpSnapshot", () => { }); describe("repositoriesService.restoreSnapshot", () => { + let originalEnableLocalAgent: boolean; + + beforeEach(() => { + originalEnableLocalAgent = config.flags.enableLocalAgent; + config.flags.enableLocalAgent = true; + }); + afterEach(() => { + config.flags.enableLocalAgent = originalEnableLocalAgent; vi.restoreAllMocks(); }); @@ -425,7 +436,7 @@ describe("repositoriesService.restoreSnapshot", () => { ); const restoreMock = vi.fn(() => - Effect.succeed({ + Promise.resolve({ message_type: "summary" as const, seconds_elapsed: 1, percent_done: 100, @@ -436,7 +447,7 @@ describe("repositoriesService.restoreSnapshot", () => { bytes_restored: 1, }), ); - vi.spyOn(restic, "restore").mockImplementation(restoreMock); + vi.spyOn(agentManager, "runRestoreCommand").mockImplementation(restoreMock); return { organizationId, @@ -446,9 +457,14 @@ describe("repositoriesService.restoreSnapshot", () => { }; }; - test("rejects restore targets inside protected roots", async () => { + test("delegates restore target protection to the selected agent", async () => { const { organizationId, userId, repositoryShortId, restoreMock } = await setupRestoreSnapshotScenario(); const targetPath = nodePath.join(os.tmpdir(), "zerobyte-restore-target"); + restoreMock.mockRejectedValueOnce( + new Error( + "Restore target path is not allowed. Restoring to this path could overwrite critical system files or application data.", + ), + ); await expect( withContext({ organizationId, userId }, () => @@ -456,7 +472,7 @@ describe("repositoriesService.restoreSnapshot", () => { ), ).rejects.toThrow("Restore target path is not allowed"); - expect(restoreMock).not.toHaveBeenCalled(); + expect(restoreMock).toHaveBeenCalledOnce(); }); test("restores to a custom target outside protected roots", async () => { @@ -472,15 +488,97 @@ describe("repositoriesService.restoreSnapshot", () => { } expect(restoreMock).toHaveBeenCalledWith( + "local", expect.objectContaining({ - backend: "local", + snapshotId: "snapshot-restore", + target: targetPath, + repositoryConfig: expect.objectContaining({ backend: "local" }), + options: expect.objectContaining({ + organizationId, + basePath: "/var/lib/zerobyte/volumes/vol123/_data", + }), }), + ); + }); + + test("routes restore to the requested target agent", async () => { + const organizationId = session.organizationId; + const repository = await createTestRepository(organizationId, { + type: "s3", + config: { + backend: "s3", + endpoint: "https://s3.example.com", + bucket: "bucket", + accessKeyId: "access-key", + secretAccessKey: "secret-key", + }, + }); + vi.spyOn(restic, "snapshots").mockReturnValue( + Effect.succeed([ + { + id: "snapshot-restore", + short_id: "snapshot-restore", + time: new Date().toISOString(), + paths: ["/var/lib/zerobyte/volumes/vol123/_data"], + hostname: "host", + }, + ]), + ); + const restoreMock = vi.fn(() => + Promise.resolve({ + message_type: "summary" as const, + files_skipped: 0, + files_restored: 1, + }), + ); + vi.spyOn(agentManager, "runRestoreCommand").mockImplementation(restoreMock); + const targetPath = await fs.mkdtemp(nodePath.join(process.cwd(), "restore-target-")); + + try { + await withContext({ organizationId, userId: session.user.id }, () => + repositoriesService.restoreSnapshot(repository.shortId, "snapshot-restore", { + targetPath, + targetAgentId: "agent-remote", + }), + ); + } finally { + await fs.rm(targetPath, { recursive: true, force: true }); + } + + expect(restoreMock).toHaveBeenCalledWith( + "agent-remote", + expect.objectContaining({ + target: targetPath, + }), + ); + }); + + test("uses controller-local restore fallback when local agent supervision is disabled", async () => { + config.flags.enableLocalAgent = false; + const { organizationId, userId, repositoryShortId, restoreMock } = await setupRestoreSnapshotScenario(); + const resticRestoreMock = vi.spyOn(restic, "restore").mockReturnValue( + Effect.succeed({ + message_type: "summary" as const, + files_skipped: 0, + files_restored: 1, + }), + ); + const targetPath = await fs.mkdtemp(nodePath.join(process.cwd(), "restore-target-")); + + try { + await withContext({ organizationId, userId }, () => + repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", { targetPath }), + ); + } finally { + await fs.rm(targetPath, { recursive: true, force: true }); + } + + expect(restoreMock).not.toHaveBeenCalled(); + expect(resticRestoreMock).toHaveBeenCalledWith( + expect.objectContaining({ backend: "local" }), "snapshot-restore", targetPath, - expect.objectContaining({ - organizationId, - basePath: "/var/lib/zerobyte/volumes/vol123/_data", - }), + expect.objectContaining({ organizationId, basePath: "/var/lib/zerobyte/volumes/vol123/_data" }), ); }); @@ -516,14 +614,15 @@ describe("repositoriesService.restoreSnapshot", () => { } expect(restoreMock).toHaveBeenCalledWith( + "local", expect.objectContaining({ - backend: "local", - }), - "snapshot-restore", - targetPath, - expect.objectContaining({ - organizationId, - basePath: "/", + snapshotId: "snapshot-restore", + target: targetPath, + repositoryConfig: expect.objectContaining({ backend: "local" }), + options: expect.objectContaining({ + organizationId, + basePath: "/", + }), }), ); }); diff --git a/app/server/modules/repositories/repositories.dto.ts b/app/server/modules/repositories/repositories.dto.ts index 79ca2d00..3a346e87 100644 --- a/app/server/modules/repositories/repositories.dto.ts +++ b/app/server/modules/repositories/repositories.dto.ts @@ -343,6 +343,7 @@ export const restoreSnapshotBody = z.object({ excludeXattr: z.array(z.string()).optional(), delete: z.boolean().optional(), targetPath: z.string().optional(), + targetAgentId: z.string().optional(), overwrite: overwriteModeSchema.optional(), }); diff --git a/app/server/modules/repositories/repositories.service.ts b/app/server/modules/repositories/repositories.service.ts index 61eeaeb4..ab8ee8aa 100644 --- a/app/server/modules/repositories/repositories.service.ts +++ b/app/server/modules/repositories/repositories.service.ts @@ -19,7 +19,7 @@ import { logger } from "@zerobyte/core/node"; import { parseRetentionCategories, type RetentionCategory } from "~/server/utils/retention-categories"; import { repoMutex } from "../../core/repository-mutex"; import { db } from "../../db/db"; -import { repositoriesTable, type Repository } from "../../db/schema"; +import { repositoriesTable, type Repository, type RepositoryInsert } from "../../db/schema"; import { cache, cacheKeys } from "../../utils/cache"; import { runEffectPromise, toMessage } from "../../utils/errors"; import { generateShortId } from "../../utils/id"; @@ -33,6 +33,8 @@ import { executeDoctor } from "./helpers/doctor"; import type { ShortId } from "~/server/utils/branded"; import { decryptRepositoryConfig, encryptRepositoryConfig } from "./repository-config-secrets"; import { getScheduleByIdOrShortId } from "../backups/helpers/backup-schedule-lookups"; +import { agentManager } from "../agents/agents-manager"; +import { LOCAL_AGENT_ID } from "../agents/constants"; import { Effect } from "effect"; const runningDoctors = new Map(); @@ -56,6 +58,21 @@ const getBlockedRestoreTargets = () => { ].filter((e) => e !== undefined); }; +const assertAllowedControllerLocalRestoreTarget = (target: string) => { + const resolvedTarget = nodePath.resolve(target); + + for (const blockedTarget of getBlockedRestoreTargets()) { + if (isPathWithin(blockedTarget, resolvedTarget)) { + throw new BadRequestError( + "Restore target path is not allowed. Restoring to this path could overwrite critical system files or application data.", + ); + } + } +}; + +const shouldUseControllerLocalRestoreFallback = (agentId: string) => + agentId === LOCAL_AGENT_ID && !appConfig.flags.enableLocalAgent; + const findRepository = async (shortId: ShortId) => { const organizationId = getOrganizationId(); return await db.query.repositoriesTable.findFirst({ @@ -327,6 +344,7 @@ const restoreSnapshot = async ( excludeXattr?: string[]; delete?: boolean; targetPath?: string; + targetAgentId?: string; overwrite?: OverwriteMode; }, ) => { @@ -338,19 +356,9 @@ const restoreSnapshot = async ( } const target = options?.targetPath || "/"; - const resolvedTarget = nodePath.resolve(target); - const blockedTargets = getBlockedRestoreTargets(); - for (const blockedTarget of blockedTargets) { - if (isPathWithin(blockedTarget, resolvedTarget)) { - throw new BadRequestError( - "Restore target path is not allowed. Restoring to this path could overwrite critical system files or application data.", - ); - } - } - - const { paths } = await getSnapshotDetails(repository.shortId, snapshotId); - const hasNonPosixSnapshotPaths = paths.some((path) => !path.startsWith("/")); + const snapshot = await getSnapshotDetails(repository.shortId, snapshotId); + const hasNonPosixSnapshotPaths = snapshot.paths.some((path) => !path.startsWith("/")); if (hasNonPosixSnapshotPaths && !options?.targetPath) { throw new BadRequestError( @@ -358,7 +366,15 @@ const restoreSnapshot = async ( ); } - const basePath = hasNonPosixSnapshotPaths ? "/" : findCommonAncestor(paths); + const basePath = hasNonPosixSnapshotPaths ? "/" : findCommonAncestor(snapshot.paths); + const executionAgentId = options?.targetAgentId ?? LOCAL_AGENT_ID; + const useControllerLocalRestoreFallback = shouldUseControllerLocalRestoreFallback(executionAgentId); + + if (!useControllerLocalRestoreFallback && repository.type === "local" && executionAgentId !== LOCAL_AGENT_ID) { + throw new BadRequestError( + "Local repository restores must run on the agent that can access the repository path.", + ); + } const releaseLock = await repoMutex.acquireShared(repository.id, `restore:${snapshotId}`); try { @@ -368,21 +384,43 @@ const restoreSnapshot = async ( snapshotId, }); - const result = await runEffectPromise( - restic.restore(repository.config, snapshotId, target, { - basePath, - ...options, + const repositoryConfig = await decryptRepositoryConfig(repository.config); + let result; + + if (useControllerLocalRestoreFallback) { + assertAllowedControllerLocalRestoreTarget(target); + result = await runEffectPromise( + restic.restore(repositoryConfig, snapshotId, target, { + basePath, + ...options, + organizationId, + onProgress: (progress) => { + serverEvents.emit("restore:progress", { + organizationId, + repositoryId: repository.shortId, + snapshotId, + ...progress, + }); + }, + }), + ); + } else { + const encryptedResticPassword = await resticDeps.getOrganizationResticPassword(organizationId); + const resticPassword = await resticDeps.resolveSecret(encryptedResticPassword); + result = await agentManager.runRestoreCommand(executionAgentId, { organizationId, - onProgress: (progress) => { - serverEvents.emit("restore:progress", { - organizationId, - repositoryId: repository.shortId, - snapshotId, - ...progress, - }); + repositoryId: repository.shortId, + snapshotId, + target, + repositoryConfig, + runtime: { password: resticPassword }, + options: { + basePath, + ...options, + organizationId, }, - }), - ); + }); + } serverEvents.emit("restore:completed", { organizationId, @@ -738,23 +776,22 @@ const updateRepository = async (shortId: ShortId, updates: UpdateRepositoryBody) const configChanged = updates.config && JSON.stringify(decryptedExisting) !== JSON.stringify(parsedConfig); const encryptedConfig = updates.config ? await encryptRepositoryConfig(parsedConfig) : existingConfig; const updatedAt = Date.now(); - const updatePayload = { + const updatePayload: Partial = { name: newName, compressionMode: updates.compressionMode ?? existing.compressionMode, updatedAt, config: encryptedConfig, - ...(configChanged - ? { - status: "unknown" as const, - lastChecked: null, - lastError: null, - doctorResult: null, - stats: null, - statsUpdatedAt: null, - } - : {}), }; + if (configChanged) { + updatePayload.status = "unknown"; + updatePayload.lastChecked = null; + updatePayload.lastError = null; + updatePayload.doctorResult = null; + updatePayload.stats = null; + updatePayload.statsUpdatedAt = null; + } + const [updated] = await db .update(repositoriesTable) .set(updatePayload) diff --git a/apps/agent/src/commands/volume.test.ts b/apps/agent/src/commands/__tests__/volume.test.ts similarity index 93% rename from apps/agent/src/commands/volume.test.ts rename to apps/agent/src/commands/__tests__/volume.test.ts index 28c3fbdf..07527838 100644 --- a/apps/agent/src/commands/volume.test.ts +++ b/apps/agent/src/commands/__tests__/volume.test.ts @@ -6,7 +6,7 @@ import { type AgentWireMessage, type VolumeCommandPayload, } from "@zerobyte/contracts/agent-protocol"; -import type { ControllerCommandContext } from "../context"; +import type { ControllerCommandContext } from "../../context"; const volumeHostMock = vi.hoisted(() => ({ createVolumeBackend: vi.fn(), @@ -20,10 +20,10 @@ const operationsMock = vi.hoisted(() => ({ testVolumeConnection: vi.fn(), })); -vi.mock("../volume-host", () => volumeHostMock); -vi.mock("../volume-host/operations", () => operationsMock); +vi.mock("../../volume-host", () => volumeHostMock); +vi.mock("../../volume-host/operations", () => operationsMock); -import { handleVolumeCommand } from "./volume"; +import { handleVolumeCommand } from "../volume"; afterEach(() => { vi.restoreAllMocks(); diff --git a/apps/agent/src/commands/backup-run.ts b/apps/agent/src/commands/backup-run.ts index 17e53e36..a4f40187 100644 --- a/apps/agent/src/commands/backup-run.ts +++ b/apps/agent/src/commands/backup-run.ts @@ -8,7 +8,7 @@ import { toMessage } from "@zerobyte/core/utils"; import type { ControllerCommandContext } from "../context"; import { resticDeps } from "../restic/deps"; import { createVolumeBackend, getVolumePath } from "../volume-host"; -import { createBackupOptions } from "./backup.helpers"; +import { createBackupOptions } from "./helpers/backup.helpers"; class VolumeReadinessError extends Data.TaggedError("VolumeReadinessError")<{ readonly _tag: "VolumeReadinessError"; diff --git a/apps/agent/src/commands/__tests__/backup.helpers.test.ts b/apps/agent/src/commands/helpers/__tests__/backup.helpers.test.ts similarity index 100% rename from apps/agent/src/commands/__tests__/backup.helpers.test.ts rename to apps/agent/src/commands/helpers/__tests__/backup.helpers.test.ts diff --git a/apps/agent/src/commands/backup.helpers.ts b/apps/agent/src/commands/helpers/backup.helpers.ts similarity index 100% rename from apps/agent/src/commands/backup.helpers.ts rename to apps/agent/src/commands/helpers/backup.helpers.ts diff --git a/apps/agent/src/commands/index.ts b/apps/agent/src/commands/index.ts index 6f73e0c6..780e8aec 100644 --- a/apps/agent/src/commands/index.ts +++ b/apps/agent/src/commands/index.ts @@ -3,6 +3,7 @@ import { handleBackupCancelCommand } from "./backup-cancel"; import { handleBackupRunCommand } from "./backup-run"; import type { ControllerCommandContext } from "../context"; import { handleHeartbeatPingCommand } from "./heartbeat-ping"; +import { handleRestoreCommand } from "./restore"; import { handleVolumeCommand } from "./volume"; export const handleControllerCommand = (context: ControllerCommandContext, message: ControllerMessage) => { @@ -16,6 +17,9 @@ export const handleControllerCommand = (context: ControllerCommandContext, messa case "volume.command": { return handleVolumeCommand(context, message.payload); } + case "restore.command": { + return handleRestoreCommand(context, message.payload); + } case "heartbeat.ping": { return handleHeartbeatPingCommand(context, message.payload); } diff --git a/apps/agent/src/commands/restore.ts b/apps/agent/src/commands/restore.ts new file mode 100644 index 00000000..5e159595 --- /dev/null +++ b/apps/agent/src/commands/restore.ts @@ -0,0 +1,73 @@ +import os from "node:os"; +import path from "node:path"; +import { Effect, Runtime } from "effect"; +import { createAgentMessage, type RestoreCommandPayload } from "@zerobyte/contracts/agent-protocol"; +import { createRestic } from "@zerobyte/core/restic/server"; +import { isPathWithin, toMessage } from "@zerobyte/core/utils"; +import { logger } from "@zerobyte/core/node"; +import type { ControllerCommandContext } from "../context"; +import { resticDeps } from "../restic/deps"; + +const REPOSITORY_BASE = process.env.ZEROBYTE_REPOSITORIES_DIR || "/var/lib/zerobyte/repositories"; +const RESTIC_PASS_FILE = process.env.RESTIC_PASS_FILE || "/var/lib/zerobyte/data/restic.pass"; + +const getBlockedRestoreTargets = () => + [REPOSITORY_BASE, path.dirname(RESTIC_PASS_FILE), os.tmpdir()].map((target) => path.resolve(target)); + +const assertAllowedRestoreTarget = (target: string) => { + const resolvedTarget = path.resolve(target); + + for (const blockedTarget of getBlockedRestoreTargets()) { + if (isPathWithin(blockedTarget, resolvedTarget)) { + throw new Error( + "Restore target path is not allowed. Restoring to this path could overwrite critical system files or application data.", + ); + } + } +}; + +export const handleRestoreCommand = (context: ControllerCommandContext, payload: RestoreCommandPayload) => { + return Effect.gen(function* () { + assertAllowedRestoreTarget(payload.target); + + const runtime = yield* Effect.runtime(); + const restic = createRestic(resticDeps(payload.runtime.password)); + const result = yield* restic.restore(payload.repositoryConfig, payload.snapshotId, payload.target, { + ...payload.options, + onProgress: (progress) => { + void Runtime.runPromise( + runtime, + context.offerOutbound( + createAgentMessage("restore.progress", { + commandId: payload.commandId, + organizationId: payload.organizationId, + repositoryId: payload.repositoryId, + snapshotId: payload.snapshotId, + progress, + }), + ), + ).catch((error) => { + logger.error(`Failed to send restore progress update: ${toMessage(error)}`); + }); + }, + }); + + yield* context.offerOutbound( + createAgentMessage("restore.commandResult", { + commandId: payload.commandId, + status: "success", + result, + }), + ); + }).pipe( + Effect.catchAll((error) => + context.offerOutbound( + createAgentMessage("restore.commandResult", { + commandId: payload.commandId, + status: "error", + error: toMessage(error), + }), + ), + ), + ); +}; diff --git a/packages/contracts/src/agent-protocol.ts b/packages/contracts/src/agent-protocol.ts index bd78ad58..57a1a3b4 100644 --- a/packages/contracts/src/agent-protocol.ts +++ b/packages/contracts/src/agent-protocol.ts @@ -5,6 +5,8 @@ import { repositoryConfigSchema, resticBackupOutputSchema, resticBackupProgressSchema, + resticRestoreOutputSchema, + restoreProgressSchema, type CompressionMode, } from "@zerobyte/core/restic"; import { @@ -96,6 +98,48 @@ const volumeCommandResponseSchema = z.object({ ]), }); +const restoreCommandSchema = z.object({ + type: z.literal("restore.command"), + payload: z.object({ + commandId: z.string(), + organizationId: z.string(), + repositoryId: z.string(), + snapshotId: z.string(), + target: z.string(), + repositoryConfig: repositoryConfigSchema, + runtime: z.object({ password: z.string() }), + options: z.object({ + basePath: z.string().optional(), + organizationId: z.string(), + include: z.array(z.string()).optional(), + selectedItemKind: z.enum(["file", "dir"]).optional(), + exclude: z.array(z.string()).optional(), + excludeXattr: z.array(z.string()).optional(), + delete: z.boolean().optional(), + overwrite: z.enum(["always", "if-changed", "if-newer", "never"]).optional(), + }), + }), +}); + +const restoreProgressMessageSchema = z.object({ + type: z.literal("restore.progress"), + payload: z.object({ + commandId: z.string(), + organizationId: z.string(), + repositoryId: z.string(), + snapshotId: z.string(), + progress: restoreProgressSchema, + }), +}); + +const restoreCommandResultSchema = z.object({ + type: z.literal("restore.commandResult"), + payload: z.discriminatedUnion("status", [ + z.object({ commandId: z.string(), status: z.literal("success"), result: resticRestoreOutputSchema }), + z.object({ commandId: z.string(), status: z.literal("error"), error: z.string() }), + ]), +}); + const heartbeatPingSchema = z.object({ type: z.literal("heartbeat.ping"), payload: z.object({ sentAt: z.number() }), @@ -165,6 +209,7 @@ const controllerMessageSchema = z.discriminatedUnion("type", [ backupRunSchema, backupCancelSchema, volumeCommandRequestSchema, + restoreCommandSchema, heartbeatPingSchema, ]); const agentMessageSchema = z.discriminatedUnion("type", [ @@ -175,6 +220,8 @@ const agentMessageSchema = z.discriminatedUnion("type", [ backupFailedSchema, backupCancelledSchema, volumeCommandResponseSchema, + restoreProgressMessageSchema, + restoreCommandResultSchema, heartbeatPongSchema, ]); @@ -189,6 +236,8 @@ export type VolumeCommandPayload = z.infer["p export type VolumeCommand = z.infer; export type VolumeCommandResult = z.infer; export type VolumeCommandResponsePayload = z.infer["payload"]; +export type RestoreCommandPayload = z.infer["payload"]; +export type RestoreCommandResponsePayload = z.infer["payload"]; export type ControllerMessage = z.infer; export type AgentMessage = z.infer; diff --git a/packages/core/src/restic/commands/__tests__/restore.test.ts b/packages/core/src/restic/commands/__tests__/restore.test.ts index a043a923..914c4063 100644 --- a/packages/core/src/restic/commands/__tests__/restore.test.ts +++ b/packages/core/src/restic/commands/__tests__/restore.test.ts @@ -1,11 +1,11 @@ import { afterEach, describe, expect, test, vi } from "vitest"; +import { Effect } from "effect"; import * as cleanupModule from "../../helpers/cleanup-temporary-keys"; import * as spawnModule from "../../../node/spawn"; import { ResticError } from "../../error"; import { restore } from "../restore"; import type { ResticDeps } from "../../types"; import type { SafeSpawnParams, SpawnResult } from "../../../node/spawn"; -import { Effect } from "effect"; const mockDeps: ResticDeps = { resolveSecret: async (s) => s, @@ -44,14 +44,19 @@ const config = { type SetupOptions = { spawnResult?: Partial; onSpawnCall?: (params: SafeSpawnParams) => void | Promise; + spawnError?: unknown; }; -const setup = ({ spawnResult = {}, onSpawnCall }: SetupOptions = {}) => { +const setup = ({ spawnResult = {}, onSpawnCall, spawnError }: SetupOptions = {}) => { let capturedArgs: string[] = []; vi.spyOn(cleanupModule, "cleanupTemporaryKeys").mockImplementation(() => Promise.resolve()); vi.spyOn(spawnModule, "safeSpawn").mockImplementation((params: SafeSpawnParams) => { capturedArgs = params.args; + if (spawnError) { + return Promise.reject(spawnError); + } + return Promise.resolve(onSpawnCall?.(params)).then(() => ({ exitCode: 0, summary: successfulRestoreSummary, @@ -89,25 +94,26 @@ afterEach(() => { vi.restoreAllMocks(); }); +const runRestore = (...args: Parameters) => Effect.runPromise(restore(...args)); +const runRestoreError = (...args: Parameters) => Effect.runPromise(Effect.flip(restore(...args))); + describe("restore command", () => { describe("path selection", () => { test("uses the common ancestor as restore root and strips includes for non-root targets", async () => { const { getRestoreArg, getOptionValues } = setup(); - await Effect.runPromise( - restore( - config, - "snapshot-456", - "/tmp/restore-target", - { - organizationId: "org-1", - include: [ - "/var/lib/zerobyte/volumes/vol123/_data/Documents/report.pdf", - "/var/lib/zerobyte/volumes/vol123/_data/Photos/summer.jpg", - ], - }, - mockDeps, - ), + await runRestore( + config, + "snapshot-456", + "/tmp/restore-target", + { + organizationId: "org-1", + include: [ + "/var/lib/zerobyte/volumes/vol123/_data/Documents/report.pdf", + "/var/lib/zerobyte/volumes/vol123/_data/Photos/summer.jpg", + ], + }, + mockDeps, ); expect(getRestoreArg()).toBe("snapshot-456:/var/lib/zerobyte/volumes/vol123/_data"); @@ -117,18 +123,16 @@ describe("restore command", () => { test("restores a selected file from its parent directory for non-root targets", async () => { const { getRestoreArg, getOptionValues } = setup(); - await Effect.runPromise( - restore( - config, - "snapshot-single-file", - "/tmp/restore-target", - { - organizationId: "org-1", - include: ["/var/lib/zerobyte/volumes/vol123/_data/archive/backup.20260301-233001.7z"], - selectedItemKind: "file", - }, - mockDeps, - ), + await runRestore( + config, + "snapshot-single-file", + "/tmp/restore-target", + { + organizationId: "org-1", + include: ["/var/lib/zerobyte/volumes/vol123/_data/archive/backup.20260301-233001.7z"], + selectedItemKind: "file", + }, + mockDeps, ); expect(getRestoreArg()).toBe("snapshot-single-file:/var/lib/zerobyte/volumes/vol123/_data/archive"); @@ -138,17 +142,15 @@ describe("restore command", () => { test("treats flag-like snapshot IDs as positional restore args", async () => { const { getArgs, getRestoreArg } = setup(); - await Effect.runPromise( - restore( - config, - "--help", - "/tmp/restore-target", - { - organizationId: "org-1", - basePath: "/var/lib/zerobyte/volumes/vol123/_data", - }, - mockDeps, - ), + await runRestore( + config, + "--help", + "/tmp/restore-target", + { + organizationId: "org-1", + basePath: "/var/lib/zerobyte/volumes/vol123/_data", + }, + mockDeps, ); const separatorIndex = getArgs().indexOf("--"); @@ -161,14 +163,12 @@ describe("restore command", () => { test("returns a parsed restore summary on success", async () => { setup(); - const result = await Effect.runPromise( - restore( - config, - "snapshot-123", - "/tmp/restore-target", - { organizationId: "org-1", basePath: "/var/lib/zerobyte/volumes/vol123/_data" }, - mockDeps, - ), + const result = await runRestore( + config, + "snapshot-123", + "/tmp/restore-target", + { organizationId: "org-1", basePath: "/var/lib/zerobyte/volumes/vol123/_data" }, + mockDeps, ); expect(result).toMatchObject({ @@ -181,32 +181,41 @@ describe("restore command", () => { test("throws ResticError when the command fails", async () => { setup({ spawnResult: { exitCode: 1, summary: "", error: "restore failed" } }); - const error = await Effect.runPromise( - Effect.flip( - restore( - config, - "snapshot-123", - "/tmp/restore-target", - { organizationId: "org-1", basePath: "/var/lib/zerobyte/volumes/vol123/_data" }, - mockDeps, - ), - ), - ); - - expect(error).toBeInstanceOf(ResticError); - }); - - test("falls back to an empty summary when restic output cannot be parsed", async () => { - setup({ spawnResult: { summary: "not-json" } }); - - const result = await Effect.runPromise( - restore( + await expect( + runRestoreError( config, "snapshot-123", "/tmp/restore-target", { organizationId: "org-1", basePath: "/var/lib/zerobyte/volumes/vol123/_data" }, mockDeps, ), + ).resolves.toBeInstanceOf(ResticError); + }); + + test("cleans up temporary keys when spawning restic rejects", async () => { + const cleanupSpy = vi.spyOn(cleanupModule, "cleanupTemporaryKeys"); + setup({ spawnError: new Error("spawn failed") }); + + await runRestoreError( + config, + "snapshot-123", + "/tmp/restore-target", + { organizationId: "org-1", basePath: "/var/lib/zerobyte/volumes/vol123/_data" }, + mockDeps, + ); + + expect(cleanupSpy).toHaveBeenCalledTimes(1); + }); + + test("falls back to an empty summary when restic output cannot be parsed", async () => { + setup({ spawnResult: { summary: "not-json" } }); + + const result = await runRestore( + config, + "snapshot-123", + "/tmp/restore-target", + { organizationId: "org-1", basePath: "/var/lib/zerobyte/volumes/vol123/_data" }, + mockDeps, ); expect(result).toEqual({ @@ -224,18 +233,16 @@ describe("restore command", () => { const progressUpdates: unknown[] = []; setup({ onSpawnCall: (params) => params.onStdout?.(validProgressLine) }); - await Effect.runPromise( - restore( - config, - "snapshot-123", - "/tmp/restore-target", - { - organizationId: "org-1", - basePath: "/var/lib/zerobyte/volumes/vol123/_data", - onProgress: (progress) => progressUpdates.push(progress), - }, - mockDeps, - ), + await runRestore( + config, + "snapshot-123", + "/tmp/restore-target", + { + organizationId: "org-1", + basePath: "/var/lib/zerobyte/volumes/vol123/_data", + onProgress: (progress) => progressUpdates.push(progress), + }, + mockDeps, ); expect(progressUpdates).toHaveLength(1); @@ -255,18 +262,16 @@ describe("restore command", () => { }, }); - await Effect.runPromise( - restore( - config, - "snapshot-123", - "/tmp/restore-target", - { - organizationId: "org-1", - basePath: "/var/lib/zerobyte/volumes/vol123/_data", - onProgress: (progress) => progressUpdates.push(progress), - }, - mockDeps, - ), + await runRestore( + config, + "snapshot-123", + "/tmp/restore-target", + { + organizationId: "org-1", + basePath: "/var/lib/zerobyte/volumes/vol123/_data", + onProgress: (progress) => progressUpdates.push(progress), + }, + mockDeps, ); expect(progressUpdates).toHaveLength(0); diff --git a/packages/core/src/restic/commands/restore.ts b/packages/core/src/restic/commands/restore.ts index 01553382..e513764b 100644 --- a/packages/core/src/restic/commands/restore.ts +++ b/packages/core/src/restic/commands/restore.ts @@ -19,7 +19,7 @@ class ResticRestoreCommandError extends Data.TaggedError("ResticRestoreCommandEr message: string; }> {} -const restoreProgressSchema = z.object({ +export const restoreProgressSchema = z.object({ message_type: z.enum(["status", "summary"]), seconds_elapsed: z.number().default(0), percent_done: z.number().default(0), @@ -48,11 +48,14 @@ export const restore = ( signal?: AbortSignal; }, deps: ResticDeps, -) => { - return Effect.tryPromise({ - try: async () => { - const repoUrl = buildRepoUrl(config); - const env = await buildEnv(config, options.organizationId, deps); +): Effect.Effect => { + return Effect.scoped( + Effect.gen(function* () { + const repoUrl = yield* Effect.try(() => buildRepoUrl(config)); + const env = yield* Effect.acquireRelease( + Effect.tryPromise(() => buildEnv(config, options.organizationId, deps)), + (env) => Effect.promise(() => cleanupTemporaryKeys(env, deps)), + ); let restoreArg = snapshotId; @@ -131,23 +134,23 @@ export const restore = ( }, 1000); logger.debug(`Executing: restic ${args.join(" ")}`); - const res = await safeSpawn({ - command: "restic", - args, - env, - signal: options.signal, - onStdout: (data) => { - if (options.onProgress) { - streamProgress(data); - } - }, - }); - - await cleanupTemporaryKeys(env, deps); + const res = yield* Effect.tryPromise(() => + safeSpawn({ + command: "restic", + args, + env, + signal: options.signal, + onStdout: (data) => { + if (options.onProgress) { + streamProgress(data); + } + }, + }), + ); if (res.exitCode !== 0) { logger.error(`Restic restore failed: ${res.error}`); - throw new ResticError(res.exitCode, res.stderr || res.error); + return yield* Effect.fail(new ResticError(res.exitCode, res.stderr || res.error)); } const lastLine = res.summary.trim(); @@ -181,16 +184,19 @@ export const restore = ( ); return result.data; - }, - catch: (error) => { - if (error instanceof ResticError) { - return error; - } + }).pipe( + Effect.catchAll((error): Effect.Effect => { + if (error instanceof ResticError) { + return Effect.fail(error); + } - return new ResticRestoreCommandError({ - cause: error, - message: toMessage(error), - }); - }, - }); + return Effect.fail( + new ResticRestoreCommandError({ + cause: error, + message: toMessage(error), + }), + ); + }), + ), + ); }; diff --git a/packages/core/src/restic/index.ts b/packages/core/src/restic/index.ts index e7928da9..843c27c2 100644 --- a/packages/core/src/restic/index.ts +++ b/packages/core/src/restic/index.ts @@ -1,6 +1,7 @@ export * from "./schemas"; export * from "./restic-dto"; export { ResticError } from "./error"; +export { restoreProgressSchema } from "./commands/restore"; export type { RestoreProgress } from "./commands/restore"; export type {