danvics/zerobyte
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

feat: move restore execution to agent RPC

Browse source
This commit is contained in:
Nicolas Meienberger 2026-05-10 22:25:31 +02:00
parent d4436b0cdc
commit 7ccbb0e2e4
No known key found for this signature in database
17 changed files with 582 additions and 207 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
app/server/modules/agents/agents-manager.ts
View file

@ -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<RestoreCommandPayload, "commandId">) => {
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 () => {

36
app/server/modules/agents/controller/server.ts
View file

@ -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<RestoreCommandPayload, "commandId">,
): Effect.Effect<RestoreCommandResponsePayload | null, Error> =>
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,
};
}

74
app/server/modules/agents/controller/session.ts
View file

@ -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<VolumeCommandResponsePayload, Error>;
description: string;
};
type PendingCommand = { deferred: Deferred.Deferred<any, Error>; description: string };
export type ControllerAgentSessionEvent =
| Exclude<AgentMessage, { type: "volume.commandResult" }>
| {
type: "agent.disconnected";
};
| Exclude<AgentMessage, { type: "volume.commandResult" | "restore.commandResult" }>
| { type: "agent.disconnected" };
export type ControllerAgentSession = {
readonly connectionId: string;
@ -46,6 +43,9 @@ export type ControllerAgentSession = {
sendBackup: (payload: BackupRunPayload) => Effect.Effect<boolean>;
sendBackupCancel: (payload: BackupCancelPayload) => Effect.Effect<boolean>;
runVolumeCommand: (command: VolumeCommand) => Effect.Effect<VolumeCommandResponsePayload, Error>;
runRestoreCommand: (
payload: Omit<RestoreCommandPayload, "commandId">,
) => Effect.Effect<RestoreCommandResponsePayload, Error>;
isReady: () => Effect.Effect<boolean>;
run: Effect.Effect<void, never, Scope.Scope>;
};
@ -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<VolumeCommandResponsePayload, Error>();
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<RestoreCommandResponsePayload, Error>();
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,

2
app/server/modules/backups/backup.helpers.ts
View file

@ -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";

133
app/server/modules/repositories/__tests__/repositories.service.test.ts
View file

@ -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<RepositoryInsert> = {}) => {
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",
}),
"snapshot-restore",
targetPath,
expect.objectContaining({
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" }),
);
});
@ -516,15 +614,16 @@ describe("repositoriesService.restoreSnapshot", () => {
}
expect(restoreMock).toHaveBeenCalledWith(
"local",
expect.objectContaining({
backend: "local",
}),
"snapshot-restore",
targetPath,
expect.objectContaining({
snapshotId: "snapshot-restore",
target: targetPath,
repositoryConfig: expect.objectContaining({ backend: "local" }),
options: expect.objectContaining({
organizationId,
basePath: "/",
}),
}),
);
});
});

1
app/server/modules/repositories/repositories.dto.ts
View file

@ -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(),
});

91
app/server/modules/repositories/repositories.service.ts
View file

@ -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<string, AbortController>();
@ -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,8 +384,13 @@ const restoreSnapshot = async (
snapshotId,
});
const result = await runEffectPromise(
restic.restore(repository.config, snapshotId, target, {
const repositoryConfig = await decryptRepositoryConfig(repository.config);
let result;
if (useControllerLocalRestoreFallback) {
assertAllowedControllerLocalRestoreTarget(target);
result = await runEffectPromise(
restic.restore(repositoryConfig, snapshotId, target, {
basePath,
...options,
organizationId,
@ -383,6 +404,23 @@ const restoreSnapshot = async (
},
}),
);
} else {
const encryptedResticPassword = await resticDeps.getOrganizationResticPassword(organizationId);
const resticPassword = await resticDeps.resolveSecret(encryptedResticPassword);
result = await agentManager.runRestoreCommand(executionAgentId, {
organizationId,
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<RepositoryInsert> = {
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)

8
apps/agent/src/commands/volume.test.ts → apps/agent/src/commands/__tests__/volume.test.ts
View file

@ -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();

2
apps/agent/src/commands/backup-run.ts
View file

@ -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";

0
apps/agent/src/commands/__tests__/backup.helpers.test.ts → apps/agent/src/commands/helpers/__tests__/backup.helpers.test.ts
View file

0
apps/agent/src/commands/backup.helpers.ts → apps/agent/src/commands/helpers/backup.helpers.ts
View file

4
apps/agent/src/commands/index.ts
View file

@ -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);
}

73
apps/agent/src/commands/restore.ts Normal file
View file

@ -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<never>();
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),
}),
),
),
);
};

