feat: download and upload snapshot offline (#382)

* feat: download snapshot offline

* feat: download chunks

* feat: download all memories

* feat: print

* feat: logs

* feat: skip size zero

* feat: use heap and stable

* feat: download chunk store

* feat: log

* feat: more chunks

* feat: rename and integrate for all modules

* feat: smaller limit for batching

* feat: add command to help

* feat: init upload

* feat: write stream to a single file

* feat: assert size and write metadata once the process is complete

* feat: read metadata

* feat: upload snapshot

* feat: check dir before loading snapshot and fix passing down args

* fix: always a new snapshot id

* feat: cleanup

* chore: rename constant

* refactor: extract and move types and schemas

* refactor: move hash utils

* chore: fmt and lint

* docs: incorrect referenced lib

Author: peterpeterparker

package-lock.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",

package.json

@@ -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<ReadCanisterSnapshotMetadataResponse> => {
+  const agent = await initAgent();
+
+  const {readCanisterSnapshotMetadata} = ICManagementCanister.create({
+    agent
+  });
+
+  return await readCanisterSnapshotMetadata(params);
+};
+
+export const readCanisterSnapshotData = async (
+  params: ReadCanisterSnapshotMetadataParams
+): Promise<read_canister_snapshot_data_response> => {
+  const agent = await initAgent();
+
+  const {readCanisterSnapshotData} = ICManagementCanister.create({
+    agent
+  });
+
+  return await readCanisterSnapshotData(params);
+};
+
+export const uploadCanisterSnapshotMetadata = async (
+  params: UploadCanisterSnapshotMetadataParams
+): Promise<upload_canister_snapshot_metadata_response> => {
+  const agent = await initAgent();
+
+  const {uploadCanisterSnapshotMetadata} = ICManagementCanister.create({
+    agent
+  });
+
+  return await uploadCanisterSnapshotMetadata(params);
+};
+
+export const uploadCanisterSnapshotData = async (
+  params: UploadCanisterSnapshotDataParams
+): Promise<void> => {
+  const agent = await initAgent();
+
+  const {uploadCanisterSnapshotData} = ICManagementCanister.create({
+    agent
+  });
+
+  await uploadCanisterSnapshotData(params);
+};

src/api/ic.api.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<void>;
-  missionControlFn: () => Promise<void>;
-  orbiterFn: () => Promise<void>;
+  satelliteFn: (args?: string[]) => Promise<void>;
+  missionControlFn: (args?: string[]) => Promise<void>;
+  orbiterFn: (args?: string[]) => Promise<void>;
 }) => {
   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);
+  }
+};

src/commands/snapshot.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.`;

src/constants/help.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;

src/constants/snapshot.constants.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('<subcomman
 
 Subcommands:
   ${magenta('create')}               Create a snapshot of your current state.
-  ${magenta('restore')}              Restore a previously created snapshot.
   ${magenta('delete')}               Delete an existing snapshot.
+  ${magenta('download')}             Download a snapshot to offline files.
+  ${magenta('upload')}               ${SNAPSHOT_UPLOAD_DESCRIPTION}
+  ${magenta('restore')}              Restore a previously created snapshot.
   
 Options:
   ${targetOption('snapshot')}

src/help/snapshot.help.ts

@@ -0,0 +1,29 @@
+import {cyan, green, magenta, yellow} from 'kleur';
+import {OPTION_HELP, OPTIONS_ENV, SNAPSHOT_UPLOAD_DESCRIPTION} from '../constants/help.constants';
+import {helpOutput} from './common.help';
+import {TITLE} from './help';
+
+const usage = `Usage: ${green('juno')} ${cyan('snapshot')} ${magenta('upload')} ${yellow('[options]')}
+
+Options:
+  ${yellow('--dir')}                 Path to the snapshot directory that contains the metadata.json and chunks.
+  ${OPTIONS_ENV}
+  ${OPTION_HELP}`;
+
+const doc = `${SNAPSHOT_UPLOAD_DESCRIPTION}
+
+\`\`\`
+${usage}
+\`\`\`
+`;
+
+const help = `${TITLE}
+
+${SNAPSHOT_UPLOAD_DESCRIPTION}
+
+${usage}
+`;
+
+export const logHelpSnapshotUpload = (args?: string[]) => {
+  console.log(helpOutput(args) === 'doc' ? doc : help);
+};

src/help/snapshot.upload.help.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);

src/index.ts

@@ -0,0 +1,65 @@
+import * as z from 'zod';
+
+const Uint8ArrayLike = z.instanceof(Uint8Array) as z.ZodType<Uint8Array>;
+
+// 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
+});