hacking on rollback

Cargo.toml

@@ -1,5 +1,5 @@
 [workspace]
-members = ["wire/key_agent", "wire/lib", "wire/cli"]
+members = ["wire/agent", "wire/lib", "wire/cli"]
 resolver = "2"
 package.edition = "2024"
 package.version = "1.0.0-beta.0"

README.md

@@ -15,7 +15,7 @@ wire
 │  │  └── Rust library containing business logic, consumed by `wire`
 │  ├── cli
 │  │  └── Rust binary, using `lib`
-│  └── key_agent
+│  └── agent
 │     └── Rust binary ran on a target node. receives key file bytes and metadata w/ protobuf over SSH stdin
 ├── doc
 │  └── a [vitepress](https://vitepress.dev/) site

doc/guides/keys.md

@@ -38,7 +38,7 @@ user must be trusted or you must add garnix as a trusted public key:
 Otherwise, you may see errors such as:
 
 ```
-error: cannot add path '/nix/store/...-wire-tool-key_agent-x86_64-linux-...' because it lacks a signature by a trusted key
+error: cannot add path '/nix/store/...-wire-tool-agent-x86_64-linux-...' because it lacks a signature by a trusted key
 ```
 
 This is a requirement because `nix copy` is used to copy the binary.

doc/guides/non-root-user.md

@@ -56,7 +56,7 @@ $ wire apply keys --on media
  INFO eval_hive: evaluating hive Flake("/path/to/hive")
 ...
  INFO media | step="Upload key @ NoFilter" progress="3/4"
-deploy-user@node:22 | Authenticate for "sudo /nix/store/.../bin/key_agent":
+deploy-user@node:22 | Authenticate for "sudo /nix/store/.../bin/agent":
 [sudo] password for deploy-user:
 ```
 

flake.nix

@@ -38,7 +38,7 @@
         ./nix/shells.nix
         ./nix/tests.nix
         ./wire/cli
-        ./wire/key_agent
+        ./wire/agent
         ./doc
         ./tests/nix
         ./runtime

justfile

@@ -4,4 +4,4 @@ build-dhat:
     cargo build --profile profiling --features dhat-heap
     @echo 'dhat binaries in target/profiling'
     @echo 'Example:'
-    @echo 'WIRE_KEY_AGENT=/nix/store/...-key_agent-0.1.0 PROJECT/target/profiling/wire apply ...'
+    @echo 'WIRE_AGENT=/nix/store/...-agent-0.1.0 PROJECT/target/profiling/wire apply ...'

nix/tests.nix

@@ -20,7 +20,7 @@
 
           installPhaseCommand = ''
             mkdir -p $out
-            cp $(ls target/debug/deps/{wire,lib,key_agent}-* | grep -v "\.d") $out
+            cp $(ls target/debug/deps/{wire,lib,agent}-* | grep -v "\.d") $out
           '';
         }
         // commonArgs

runtime/module/config.nix

@@ -22,7 +22,74 @@
         }
       ) config.deployment.keys;
 
-      services = lib.mapAttrs' (
+      system.activationScripts.setup-wire-rollback.text = ''
+        mkdir -p /var/lib/wire-rollback
+        chmod 700 /var/lib/wire-rollback
+      '';
+
+      services = {
+        wire-rollback = {
+          enable = config.deployment.rollback;
+          description = "Rolls back the NixOS profile if `/var/lib/wire-rollback/heartbeat` is not created in 30
+          seconds after this service starts.";
+          documentation = [
+            "https://wire.althaea.zone/guides/rollback"
+          ];
+          path = [
+            pkgs.coreutils
+          ];
+          wantedBy = [ "multi-user.target" ];
+          script = ''
+            set -euo pipefail
+
+            goal=$(<"/var/lib/wire-rollback/goal")
+
+            case $goal in
+                "check" | "switch" | "boot" | "test" | "dry-activate")
+                    echo "<5>using goal $goal"
+                    ;;
+                *)
+                    echo "<3>'$goal' is not a valid goal."
+                    exit 1
+                    ;;
+            esac
+
+            sleep 30
+
+            if [ -f "/var/lib/wire-rollback/heartbeat" ]; then
+                exit 0
+            fi
+
+            echo "<1>/var/lib/wire-rollback/heartbeat does not exist, rolling back system"
+
+            # set current system
+            nix-env --rollback --profile /nix/var/nix/profiles/system
+            # get the path to the system we are now rolling back to
+            system=$(readlink -f /nix/var/nix/profiles/system)
+
+            echo "<5>rolling back to $system"
+
+            # switch to the system using goal
+            "$system/bin/switch-to-configuration $goal"
+          '';
+          unitConfig = {
+            ConditionPathExists = [
+              "/var/lib/wire-rollback/goal"
+              "!/var/lib/wire-rollback/heartbeat"
+            ];
+          };
+          serviceConfig = {
+            Type = "oneshot";
+            Restart = "no";
+            StateDirectory = "wire-rollback";
+            NotifyAccess = "all";
+            RemainAfterExit = "yes";
+
+            ExecStopPost = "${pkgs.coreutils}/bin/rm -f /var/lib/wire-rollback/goal";
+          };
+        };
+      }
+      // (lib.mapAttrs' (
         _name: value:
         lib.nameValuePair "${value.name}-key" {
           description = "Service that requires ${value.path}";
@@ -55,7 +122,7 @@
             RemainAfterExit = "yes";
           };
         }
-      ) config.deployment.keys;
+      ) config.deployment.keys);
     };
 
     deployment = {

runtime/module/options.nix

@@ -50,6 +50,12 @@ in
       default = { };
     };
 
+    rollback = lib.mkOption {
+      type = types.bool;
+      default = true;
+      description = "Attempt to rollback this node if it cannot be contacted after activation.";
+    };
+
     buildOnTarget = lib.mkOption {
       type = types.bool;
       default = false;

tests/nix/default.nix

@@ -30,6 +30,7 @@ in
     ./suite/test_local_deploy
     ./suite/test_keys
     ./suite/test_stdin
+    ./suite/test_rollback
   ];
   options.wire.testing = mkOption {
     type = attrsOf (

tests/nix/suite/test_rollback/default.nix

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: AGPL-3.0-or-later
+# Copyright 2024-2025 wire Contributors
+
+{
+  wire.testing.test_remote_deploy = {
+    nodes.deployer = {
+      _wire.deployer = true;
+    };
+    nodes.receiver = {
+      _wire.receiver = true;
+    };
+    testScript = ''
+      with subtest("Deploy broken config"):
+        deployer.fail(f"wire apply --on receiver --no-progress --path {TEST_DIR}/hive.nix --no-keys -vvv >&2")
+
+      with subtest("Configuration must revert"):
+        receiver.wait_for_unit("sshd.service")
+    '';
+  };
+}

tests/nix/suite/test_rollback/hive.nix

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: AGPL-3.0-or-later
+# Copyright 2024-2025 wire Contributors
+
+let
+  inherit (import ../utils.nix { testName = "test_rollback-@IDENT@"; }) makeHive mkHiveNode;
+in
+makeHive {
+  meta.nixpkgs = import <nixpkgs> { localSystem = "x86_64-linux"; };
+
+  receiver = mkHiveNode { hostname = "receiver"; } {
+    environment.etc."identity".text = "first";
+
+    services.openssh.enable = false;
+  };
+}

wire/key_agent/Cargo.toml → wire/agent/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "key_agent"
+name = "agent"
 edition.workspace = true
 version.workspace = true
 

wire/key_agent/default.nix → wire/agent/default.nix

@@ -8,9 +8,9 @@
     {
       packages = {
         agent = buildRustProgram {
-          name = "key_agent";
-          pname = "wire-tool-key_agent-${system}";
-          cargoExtraArgs = "-p key_agent";
+          name = "agent";
+          pname = "wire-tool-agent-${system}";
+          cargoExtraArgs = "-p agent";
         };
       };
     };

wire/key_agent/src/keys.proto → wire/agent/src/keys.proto

@@ -3,7 +3,7 @@
 
 syntax = "proto3";
 
-package key_agent.keys;
+package agent.keys;
 
 message KeySpec {
   string destination = 1;

wire/key_agent/src/lib.rs → wire/agent/src/lib.rs

@@ -2,5 +2,5 @@
 // Copyright 2024-2025 wire Contributors
 
 pub mod keys {
-    include!(concat!(env!("OUT_DIR"), "/key_agent.keys.rs"));
+    include!(concat!(env!("OUT_DIR"), "/agent.keys.rs"));
 }

wire/key_agent/src/main.rs → wire/agent/src/main.rs

@@ -2,10 +2,10 @@
 // Copyright 2024-2025 wire Contributors
 
 #![deny(clippy::pedantic)]
+use agent::keys::KeySpec;
 use base64::Engine;
 use base64::prelude::BASE64_STANDARD;
 use futures_util::stream::StreamExt;
-use key_agent::keys::KeySpec;
 use nix::unistd::{Group, User};
 use prost::Message;
 use prost::bytes::Bytes;

wire/cli/default.nix

@@ -12,7 +12,7 @@
     let
       cleanSystem = system: lib.replaceStrings [ "-" ] [ "_" ] system;
       agents = lib.strings.concatMapStrings (
-        system: "--set WIRE_KEY_AGENT_${cleanSystem system} ${(getSystem system).packages.agent} "
+        system: "--set WIRE_AGENT_${cleanSystem system} ${(getSystem system).packages.agent} "
       ) (import inputs.linux-systems);
     in
     {
@@ -70,7 +70,7 @@
             pkgs.makeWrapper
           ];
           postBuild = ''
-            wrapProgram $out/bin/wire --set WIRE_KEY_AGENT_${cleanSystem system} ${self'.packages.agent}
+            wrapProgram $out/bin/wire --set WIRE_AGENT_${cleanSystem system} ${self'.packages.agent}
           '';
           meta.mainProgram = "wire";
         };

wire/lib/Cargo.toml

@@ -17,7 +17,7 @@ tracing = { workspace = true }
 im = { workspace = true }
 thiserror = "2.0.17"
 derive_more = { version = "2.0.1", features = ["display"] }
-key_agent = { path = "../key_agent" }
+agent = { path = "../agent" }
 futures = "0.3.31"
 prost = { workspace = true }
 gethostname = "1.1.0"

wire/lib/src/errors.rs

@@ -87,6 +87,20 @@ pub enum ActivationError {
     )]
     #[error("failed to run switch-to-configuration {0} on node {1}")]
     SwitchToConfigurationError(SwitchToConfigurationGoal, Name, #[source] CommandError),
+
+    #[diagnostic(
+        code(wire::activation::Heartbeat),
+        url("{DOCS_URL}#{}", self.code().unwrap())
+    )]
+    #[error("failed to touch /var/lib/wire-rollback/heartbeat on node {name}")]
+    FailedHeartbeatError {
+        name: Name,
+        #[source]
+        activation_failure: CommandError,
+
+        #[related]
+        related_errors: Vec<CommandError>,
+    },
 }
 
 #[derive(Debug, Diagnostic, Error)]

wire/lib/src/hive/node.rs

@@ -154,11 +154,19 @@ impl Display for Target {
     }
 }
 
+const fn rollback_default() -> bool {
+    true
+}
+
 #[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq, Hash)]
 pub struct Node {
     #[serde(rename = "target")]
     pub target: Target,
 
+    /// default value as this is a new attribute
+    #[serde(rename = "rollback", default = "rollback_default")]
+    pub rollback: bool,
+
     #[serde(rename = "buildOnTarget")]
     pub build_remotely: bool,
 
@@ -192,6 +200,7 @@ impl Default for Node {
             allow_local_deployment: true,
             build_remotely: false,
             host_platform: "x86_64-linux".into(),
+            rollback: rollback_default(),
         }
     }
 }
@@ -291,9 +300,10 @@ pub struct StepState {
     pub evaluation: Option<Derivation>,
     pub evaluation_rx: Option<oneshot::Receiver<Result<Derivation, HiveLibError>>>,
     pub build: Option<String>,
-    pub key_agent_directory: Option<String>,
+    pub agent_directory: Option<String>,
 }
 
+#[allow(clippy::struct_excessive_bools)]
 pub struct Context<'a> {
     pub name: &'a Name,
     pub node: &'a mut Node,