49
packages/contracts/src/agent-protocol.ts
View file

@ -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<typeof volumeCommandRequestSchema>["p
export type VolumeCommand = z.infer<typeof volumeCommandSchema>;
export type VolumeCommandResult = z.infer<typeof volumeCommandResultSchema>;
export type VolumeCommandResponsePayload = z.infer<typeof volumeCommandResponseSchema>["payload"];
export type RestoreCommandPayload = z.infer<typeof restoreCommandSchema>["payload"];
export type RestoreCommandResponsePayload = z.infer<typeof restoreCommandResultSchema>["payload"];
export type ControllerMessage = z.infer<typeof controllerMessageSchema>;
export type AgentMessage = z.infer<typeof agentMessageSchema>;

61
packages/core/src/restic/commands/__tests__/restore.test.ts
View file

@ -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<SpawnResult>;
onSpawnCall?: (params: SafeSpawnParams) => void | Promise<void>;
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,13 +94,15 @@ afterEach(() => {
vi.restoreAllMocks();
});
const runRestore = (...args: Parameters<typeof restore>) => Effect.runPromise(restore(...args));
const runRestoreError = (...args: Parameters<typeof restore>) => 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(
await runRestore(
config,
"snapshot-456",
"/tmp/restore-target",
@ -107,7 +114,6 @@ describe("restore command", () => {
],
},
mockDeps,
),
);
expect(getRestoreArg()).toBe("snapshot-456:/var/lib/zerobyte/volumes/vol123/_data");
@ -117,8 +123,7 @@ 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(
await runRestore(
config,
"snapshot-single-file",
"/tmp/restore-target",
@ -128,7 +133,6 @@ describe("restore command", () => {
selectedItemKind: "file",
},
mockDeps,
),
);
expect(getRestoreArg()).toBe("snapshot-single-file:/var/lib/zerobyte/volumes/vol123/_data/archive");
@ -138,8 +142,7 @@ describe("restore command", () => {
test("treats flag-like snapshot IDs as positional restore args", async () => {
const { getArgs, getRestoreArg } = setup();
await Effect.runPromise(
restore(
await runRestore(
config,
"--help",
"/tmp/restore-target",
@ -148,7 +151,6 @@ describe("restore command", () => {
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(
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(
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(error).toBeInstanceOf(ResticError);
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 Effect.runPromise(
restore(
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,8 +233,7 @@ describe("restore command", () => {
const progressUpdates: unknown[] = [];
setup({ onSpawnCall: (params) => params.onStdout?.(validProgressLine) });
await Effect.runPromise(
restore(
await runRestore(
config,
"snapshot-123",
"/tmp/restore-target",
@ -235,7 +243,6 @@ describe("restore command", () => {
onProgress: (progress) => progressUpdates.push(progress),
},
mockDeps,
),
);
expect(progressUpdates).toHaveLength(1);
@ -255,8 +262,7 @@ describe("restore command", () => {
},
});
await Effect.runPromise(
restore(
await runRestore(
config,
"snapshot-123",
"/tmp/restore-target",
@ -266,7 +272,6 @@ describe("restore command", () => {
onProgress: (progress) => progressUpdates.push(progress),
},
mockDeps,
),
);
expect(progressUpdates).toHaveLength(0);

42
packages/core/src/restic/commands/restore.ts
View file

@ -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<ResticRestoreOutputDto, ResticError | ResticRestoreCommandError> => {
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,7 +134,8 @@ export const restore = (
}, 1000);
logger.debug(`Executing: restic ${args.join(" ")}`);
const res = await safeSpawn({
const res = yield* Effect.tryPromise(() =>
safeSpawn({
command: "restic",
args,
env,
@ -141,13 +145,12 @@ export const restore = (
streamProgress(data);
}
},
});
await cleanupTemporaryKeys(env, deps);
}),
);
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) => {
}).pipe(
Effect.catchAll((error): Effect.Effect<never, ResticError | ResticRestoreCommandError> => {
if (error instanceof ResticError) {
return error;
return Effect.fail(error);
}
return new ResticRestoreCommandError({
return Effect.fail(
new ResticRestoreCommandError({
cause: error,
message: toMessage(error),
});
},
});
}),
);
}),
),
);
};

1
packages/core/src/restic/index.ts
View file

@ -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 {