feat: post-install actions

CHANGELOG.adoc

@@ -5,6 +5,10 @@
 
 == DFX
 
+=== feat: Post-installation tasks
+
+You can now add your own custom post-installation/post-deployment tasks to any canister type. The new `+post-install+` key for canister objects in `+dfx.json+` can be a command or list of commands, similar to the `+build+` key of `+custom+` canisters, and receives all the same environment variables. For example, to replicate the upload task performed with `+assets+` canisters, you might set `+"post-install": "icx-asset sync $CANISTER_ID dist"+`.
+
 === feat: assets are no longer copied from source directories before being uploaded to asset canister
 
 Assets are now uploaded directly from their source directories, rather than first being copied

e2e/assets/post_install/dfx.json

@@ -0,0 +1,24 @@
+{
+  "version": 1,
+  "canisters": {
+    "postinstall": {
+      "main": "main.mo",
+      "post_install": "echo hello-file"
+    },
+    "postinstall_script": {
+      "main": "main.mo",
+      "post_install": "postinstall.sh",
+      "dependencies": ["postinstall"]
+    }
+  },
+  "defaults": {
+    "build": {
+      "output": "canisters/"
+    }
+  },
+  "networks": {
+    "local": {
+      "bind": "127.0.0.1:8000"
+    }
+  }
+}

e2e/assets/post_install/main.mo

@@ -0,0 +1,3 @@
+actor {
+
+}

e2e/assets/post_install/patch.bash

@@ -0,0 +1,2 @@
+# Do nothing
+

e2e/assets/post_install/postinstall.sh

@@ -0,0 +1,2 @@
+#!/usr/bin/env bash
+echo hello-script

e2e/tests-dfx/install.bash

@@ -81,3 +81,54 @@ teardown() {
     dfx canister create --all
     assert_command_fail dfx canister install --all --wasm "${archive:?}/wallet/0.10.0/wallet.wasm"
 }
+
+@test "install runs post-install tasks" {
+    install_asset post_install
+    dfx_start
+
+    assert_command dfx canister create --all
+    assert_command dfx build
+
+    assert_command dfx canister install postinstall
+    assert_match 'hello-file'
+
+    assert_command dfx canister install postinstall_script
+    assert_match 'hello-script'
+
+    echo 'return 1' >> postinstall.sh
+    assert_command_fail dfx canister install postinstall_script --mode upgrade
+    assert_match 'hello-script'
+}
+
+@test "post-install tasks receive environment variables" {
+    install_asset post_install
+    dfx_start
+    echo "echo hello \$CANISTER_ID" >> postinstall.sh
+
+    assert_command dfx canister create --all
+    assert_command dfx build
+    id=$(dfx canister id postinstall_script)
+
+    assert_command dfx canister install --all
+    assert_match "hello $id"
+    assert_command dfx canister install postinstall_script --mode upgrade
+    assert_match "hello $id"
+
+    assert_command dfx deploy
+    assert_match "hello $id"
+    assert_command dfx deploy postinstall_script
+    assert_match "hello $id"
+}
+
+@test "post-install tasks discover dependencies" {
+    install_asset post_install
+    dfx_start
+    echo "echo hello \$CANISTER_ID_postinstall" >> postinstall.sh
+
+    assert_command dfx canister create --all
+    assert_command dfx build
+    id=$(dfx canister id postinstall)
+
+    assert_command dfx canister install postinstall_script
+    assert_match "hello $id"
+}

src/dfx/src/commands/canister/install.rs

@@ -127,6 +127,7 @@ pub async fn exec(
                 call_sender,
                 installed_module_hash,
                 opts.upgrade_unchanged,
+                None,
             )
             .await
         }
@@ -167,6 +168,7 @@ pub async fn exec(
                     call_sender,
                     installed_module_hash,
                     opts.upgrade_unchanged,
+                    None,
                 )
                 .await?;
             }

src/dfx/src/config/dfinity.rs

@@ -1,5 +1,6 @@
 #![allow(dead_code)]
 use crate::lib::error::{BuildError, DfxError, DfxResult};
+use crate::util::SerdeVec;
 use crate::{error_invalid_argument, error_invalid_config, error_invalid_data};
 
 use anyhow::{anyhow, Context};
@@ -71,6 +72,9 @@ pub struct ConfigCanistersCanister {
     #[serde(default)]
     pub remote: Option<ConfigCanistersCanisterRemote>,
 
+    #[serde(default)]
+    pub post_install: SerdeVec<String>,
+
     #[serde(flatten)]
     pub extras: BTreeMap<String, Value>,
 }

