diff --git a/package-lock.json b/package-lock.json index 80cee7d3..ae272327 100644 --- a/package-lock.json +++ b/package-lock.json @@ -12,7 +12,7 @@ "@dfinity/agent": "^3.2.6", "@dfinity/auth-client": "^3.2.6", "@dfinity/candid": "^3.2.6", - "@dfinity/ic-management": "^7.0.1", + "@dfinity/ic-management": "^7.0.2-beta-2025-09-30", "@dfinity/identity": "^3.2.6", "@dfinity/principal": "^3.2.6", "@dfinity/zod-schemas": "^2.1.0", @@ -614,15 +614,15 @@ "license": "Apache-2.0" }, "node_modules/@dfinity/ic-management": { - "version": "7.0.1", - "resolved": "https://registry.npmjs.org/@dfinity/ic-management/-/ic-management-7.0.1.tgz", - "integrity": "sha512-zatpUqzf9k3bYkLeikurwYXOwzvJgKf+y8+u1Vfb/cS58gOu4OMB2buc9/Rw5frtN9TRbe/vsTQSodnLLlbZIw==", + "version": "7.0.2-beta-2025-09-30", + "resolved": "https://registry.npmjs.org/@dfinity/ic-management/-/ic-management-7.0.2-beta-2025-09-30.tgz", + "integrity": "sha512-UcNNIFKulBW+hd9zgDIL2pzPI/MJ6k9rg0uDbPS66q0lVebZ3xvNUP83JTwn143K2PNsYr3SBt9zC/l1OxtvYQ==", "license": "Apache-2.0", "peerDependencies": { - "@dfinity/agent": "^3", - "@dfinity/candid": "^3", - "@dfinity/principal": "^3", - "@dfinity/utils": "^3" + "@dfinity/agent": "*", + "@dfinity/candid": "*", + "@dfinity/principal": "*", + "@dfinity/utils": "*" } }, "node_modules/@dfinity/identity": { @@ -6855,9 +6855,9 @@ "integrity": "sha512-GPJpH73kDEKbUBdUjY80lz7cq9l0vm1h/7ppejPV6O0ZTqCLrYspssYvqjRmK4aNnJ/SKXsP0rg9LYX7zpegaA==" }, "@dfinity/ic-management": { - "version": "7.0.1", - "resolved": "https://registry.npmjs.org/@dfinity/ic-management/-/ic-management-7.0.1.tgz", - "integrity": "sha512-zatpUqzf9k3bYkLeikurwYXOwzvJgKf+y8+u1Vfb/cS58gOu4OMB2buc9/Rw5frtN9TRbe/vsTQSodnLLlbZIw==", + "version": "7.0.2-beta-2025-09-30", + "resolved": "https://registry.npmjs.org/@dfinity/ic-management/-/ic-management-7.0.2-beta-2025-09-30.tgz", + "integrity": "sha512-UcNNIFKulBW+hd9zgDIL2pzPI/MJ6k9rg0uDbPS66q0lVebZ3xvNUP83JTwn143K2PNsYr3SBt9zC/l1OxtvYQ==", "requires": {} }, "@dfinity/identity": { diff --git a/package.json b/package.json index d5e95674..fc580b22 100644 --- a/package.json +++ b/package.json @@ -26,7 +26,7 @@ "@dfinity/agent": "^3.2.6", "@dfinity/auth-client": "^3.2.6", "@dfinity/candid": "^3.2.6", - "@dfinity/ic-management": "^7.0.1", + "@dfinity/ic-management": "^7.0.2-beta-2025-09-30", "@dfinity/identity": "^3.2.6", "@dfinity/principal": "^3.2.6", "@dfinity/zod-schemas": "^2.1.0", diff --git a/src/api/ic.api.ts b/src/api/ic.api.ts index fc3b4fed..12039da8 100644 --- a/src/api/ic.api.ts +++ b/src/api/ic.api.ts @@ -1,8 +1,17 @@ import {ICManagementCanister} from '@dfinity/ic-management'; import type { list_canister_snapshots_result, - snapshot_id + read_canister_snapshot_data_response, + snapshot_id, + upload_canister_snapshot_metadata_response } from '@dfinity/ic-management/dist/candid/ic-management'; +import { + type ReadCanisterSnapshotMetadataParams, + type SnapshotParams, + type UploadCanisterSnapshotDataParams, + type UploadCanisterSnapshotMetadataParams +} from '@dfinity/ic-management/dist/types/types/snapshot.params'; +import type {ReadCanisterSnapshotMetadataResponse} from '@dfinity/ic-management/dist/types/types/snapshot.responses'; import type {Principal} from '@dfinity/principal'; import {initAgent} from './agent.api'; @@ -76,3 +85,51 @@ export const deleteCanisterSnapshot = async (params: { await deleteCanisterSnapshot(params); }; + +export const readCanisterSnapshotMetadata = async ( + params: SnapshotParams +): Promise => { + const agent = await initAgent(); + + const {readCanisterSnapshotMetadata} = ICManagementCanister.create({ + agent + }); + + return await readCanisterSnapshotMetadata(params); +}; + +export const readCanisterSnapshotData = async ( + params: ReadCanisterSnapshotMetadataParams +): Promise => { + const agent = await initAgent(); + + const {readCanisterSnapshotData} = ICManagementCanister.create({ + agent + }); + + return await readCanisterSnapshotData(params); +}; + +export const uploadCanisterSnapshotMetadata = async ( + params: UploadCanisterSnapshotMetadataParams +): Promise => { + const agent = await initAgent(); + + const {uploadCanisterSnapshotMetadata} = ICManagementCanister.create({ + agent + }); + + return await uploadCanisterSnapshotMetadata(params); +}; + +export const uploadCanisterSnapshotData = async ( + params: UploadCanisterSnapshotDataParams +): Promise => { + const agent = await initAgent(); + + const {uploadCanisterSnapshotData} = ICManagementCanister.create({ + agent + }); + + await uploadCanisterSnapshotData(params); +}; diff --git a/src/commands/snapshot.ts b/src/commands/snapshot.ts index 3ff62092..30b1c3ad 100644 --- a/src/commands/snapshot.ts +++ b/src/commands/snapshot.ts @@ -1,20 +1,27 @@ import {nextArg} from '@junobuild/cli-tools'; import {red} from 'kleur'; import {logHelpSnapshot} from '../help/snapshot.help'; +import {logHelpSnapshotUpload} from '../help/snapshot.upload.help'; import { createSnapshotMissionControl, deleteSnapshotMissionControl, - restoreSnapshotMissionControl + downloadSnapshotMissionControl, + restoreSnapshotMissionControl, + uploadSnapshotMissionControl } from '../services/modules/snapshot/snapshot.mission-control.services'; import { createSnapshotOrbiter, deleteSnapshotOrbiter, - restoreSnapshotOrbiter + downloadSnapshotOrbiter, + restoreSnapshotOrbiter, + uploadSnapshotOrbiter } from '../services/modules/snapshot/snapshot.orbiter.services'; import { createSnapshotSatellite, deleteSnapshotSatellite, - restoreSnapshotSatellite + downloadSnapshotSatellite, + restoreSnapshotSatellite, + uploadSnapshotSatellite } from '../services/modules/snapshot/snapshot.satellite.services'; export const snapshot = async (args?: string[]) => { @@ -45,6 +52,22 @@ export const snapshot = async (args?: string[]) => { orbiterFn: deleteSnapshotOrbiter }); break; + case 'download': + await executeSnapshotFn({ + args, + satelliteFn: downloadSnapshotSatellite, + missionControlFn: downloadSnapshotMissionControl, + orbiterFn: downloadSnapshotOrbiter + }); + break; + case 'upload': + await executeSnapshotFn({ + args, + satelliteFn: uploadSnapshotSatellite, + missionControlFn: uploadSnapshotMissionControl, + orbiterFn: uploadSnapshotOrbiter + }); + break; default: console.log(red('Unknown subcommand.')); logHelpSnapshot(args); @@ -58,27 +81,39 @@ const executeSnapshotFn = async ({ orbiterFn }: { args?: string[]; - satelliteFn: () => Promise; - missionControlFn: () => Promise; - orbiterFn: () => Promise; + satelliteFn: (args?: string[]) => Promise; + missionControlFn: (args?: string[]) => Promise; + orbiterFn: (args?: string[]) => Promise; }) => { const target = nextArg({args, option: '-t'}) ?? nextArg({args, option: '--target'}); switch (target) { case 's': case 'satellite': - await satelliteFn(); + await satelliteFn(args); break; case 'm': case 'mission-control': - await missionControlFn(); + await missionControlFn(args); break; case 'o': case 'orbiter': - await orbiterFn(); + await orbiterFn(args); break; default: console.log(red('Unknown target.')); logHelpSnapshot(args); } }; + +export const helpSnapshot = (args?: string[]) => { + const [subCommand] = args ?? []; + + switch (subCommand) { + case 'upload': + logHelpSnapshotUpload(args); + break; + default: + logHelpSnapshot(args); + } +}; diff --git a/src/constants/help.constants.ts b/src/constants/help.constants.ts index 7b320fb2..a570bb55 100644 --- a/src/constants/help.constants.ts +++ b/src/constants/help.constants.ts @@ -48,6 +48,8 @@ export const CHANGES_LIST_DESCRIPTION = 'List all submitted or applied changes.' export const CHANGES_APPLY_DESCRIPTION = 'Apply a submitted change.'; export const CHANGES_REJECT_DESCRIPTION = 'Reject a change.'; +export const SNAPSHOT_UPLOAD_DESCRIPTION = 'Upload a snapshot from offline files.'; + export const OPTION_KEEP_STAGED = `${yellow('-k, --keep-staged')} Keep staged assets in memory after applying the change.`; export const OPTION_HASH = `${yellow('--hash')} The expected hash of all included changes (for verification).`; export const OPTION_HELP = `${yellow('-h, --help')} Output usage information.`; diff --git a/src/constants/snapshot.constants.ts b/src/constants/snapshot.constants.ts new file mode 100644 index 00000000..dd56969f --- /dev/null +++ b/src/constants/snapshot.constants.ts @@ -0,0 +1,7 @@ +import {join} from 'node:path'; + +export const SNAPSHOTS_PATH = join(process.cwd(), '.snapshots'); + +// https://forum.dfinity.org/t/canister-snapshot-up-download/57397?u=peterparker +// Same value as INSTALL_MAX_CHUNK_SIZE in @junobuild/admin +export const SNAPSHOT_MAX_CHUNK_SIZE = 1_000_000n; diff --git a/src/help/snapshot.help.ts b/src/help/snapshot.help.ts index 8d687a63..6e57624d 100644 --- a/src/help/snapshot.help.ts +++ b/src/help/snapshot.help.ts @@ -1,5 +1,10 @@ import {cyan, green, magenta, yellow} from 'kleur'; -import {OPTIONS_ENV, OPTION_HELP, SNAPSHOT_DESCRIPTION} from '../constants/help.constants'; +import { + OPTIONS_ENV, + OPTION_HELP, + SNAPSHOT_DESCRIPTION, + SNAPSHOT_UPLOAD_DESCRIPTION +} from '../constants/help.constants'; import {helpOutput} from './common.help'; import {TITLE} from './help'; import {TARGET_OPTION_NOTE, targetOption} from './target.help'; @@ -8,8 +13,10 @@ const usage = `Usage: ${green('juno')} ${cyan('snapshot')} ${magenta(' { + console.log(helpOutput(args) === 'doc' ? doc : help); +}; diff --git a/src/index.ts b/src/index.ts index b62b4399..e85a213c 100644 --- a/src/index.ts +++ b/src/index.ts @@ -12,7 +12,7 @@ import {functions, helpFunctions} from './commands/functions'; import {helpHosting, hosting} from './commands/hosting'; import {open} from './commands/open'; import {helpRun, run as runCmd} from './commands/run'; -import {snapshot} from './commands/snapshot'; +import {helpSnapshot, snapshot} from './commands/snapshot'; import {startStop} from './commands/start-stop'; import {status} from './commands/status'; import {upgrade} from './commands/upgrade'; @@ -22,7 +22,6 @@ import {help} from './help/help'; import {logHelpLogin} from './help/login.help'; import {logHelpLogout} from './help/logout.help'; import {logHelpOpen} from './help/open.help'; -import {logHelpSnapshot} from './help/snapshot.help'; import {logHelpStart} from './help/start.help'; import {logHelpStatus} from './help/status.help'; import {logHelpStop} from './help/stop.help'; @@ -94,7 +93,7 @@ export const run = async () => { helpFunctions(args); break; case 'snapshot': - logHelpSnapshot(args); + helpSnapshot(args); break; case 'init': helpInit(args); diff --git a/src/schema/snapshot.schema.ts b/src/schema/snapshot.schema.ts new file mode 100644 index 00000000..12b78100 --- /dev/null +++ b/src/schema/snapshot.schema.ts @@ -0,0 +1,65 @@ +import * as z from 'zod'; + +const Uint8ArrayLike = z.instanceof(Uint8Array) as z.ZodType; + +// A Zod schema for the ic-management ReadCanisterSnapshotMetadataResponse type +export const ReadCanisterSnapshotMetadataResponseSchema = z.strictObject({ + globals: z.array( + z.union([ + z.object({f32: z.number()}), + z.object({f64: z.number()}), + z.object({i32: z.number()}), + z.object({i64: z.bigint()}), + z.object({v128: z.bigint()}) + ]) + ), + canisterVersion: z.bigint(), + source: z.union([ + z.object({metadataUpload: z.unknown()}), + z.object({takenFromCanister: z.unknown()}) + ]), + certifiedData: z.union([Uint8ArrayLike, z.array(z.number())]), + globalTimer: z.union([z.object({active: z.bigint()}), z.object({inactive: z.null()})]).optional(), + onLowWasmMemoryHookStatus: z + .union([ + z.object({conditionNotSatisfied: z.null()}), + z.object({executed: z.null()}), + z.object({ready: z.null()}) + ]) + .optional(), + wasmModuleSize: z.bigint(), + stableMemorySize: z.bigint(), + wasmChunkStore: z.array( + z.object({ + hash: z.union([Uint8ArrayLike, z.array(z.number())]) + }) + ), + takenAtTimestamp: z.bigint(), + wasmMemorySize: z.bigint() +}); + +export const SnapshotFilenameSchema = z.enum([ + 'wasm-code.bin', + 'heap.bin', + 'stable.bin', + 'chunks-store.bin' +]); + +export const SnapshotFileSchema = z.strictObject({ + filename: SnapshotFilenameSchema, + size: z.bigint(), + hash: z.hash('sha256') +}); + +const SnapshotDataSchema = z.strictObject({ + wasmModule: SnapshotFileSchema.nullable(), + wasmMemory: SnapshotFileSchema.nullable(), + stableMemory: SnapshotFileSchema.nullable(), + wasmChunkStore: SnapshotFileSchema.nullable() +}); + +export const SnapshotMetadataSchema = z.strictObject({ + snapshotId: z.string(), + data: SnapshotDataSchema, + metadata: ReadCanisterSnapshotMetadataResponseSchema +}); diff --git a/src/services/modules/snapshot/snapshot.download.services.ts b/src/services/modules/snapshot/snapshot.download.services.ts new file mode 100644 index 00000000..71d8574f --- /dev/null +++ b/src/services/modules/snapshot/snapshot.download.services.ts @@ -0,0 +1,404 @@ +import {type CanisterSnapshotMetadataKind, encodeSnapshotId} from '@dfinity/ic-management'; +import type {ReadCanisterSnapshotMetadataResponse} from '@dfinity/ic-management/dist/types/types/snapshot.responses'; +import {arrayOfNumberToUint8Array, jsonReplacer} from '@dfinity/utils'; +import {red} from 'kleur'; +import {createHash} from 'node:crypto'; +import {createWriteStream, existsSync, statSync} from 'node:fs'; +import {mkdir, writeFile} from 'node:fs/promises'; +import {join, relative} from 'node:path'; +import {Readable, Transform} from 'node:stream'; +import {pipeline} from 'node:stream/promises'; +import ora from 'ora'; +import {readCanisterSnapshotData, readCanisterSnapshotMetadata} from '../../../api/ic.api'; +import {SNAPSHOTS_PATH} from '../../../constants/snapshot.constants'; +import {type AssetKey} from '../../../types/asset-key'; +import { + type DownloadSnapshotParams, + type SnapshotBatchResult, + type SnapshotFile, + type SnapshotFilename, + type SnapshotLog, + type SnapshotMetadata +} from '../../../types/snapshot'; +import {displaySegment} from '../../../utils/display.utils'; +import {type BuildChunkFn, prepareDataChunks} from '../../../utils/snapshot.utils'; + +type DataChunk = CanisterSnapshotMetadataKind; +type DownloadedChunk = Uint8Array; + +class SnapshotFsFolderError extends Error { + constructor(public readonly folder: string) { + super(); + } +} + +class SnapshotFsSizeError extends Error { + constructor( + public readonly filename: string, + public readonly expectedSize: bigint, + public readonly downloadedSize: bigint + ) { + super(); + } +} + +interface DownloadSnapshotBatchResult extends SnapshotBatchResult { + downloadedChunks: DownloadedChunk[]; +} + +export const downloadExistingSnapshot = async ({ + segment, + ...params +}: DownloadSnapshotParams & { + segment: AssetKey; +}): Promise => { + const spinner = ora().start(); + + try { + const result = await downloadSnapshotMetadataAndMemory({ + ...params, + log: (text) => (spinner.text = text) + }); + + spinner.stop(); + + const {snapshotIdText, folder} = result; + + console.log( + `✅ The snapshot ${snapshotIdText} for your ${displaySegment(segment)} has been downloaded.` + ); + console.log(`🗂️ Files saved to ${folder}`); + } catch (error: unknown) { + spinner.stop(); + + if (error instanceof SnapshotFsFolderError) { + console.log( + `${red('Cannot proceed with download.')}\nDestination ${error.folder} already exists.` + ); + return; + } + + if (error instanceof SnapshotFsSizeError) { + console.log( + `${red('Download size mismatch.')}\n${error.filename}: expected ${error.expectedSize} bytes, got ${error.downloadedSize} bytes.` + ); + return; + } + + throw error; + } +}; + +const downloadSnapshotMetadataAndMemory = async ({ + snapshotId, + log, + ...rest +}: DownloadSnapshotParams & SnapshotLog): Promise<{ + status: 'success'; + snapshotIdText: string; + folder: string; +}> => { + const snapshotIdText = `0x${encodeSnapshotId(snapshotId)}`; + const folder = join(SNAPSHOTS_PATH, snapshotIdText); + + if (existsSync(folder)) { + throw new SnapshotFsFolderError(folder); + } + + await mkdir(folder, {recursive: true}); + + // 1. Load the snapshot metadata + const {metadata} = await loadMetadata({snapshotId, log, snapshotIdText, ...rest}); + + // 2. Download the snapshot data (WASM program code, heap and stable memory, WASM chunks store) + const {wasmModuleSize, wasmMemorySize, stableMemorySize, wasmChunkStore} = metadata; + + const wasmModuleResult = await assertSizeAndDownloadChunks({ + folder, + filename: 'wasm-code.bin', + snapshotId, + size: wasmModuleSize, + build: (param) => ({wasmModule: param}), + log, + ...rest + }); + + const wasmMemoryResult = await assertSizeAndDownloadChunks({ + folder, + filename: 'heap.bin', + snapshotId, + size: wasmMemorySize, + build: (param) => ({wasmMemory: param}), + log, + ...rest + }); + + const stableMemoryResult = await assertSizeAndDownloadChunks({ + folder, + filename: 'stable.bin', + snapshotId, + size: stableMemorySize, + build: (param) => ({stableMemory: param}), + log, + ...rest + }); + + const wasmChunkStoreResult = await assertAndDownloadWasmChunks({ + folder, + filename: 'chunks-store.bin', + snapshotId, + wasmChunkStore, + log, + ...rest + }); + + // 3. Save the metadata of the offline snapshot + await saveMetadata({ + log, + folder, + metadata: { + snapshotId: snapshotIdText, + metadata, + data: { + wasmModule: wasmModuleResult.status === 'ok' ? wasmModuleResult.snapshotFile : null, + wasmMemory: wasmMemoryResult.status === 'ok' ? wasmMemoryResult.snapshotFile : null, + stableMemory: stableMemoryResult.status === 'ok' ? stableMemoryResult.snapshotFile : null, + wasmChunkStore: + wasmChunkStoreResult.status === 'ok' ? wasmChunkStoreResult.snapshotFile : null + } + } + }); + + return {status: 'success', snapshotIdText, folder}; +}; + +const loadMetadata = async ({ + snapshotIdText, + log, + ...rest +}: DownloadSnapshotParams & {snapshotIdText: string} & SnapshotLog): Promise<{ + metadata: ReadCanisterSnapshotMetadataResponse; +}> => { + log(`Downloading snapshot metadata ${snapshotIdText}...`); + + const metadata = await readCanisterSnapshotMetadata(rest); + return {metadata}; +}; + +const saveMetadata = async ({ + folder, + log, + metadata +}: {folder: string; metadata: SnapshotMetadata} & SnapshotLog) => { + log(`Saving snapshot metadata...`); + + const destination = join(folder, 'metadata.json'); + await writeFile(destination, JSON.stringify(metadata, jsonReplacer, 2), 'utf-8'); +}; + +const assertSizeAndDownloadChunks = async ({ + size, + log, + filename, + ...params +}: DownloadSnapshotParams & { + folder: string; + filename: SnapshotFilename; + size: bigint; + build: BuildChunkFn; +} & SnapshotLog): Promise<{status: 'ok'; snapshotFile: SnapshotFile} | {status: 'skip'}> => { + if (size === 0n) { + log(`No chunks to download for ${filename} (size = 0). Skipping.`); + await sleep(); + return {status: 'skip'}; + } + + const {size: downloadedSize, hash} = await downloadMemoryChunks({ + size, + log, + filename, + ...params + }); + + if (downloadedSize !== size) { + throw new SnapshotFsSizeError(filename, size, downloadedSize); + } + + return { + status: 'ok', + snapshotFile: { + filename, + size: downloadedSize, + hash + } + }; +}; + +// eslint-disable-next-line promise/avoid-new +const sleep = async () => await new Promise((resolve) => setTimeout(resolve, 2500)); + +const assertAndDownloadWasmChunks = async ({ + wasmChunkStore, + log, + ...params +}: DownloadSnapshotParams & {folder: string; filename: SnapshotFilename} & Pick< + ReadCanisterSnapshotMetadataResponse, + 'wasmChunkStore' + > & + SnapshotLog): Promise<{status: 'ok'; snapshotFile: SnapshotFile} | {status: 'skip'}> => { + if (wasmChunkStore.length === 0) { + log('Nothing to download from the WASM chunks store (length = 0). Skipping.'); + await sleep(); + return {status: 'skip'}; + } + + const snapshotFile = await downloadWasmChunks({ + wasmChunkStore, + log, + ...params + }); + + return {status: 'ok', snapshotFile}; +}; + +const downloadMemoryChunks = async ({ + size, + build, + log, + folder, + filename, + ...params +}: DownloadSnapshotParams & { + folder: string; + filename: SnapshotFilename; + size: bigint; + build: BuildChunkFn; +} & SnapshotLog): Promise => { + const {chunks} = prepareDataChunks({size, build}); + + const readable = Readable.from( + batchDownloadChunks({ + chunks, + ...params + }) + ); + + return await downloadAndWrite({readable, folder, filename, log}); +}; + +const downloadWasmChunks = async ({ + wasmChunkStore, + log, + folder, + filename, + ...params +}: DownloadSnapshotParams & {folder: string; filename: SnapshotFilename} & Pick< + ReadCanisterSnapshotMetadataResponse, + 'wasmChunkStore' + > & + SnapshotLog): Promise => { + const readable = Readable.from( + batchDownloadChunks({ + chunks: wasmChunkStore.map((chunk, orderId) => ({wasmChunk: chunk, orderId})), + limit: 12, + ...params + }) + ); + + return await downloadAndWrite({readable, folder, filename, log}); +}; + +const downloadAndWrite = async ({ + readable, + folder, + filename, + log +}: { + readable: Readable; + folder: string; + filename: SnapshotFilename; +} & SnapshotLog): Promise => { + // Note: we would not win much at gzipping the data. + const destination = join(folder, filename); + + log(`Downloading chunks to ${relative(process.cwd(), destination)}...`); + + // A transformer use to flatten the back of chunks when writing to the file + const transformer = new Transform({ + objectMode: true, + writableObjectMode: true, + readableObjectMode: false, + transform({downloadedChunks: chunks, progress}: DownloadSnapshotBatchResult, _enc, cb) { + try { + log(`Chunks ${progress.done}/${progress.total} downloaded. Continuing...`); + + for (const chunk of chunks) { + this.push(chunk); + } + cb(); + } catch (err: unknown) { + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + cb(err as Error); + } + } + }); + + // We want to compute a sha256 to assert the file in the upload process + const hash = createHash('sha256'); + + const hasher = new Transform({ + transform(chunk: DownloadedChunk, _enc, cb) { + hash.update(chunk); + cb(null, chunk); + } + }); + + await pipeline(readable, transformer, hasher, createWriteStream(destination)); + + return { + filename, + size: BigInt(statSync(destination).size), + hash: hash.digest('hex') + }; +}; + +async function* batchDownloadChunks({ + chunks, + limit = 20, + ...params +}: DownloadSnapshotParams & { + chunks: DataChunk[]; + limit?: number; +}): AsyncGenerator { + const total = chunks.length; + + for (let i = 0; i < total; i = i + limit) { + const batch = chunks.slice(i, i + limit); + const downloadedChunks = await Promise.all( + batch.map( + async (requestChunk) => + await downloadChunk({ + ...params, + chunk: requestChunk + }) + ) + ); + yield {downloadedChunks, progress: {index: i, done: Math.min(i + limit, total), total}}; + } +} + +const downloadChunk = async ({ + chunk: kind, + ...rest +}: DownloadSnapshotParams & { + chunk: DataChunk; +}): Promise => { + const {chunk: downloadedChunk} = await readCanisterSnapshotData({ + ...rest, + kind + }); + + return downloadedChunk instanceof Uint8Array + ? downloadedChunk + : arrayOfNumberToUint8Array(downloadedChunk); +}; diff --git a/src/services/modules/snapshot/snapshot.mission-control.services.ts b/src/services/modules/snapshot/snapshot.mission-control.services.ts index d064dbfb..36b64417 100644 --- a/src/services/modules/snapshot/snapshot.mission-control.services.ts +++ b/src/services/modules/snapshot/snapshot.mission-control.services.ts @@ -2,7 +2,13 @@ import {isNullish} from '@dfinity/utils'; import {red} from 'kleur'; import {getCliMissionControl} from '../../../configs/cli.config'; import type {AssetKey} from '../../../types/asset-key'; -import {createSnapshot, deleteSnapshot, restoreSnapshot} from './snapshot.services'; +import { + createSnapshot, + deleteSnapshot, + downloadSnapshot, + restoreSnapshot, + uploadSnapshot +} from './snapshot.services'; export const createSnapshotMissionControl = async () => { await executeSnapshotFn({ @@ -22,6 +28,20 @@ export const deleteSnapshotMissionControl = async () => { }); }; +export const downloadSnapshotMissionControl = async () => { + await executeSnapshotFn({ + fn: downloadSnapshot + }); +}; + +export const uploadSnapshotMissionControl = async (args?: string[]) => { + await executeSnapshotFn({ + fn: async (params) => { + await uploadSnapshot({...params, args}); + } + }); +}; + const executeSnapshotFn = async ({ fn }: { diff --git a/src/services/modules/snapshot/snapshot.orbiter.services.ts b/src/services/modules/snapshot/snapshot.orbiter.services.ts index b4df3aaf..d2bb6c17 100644 --- a/src/services/modules/snapshot/snapshot.orbiter.services.ts +++ b/src/services/modules/snapshot/snapshot.orbiter.services.ts @@ -1,6 +1,12 @@ import {getCliOrbiters} from '../../../configs/cli.config'; import type {AssetKey} from '../../../types/asset-key'; -import {createSnapshot, deleteSnapshot, restoreSnapshot} from './snapshot.services'; +import { + createSnapshot, + deleteSnapshot, + downloadSnapshot, + restoreSnapshot, + uploadSnapshot +} from './snapshot.services'; export const createSnapshotOrbiter = async () => { await executeSnapshotFn({ @@ -20,6 +26,20 @@ export const deleteSnapshotOrbiter = async () => { }); }; +export const downloadSnapshotOrbiter = async () => { + await executeSnapshotFn({ + fn: downloadSnapshot + }); +}; + +export const uploadSnapshotOrbiter = async (args?: string[]) => { + await executeSnapshotFn({ + fn: async (params) => { + await uploadSnapshot({...params, args}); + } + }); +}; + const executeSnapshotFn = async ({ fn }: { diff --git a/src/services/modules/snapshot/snapshot.satellite.services.ts b/src/services/modules/snapshot/snapshot.satellite.services.ts index 2e1dbdd9..53af5658 100644 --- a/src/services/modules/snapshot/snapshot.satellite.services.ts +++ b/src/services/modules/snapshot/snapshot.satellite.services.ts @@ -3,7 +3,13 @@ import {noJunoConfig} from '../../../configs/juno.config'; import type {AssetKey} from '../../../types/asset-key'; import {consoleNoConfigFound} from '../../../utils/msg.utils'; import {assertConfigAndLoadSatelliteContext} from '../../../utils/satellite.utils'; -import {createSnapshot, deleteSnapshot, restoreSnapshot} from './snapshot.services'; +import { + createSnapshot, + deleteSnapshot, + downloadSnapshot, + restoreSnapshot, + uploadSnapshot +} from './snapshot.services'; export const createSnapshotSatellite = async () => { await executeSnapshotFn({ @@ -23,6 +29,20 @@ export const deleteSnapshotSatellite = async () => { }); }; +export const downloadSnapshotSatellite = async () => { + await executeSnapshotFn({ + fn: downloadSnapshot + }); +}; + +export const uploadSnapshotSatellite = async (args?: string[]) => { + await executeSnapshotFn({ + fn: async (params) => { + await uploadSnapshot({...params, args}); + } + }); +}; + const executeSnapshotFn = async ({ fn }: { diff --git a/src/services/modules/snapshot/snapshot.services.ts b/src/services/modules/snapshot/snapshot.services.ts index dff42afa..57bda4ab 100644 --- a/src/services/modules/snapshot/snapshot.services.ts +++ b/src/services/modules/snapshot/snapshot.services.ts @@ -1,8 +1,10 @@ import type {snapshot_id} from '@dfinity/ic-management'; import {encodeSnapshotId} from '@dfinity/ic-management'; import {Principal} from '@dfinity/principal'; -import {isNullish, nonNullish} from '@dfinity/utils'; -import {red} from 'kleur'; +import {isEmptyString, isNullish, nonNullish} from '@dfinity/utils'; +import {nextArg} from '@junobuild/cli-tools'; +import {red, yellow} from 'kleur'; +import {existsSync, lstatSync} from 'node:fs'; import ora from 'ora'; import { deleteCanisterSnapshot, @@ -13,6 +15,8 @@ import { import type {AssetKey} from '../../../types/asset-key'; import {displaySegment} from '../../../utils/display.utils'; import {confirmAndExit} from '../../../utils/prompt.utils'; +import {downloadExistingSnapshot} from './snapshot.download.services'; +import {uploadExistingSnapshot} from './snapshot.upload.services'; export const createSnapshot = async ({ canisterId: cId, @@ -23,17 +27,11 @@ export const createSnapshot = async ({ }) => { const canisterId = Principal.fromText(cId); - const existingSnapshotId = await loadSnapshot({canisterId}); - - if (nonNullish(existingSnapshotId)) { - await confirmAndExit( - `A snapshot for your ${displaySegment(segment)} already exists with ID 0x${encodeSnapshotId(existingSnapshotId)}. Do you want to overwrite it?` - ); - } + const {snapshotId} = await loadSnapshotAndAssertOverwrite({canisterId, segment}); await takeSnapshot({ canisterId, - snapshotId: existingSnapshotId, + snapshotId, segment }); }; @@ -47,13 +45,14 @@ export const restoreSnapshot = async ({ }) => { const canisterId = Principal.fromText(cId); - const existingSnapshotId = await loadSnapshot({canisterId}); + const result = await loadSnapshotAndAssertExist({canisterId, segment}); - if (isNullish(existingSnapshotId)) { - console.log(red(`No snapshot found for your ${displaySegment(segment)}.`)); + if (result.result === 'not_found') { return; } + const {snapshotId: existingSnapshotId} = result; + await confirmAndExit( `Restoring the snapshot 0x${encodeSnapshotId(existingSnapshotId)} will permanently overwrite the current state of your ${displaySegment(segment)}. Are you sure you want to proceed?` ); @@ -74,13 +73,14 @@ export const deleteSnapshot = async ({ }) => { const canisterId = Principal.fromText(cId); - const existingSnapshotId = await loadSnapshot({canisterId}); + const result = await loadSnapshotAndAssertExist({canisterId, segment}); - if (isNullish(existingSnapshotId)) { - console.log(red(`No snapshot found for your ${displaySegment(segment)}.`)); + if (result.result === 'not_found') { return; } + const {snapshotId: existingSnapshotId} = result; + await confirmAndExit( `Are you sure you want to delete the snapshot 0x${encodeSnapshotId(existingSnapshotId)} of your ${displaySegment(segment)}?` ); @@ -92,6 +92,73 @@ export const deleteSnapshot = async ({ }); }; +export const downloadSnapshot = async ({ + canisterId: cId, + segment +}: { + canisterId: string; + segment: AssetKey; +}) => { + const canisterId = Principal.fromText(cId); + + const result = await loadSnapshotAndAssertExist({canisterId, segment}); + + if (result.result === 'not_found') { + return; + } + + const {snapshotId: existingSnapshotId} = result; + + await downloadExistingSnapshot({ + canisterId, + snapshotId: existingSnapshotId, + segment + }); +}; + +export const uploadSnapshot = async ({ + canisterId: cId, + segment, + args +}: { + canisterId: string; + segment: AssetKey; + args?: string[]; +}) => { + const canisterId = Principal.fromText(cId); + + const folder = nextArg({args, option: '--dir'}); + + if (isEmptyString(folder)) { + console.log( + `You did not provide a ${yellow('directory')} that contains metadata.json and chunks to upload.` + ); + return; + } + + if (!existsSync(folder)) { + console.log(`The directory ${yellow('directory')} does not exist.`); + return; + } + + if (!lstatSync(folder).isDirectory()) { + console.log(red(`${folder} is not a directory.`)); + return; + } + + // TODO: extract assertions + // TODO: more assertion like is there a metadata.json and chunk files + + const {snapshotId} = await loadSnapshotAndAssertOverwrite({canisterId, segment}); + + await uploadExistingSnapshot({ + canisterId, + snapshotId, + segment, + folder + }); +}; + const restoreExistingSnapshot = async ({ segment, ...rest @@ -166,3 +233,38 @@ const loadSnapshot = async ({ spinner.stop(); } }; + +const loadSnapshotAndAssertExist = async ({ + canisterId, + segment +}: { + canisterId: Principal; + segment: AssetKey; +}): Promise<{result: 'ok'; snapshotId: snapshot_id} | {result: 'not_found'}> => { + const existingSnapshotId = await loadSnapshot({canisterId}); + + if (isNullish(existingSnapshotId)) { + console.log(red(`No snapshot found for your ${displaySegment(segment)}.`)); + return {result: 'not_found'}; + } + + return {result: 'ok', snapshotId: existingSnapshotId}; +}; + +const loadSnapshotAndAssertOverwrite = async ({ + canisterId, + segment +}: { + canisterId: Principal; + segment: AssetKey; +}): Promise<{snapshotId: snapshot_id | undefined}> => { + const existingSnapshotId = await loadSnapshot({canisterId}); + + if (nonNullish(existingSnapshotId)) { + await confirmAndExit( + `A snapshot for your ${displaySegment(segment)} already exists with ID 0x${encodeSnapshotId(existingSnapshotId)}. Do you want to overwrite it?` + ); + } + + return {snapshotId: existingSnapshotId}; +}; diff --git a/src/services/modules/snapshot/snapshot.upload.services.ts b/src/services/modules/snapshot/snapshot.upload.services.ts new file mode 100644 index 00000000..a5c36af7 --- /dev/null +++ b/src/services/modules/snapshot/snapshot.upload.services.ts @@ -0,0 +1,299 @@ +import {encodeSnapshotId, type snapshot_id} from '@dfinity/ic-management'; +import {type UploadCanisterSnapshotDataKind} from '@dfinity/ic-management/dist/types/types/snapshot.params'; +import {arrayBufferToUint8Array, isNullish, jsonReviver} from '@dfinity/utils'; +import {red} from 'kleur'; +import {lstatSync} from 'node:fs'; +import {type FileHandle, open as openFile, readFile} from 'node:fs/promises'; +import {join, relative} from 'node:path'; +import ora from 'ora'; +import {uploadCanisterSnapshotData, uploadCanisterSnapshotMetadata} from '../../../api/ic.api'; +import {SnapshotMetadataSchema} from '../../../schema/snapshot.schema'; +import type {AssetKey} from '../../../types/asset-key'; +import { + type ReadCanisterSnapshotMetadataResponse, + type SnapshotBatchResult, + type SnapshotFile, + type SnapshotLog, + type SnapshotMetadata, + type UploadSnapshotParams +} from '../../../types/snapshot'; +import {displaySegment} from '../../../utils/display.utils'; +import {computeLargeFileHash} from '../../../utils/hash.utils'; +import {type BuildChunkFn, prepareDataChunks} from '../../../utils/snapshot.utils'; + +interface DataChunk { + kind: UploadCanisterSnapshotDataKind; + offset: number; + size: number; +} + +class SnapshotAssertError extends Error {} +class SnapshotFsReadError extends Error {} + +export const uploadExistingSnapshot = async ({ + segment, + ...params +}: UploadSnapshotParams & { + segment: AssetKey; + folder: string; +}): Promise => { + const spinner = ora('Uploading the snapshot...').start(); + + try { + const result = await uploadSnapshotMetadataAndMemory({ + ...params, + log: (text) => (spinner.text = text) + }); + + spinner.stop(); + + const {snapshotIdText} = result; + + console.log( + `✅ The snapshot ${snapshotIdText} for your ${displaySegment(segment)} has been uploaded.` + ); + } catch (error: unknown) { + spinner.stop(); + + if (error instanceof SnapshotFsReadError || error instanceof SnapshotAssertError) { + console.log(red(error.message)); + return; + } + + throw error; + } +}; + +const uploadSnapshotMetadataAndMemory = async ({ + log, + folder, + snapshotId: replaceSnapshotId, + ...rest +}: UploadSnapshotParams & {folder: string} & SnapshotLog): Promise<{snapshotIdText: string}> => { + // 1. Read the snapshot metadata + const { + metadata: { + metadata, + data: {wasmModule, wasmMemory, stableMemory, wasmChunkStore} + } + } = await readMetadata({folder, log}); + + // 2. Upload the snapshot metadata - i.e. we need to upload first the metadata + // because an existing snapshot is required to upload the chunk. Technically, + // we can probably assert here, in case of replacement, if the existing snapshot + // is equals to the one we upload but for simplicity reasons we upload every time. + const {snapshotId} = await uploadMetadata({ + ...rest, + snapshotId: replaceSnapshotId, + metadata, + log + }); + + const snapshotIdText = `0x${encodeSnapshotId(snapshotId)}`; + + // 3. Upload chunks + await assertAndUploadChunks({ + ...rest, + snapshotId, + log, + folder, + file: wasmModule, + build: ({offset, size}) => ({ + kind: {wasmModule: {offset}}, + offset: Number(offset), + size: Number(size) + }) + }); + + await assertAndUploadChunks({ + ...rest, + snapshotId, + log, + folder, + file: wasmMemory, + build: ({offset, size}) => ({ + kind: {wasmMemory: {offset}}, + offset: Number(offset), + size: Number(size) + }) + }); + + await assertAndUploadChunks({ + ...rest, + snapshotId, + log, + folder, + file: stableMemory, + build: ({offset, size}) => ({ + kind: {stableMemory: {offset}}, + offset: Number(offset), + size: Number(size) + }) + }); + + await assertAndUploadChunks({ + ...rest, + snapshotId, + log, + folder, + file: wasmChunkStore, + build: ({offset, size}) => ({ + kind: {wasmChunk: null}, + offset: Number(offset), + size: Number(size) + }) + }); + + return {snapshotIdText}; +}; + +const readMetadata = async ({ + folder, + log +}: {folder: string} & SnapshotLog): Promise<{metadata: SnapshotMetadata}> => { + log('Loading metadata...'); + + const source = join(folder, 'metadata.json'); + + const data = await readFile(source, 'utf-8'); + const metadata = JSON.parse(data, jsonReviver); + + return {metadata: SnapshotMetadataSchema.parse(metadata)}; +}; + +const uploadMetadata = async ({ + log, + metadata: { + globals, + certifiedData, + globalTimer, + onLowWasmMemoryHookStatus, + wasmModuleSize, + stableMemorySize, + wasmMemorySize + }, + ...rest +}: UploadSnapshotParams & {metadata: ReadCanisterSnapshotMetadataResponse} & SnapshotLog): Promise<{ + snapshotId: snapshot_id; +}> => { + log('Uploading snapshot metadata...'); + + const {snapshot_id: snapshotId} = await uploadCanisterSnapshotMetadata({ + // Handpicked data to avoid sending unexpected data over the wire + metadata: { + globals, + certifiedData, + globalTimer, + onLowWasmMemoryHookStatus, + wasmModuleSize, + stableMemorySize, + wasmMemorySize + }, + ...rest + }); + + return {snapshotId}; +}; + +const assertAndUploadChunks = async ({ + folder, + file, + log, + build, + ...rest +}: Required & { + folder: string; + file: SnapshotFile | null; + build: BuildChunkFn; +} & SnapshotLog): Promise<{status: 'success' | 'skip'}> => { + if (isNullish(file)) { + // We do not log a message to not make the developer think something is missing. + return {status: 'skip'}; + } + + const {filename, size, hash} = file; + + const source = join(folder, filename); + + const actualSize = BigInt(lstatSync(source).size); + if (size !== actualSize) { + throw new SnapshotAssertError( + `Size mismatch for ${filename}: expected ${size} bytes, got ${actualSize} bytes` + ); + } + + const actualHash = await computeLargeFileHash(source); + if (hash !== actualHash) { + throw new SnapshotAssertError( + `Hash mismatch for ${filename}: expected ${hash}, got ${actualHash}` + ); + } + + const {chunks} = prepareDataChunks({size, build}); + + log(`Uploading chunks from ${relative(process.cwd(), source)}...`); + + const sourceHandler = await openFile(source); + + try { + for await (const {progress} of batchUploadChunks({ + sourceHandler, + chunks, + ...rest + })) { + log(`Chunks ${progress.done}/${progress.total} uploaded. Continuing...`); + } + } finally { + await sourceHandler.close(); + } + + return {status: 'success'}; +}; + +async function* batchUploadChunks({ + sourceHandler, + chunks, + limit = 20, + ...params +}: Required & { + sourceHandler: FileHandle; + chunks: DataChunk[]; + limit?: number; +}): AsyncGenerator { + const total = chunks.length; + + for (let i = 0; i < total; i = i + limit) { + const batch = chunks.slice(i, i + limit); + await Promise.all( + batch.map(async (requestChunk) => { + await uploadChunk({ + ...params, + sourceHandler, + chunk: requestChunk + }); + }) + ); + yield {progress: {index: i, done: Math.min(i + limit, total), total}}; + } +} + +const uploadChunk = async ({ + chunk: {kind, size, offset}, + sourceHandler, + ...rest +}: Required & { + sourceHandler: FileHandle; + chunk: DataChunk; +}): Promise => { + const {buffer, bytesRead} = await sourceHandler.read(Buffer.alloc(size), 0, size, offset); + + if (bytesRead !== size) { + throw new SnapshotFsReadError(`Unexpected bytes read: expected ${size} but got ${bytesRead}`); + } + + await uploadCanisterSnapshotData({ + ...rest, + kind, + chunk: arrayBufferToUint8Array(buffer.buffer) + }); +}; diff --git a/src/types/snapshot.ts b/src/types/snapshot.ts new file mode 100644 index 00000000..ecf8cd3b --- /dev/null +++ b/src/types/snapshot.ts @@ -0,0 +1,34 @@ +import type {snapshot_id} from '@dfinity/ic-management'; +import type {Principal} from '@dfinity/principal'; +import type * as z from 'zod'; +import { + type ReadCanisterSnapshotMetadataResponseSchema, + type SnapshotFilenameSchema, + type SnapshotFileSchema, + type SnapshotMetadataSchema +} from '../schema/snapshot.schema'; + +export type SnapshotFilename = z.infer; +export type SnapshotFile = z.infer; +export type SnapshotMetadata = z.infer; +export type ReadCanisterSnapshotMetadataResponse = z.infer< + typeof ReadCanisterSnapshotMetadataResponseSchema +>; + +export interface DownloadSnapshotParams { + canisterId: Principal; + snapshotId: snapshot_id; +} + +export type UploadSnapshotParams = Omit & + Partial>; + +// A handy wrapper to pass down a function that updates +// the spinner log. +export interface SnapshotLog { + log: (text: string) => void; +} + +export interface SnapshotBatchResult { + progress: {index: number; done: number; total: number}; +} diff --git a/src/utils/hash.utils.ts b/src/utils/hash.utils.ts new file mode 100644 index 00000000..908fc5cb --- /dev/null +++ b/src/utils/hash.utils.ts @@ -0,0 +1,20 @@ +import {createHash} from 'node:crypto'; +import {createReadStream} from 'node:fs'; +import {Writable} from 'node:stream'; +import {pipeline} from 'node:stream/promises'; + +export const computeLargeFileHash = async (filepath: string): Promise => { + const hash = createHash('sha256'); + + await pipeline( + createReadStream(filepath), + new Writable({ + write(chunk: Uint8Array, _enc, cb) { + hash.update(chunk); + cb(); + } + }) + ); + + return hash.digest('hex'); +}; diff --git a/src/utils/snapshot.utils.ts b/src/utils/snapshot.utils.ts new file mode 100644 index 00000000..9ea500c3 --- /dev/null +++ b/src/utils/snapshot.utils.ts @@ -0,0 +1,27 @@ +import {SNAPSHOT_MAX_CHUNK_SIZE} from '../constants/snapshot.constants'; + +export type BuildChunkFn = (params: {offset: bigint; size: bigint}) => T; + +export const prepareDataChunks = ({ + size: totalSize, + build +}: { + size: bigint; + build: BuildChunkFn; +}): {chunks: T[]} => { + const chunks: T[] = []; + + for (let offset = 0n; offset < totalSize; offset += SNAPSHOT_MAX_CHUNK_SIZE) { + const size = + offset + SNAPSHOT_MAX_CHUNK_SIZE <= totalSize ? SNAPSHOT_MAX_CHUNK_SIZE : totalSize - offset; + + chunks.push({ + ...build({ + offset, + size + }) + }); + } + + return {chunks}; +};