src/dfx/src/lib/builders/mod.rs

@@ -250,7 +250,7 @@ fn ensure_trailing_newline(s: String) -> String {
 
 type Env<'a> = (Cow<'static, str>, Cow<'a, OsStr>);
 
-fn environment_variables<'a>(
+pub fn environment_variables<'a>(
     info: &CanisterInfo,
     network_name: &'a str,
     pool: &'a CanisterPool,

src/dfx/src/lib/canister_info.rs

@@ -51,6 +51,8 @@ pub struct CanisterInfo {
     packtool: Option<String>,
     args: Option<String>,
 
+    post_install: Vec<String>,
+
     extras: BTreeMap<String, serde_json::Value>,
 }
 
@@ -113,6 +115,8 @@ impl CanisterInfo {
             .cloned()
             .unwrap_or_else(|| "motoko".to_owned());
 
+        let post_install = canister_config.post_install.clone().into_vec();
+
         let canister_info = CanisterInfo {
             name: name.to_string(),
             canister_type,
@@ -127,6 +131,8 @@ impl CanisterInfo {
             packtool: build_defaults.get_packtool(),
             args: build_defaults.get_args(),
             extras,
+
+            post_install,
         };
 
         let canister_args: Option<String> = canister_info.get_extra_optional("args")?;
@@ -233,6 +239,10 @@ impl CanisterInfo {
         &self.packtool
     }
 
+    pub fn get_post_install(&self) -> &[String] {
+        &self.post_install
+    }
+
     pub fn get_args(&self) -> &Option<String> {
         &self.args
     }

src/dfx/src/lib/operations/canister/deploy_canisters.rs

@@ -91,7 +91,7 @@ pub async fn deploy_canisters(
     )
     .await?;
 
-    build_canisters(env, &canisters_to_build, &config)?;
+    let pool = build_canisters(env, &canisters_to_build, &config)?;
 
     install_canisters(
         env,
@@ -104,6 +104,7 @@ pub async fn deploy_canisters(
         upgrade_unchanged,
         timeout,
         call_sender,
+        pool,
     )
     .await?;
 
@@ -193,12 +194,17 @@ async fn register_canisters(
 }
 
 #[context("Failed to build call canisters.")]
-fn build_canisters(env: &dyn Environment, canister_names: &[String], config: &Config) -> DfxResult {
+fn build_canisters(
+    env: &dyn Environment,
+    canister_names: &[String],
+    config: &Config,
+) -> DfxResult<CanisterPool> {
     info!(env.get_logger(), "Building canisters...");
     let build_mode_check = false;
     let canister_pool = CanisterPool::load(env, build_mode_check, canister_names)?;
 
-    canister_pool.build_or_fail(&BuildConfig::from_config(config)?)
+    canister_pool.build_or_fail(&BuildConfig::from_config(config)?)?;
+    Ok(canister_pool)
 }
 
 #[allow(clippy::too_many_arguments)]
@@ -214,6 +220,7 @@ async fn install_canisters(
     upgrade_unchanged: bool,
     timeout: Duration,
     call_sender: &CallSender,
+    pool: CanisterPool,
 ) -> DfxResult {
     info!(env.get_logger(), "Installing canisters...");
 
@@ -266,6 +273,7 @@ async fn install_canisters(
             call_sender,
             installed_module_hash,
             upgrade_unchanged,
+            Some(&pool),
         )
         .await?;
     }

src/dfx/src/lib/operations/canister/install_canister.rs

@@ -1,11 +1,14 @@
+use crate::lib::builders::environment_variables;
 use crate::lib::canister_info::CanisterInfo;
 use crate::lib::environment::Environment;
 use crate::lib::error::DfxResult;
 use crate::lib::identity::identity_utils::CallSender;
 use crate::lib::identity::Identity;
 use crate::lib::installers::assets::post_install_store_assets;
+use crate::lib::models::canister::CanisterPool;
 use crate::lib::models::canister_id_store::CanisterIdStore;
 use crate::lib::named_canister;
+use crate::lib::network::network_descriptor::NetworkDescriptor;
 use crate::lib::waiter::waiter_with_timeout;
 use crate::util::assets::wallet_wasm;
 use crate::util::{expiry_duration, read_module_metadata};
@@ -18,10 +21,12 @@ use ic_utils::call::AsyncCall;
 use ic_utils::interfaces::management_canister::builders::{CanisterInstall, InstallMode};
 use ic_utils::interfaces::ManagementCanister;
 use ic_utils::Argument;
+use itertools::Itertools;
 use openssl::sha::Sha256;
 use slog::info;
 use std::collections::HashSet;
 use std::io::stdin;
+use std::process::{Command, Stdio};
 use std::time::Duration;
 
 #[allow(clippy::too_many_arguments)]
@@ -37,6 +42,7 @@ pub async fn install_canister(
     call_sender: &CallSender,
     installed_module_hash: Option<Vec<u8>>,
     upgrade_unchanged: bool,
+    pool: Option<&CanisterPool>,
 ) -> DfxResult {
     let log = env.get_logger();
     let network = env.get_network_descriptor();
@@ -161,6 +167,79 @@ pub async fn install_canister(
         post_install_store_assets(canister_info, agent, timeout).await?;
     }
 
+    if !canister_info.get_post_install().is_empty() {
+        run_post_install_tasks(env, canister_info, network, pool)?;
+    }
+
+    Ok(())
+}
+
+#[context("Failed to run post-install tasks")]
+fn run_post_install_tasks(
+    env: &dyn Environment,
+    canister: &CanisterInfo,
+    network: &NetworkDescriptor,
+    pool: Option<&CanisterPool>,
+) -> DfxResult {
+    let tmp;
+    let pool = match pool {
+        Some(pool) => pool,
+        None => {
+            tmp = env
+                .get_config_or_anyhow()?
+                .get_config()
+                .get_canister_names_with_dependencies(Some(canister.get_name()))
+                .and_then(|deps| CanisterPool::load(env, false, &deps))
+                .context("Error collecting canisters for post-install task")?;
+            &tmp
+        }
+    };
+    let dependencies = pool
+        .get_canister_list()
+        .iter()
+        .map(|can| can.canister_id())
+        .collect_vec();
+    for task in canister.get_post_install() {
+        run_post_install_task(canister, task, network, pool, &dependencies)?;
+    }
+    Ok(())
+}
+
+#[context("Failed to run post-install task {task}")]
+fn run_post_install_task(
+    canister: &CanisterInfo,
+    task: &str,
+    network: &NetworkDescriptor,
+    pool: &CanisterPool,
+    dependencies: &[Principal],
+) -> DfxResult {
+    let cwd = canister.get_workspace_root();
+    let words = shell_words::split(task)
+        .with_context(|| format!("Error interpreting post-install task `{task}`"))?;
+    let canonicalized = cwd
+        .join(&words[0])
+        .canonicalize()
+        .or_else(|_| which::which(&words[0]))
+        .map_err(|_| anyhow!("Cannot find command or file {}", &words[0]))?;
+    let mut command = Command::new(&canonicalized);
+    command.args(&words[1..]);
+    let vars = environment_variables(canister, &network.name, pool, dependencies);
+    for (key, val) in vars {
+        command.env(&*key, val);
+    }
+    command
+        .current_dir(cwd)
+        .stdout(Stdio::inherit())
+        .stderr(Stdio::inherit());
+    let status = command.status()?;
+    if !status.success() {
+        match status.code() {
+            Some(code) => {
+                bail!("The post-install task `{task}` failed with exit code {code}")
+            }
+            None => bail!("The post-install task `{task}` was terminated by a signal"),
+        }
+    }
     Ok(())
 }
 

src/dfx/src/util/mod.rs

@@ -8,6 +8,7 @@ use candid::{parser::value::IDLValue, IDLArgs};
 use fn_error_context::context;
 use net2::TcpListenerExt;
 use net2::{unix::UnixTcpBuilderExt, TcpBuilder};
+use serde::{Deserialize, Serialize};
 use std::net::{IpAddr, SocketAddr};
 use std::time::Duration;
 
@@ -220,3 +221,25 @@ pub fn blob_from_arguments(
         v => Err(error_unknown!("Invalid type: {}", v)),
     }
 }
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+#[serde(untagged)]
+pub enum SerdeVec<T> {
+    One(T),
+    Many(Vec<T>),
+}
+
+impl<T> SerdeVec<T> {
+    pub fn into_vec(self) -> Vec<T> {
+        match self {
+            Self::One(t) => vec![t],
+            Self::Many(ts) => ts,
+        }
+    }
+}
+
+impl<T> Default for SerdeVec<T> {
+    fn default() -> Self {
+        Self::Many(vec![])
+    }
+}