From 73015da5a31f06594b8f53c4d46260132ecd1d2c Mon Sep 17 00:00:00 2001
From: "Timothy J. Aveni" <tja@berkeley.edu>
Date: Sat, 20 Jun 2026 12:54:05 -0700
Subject: [PATCH] Apply NixOS flakes from EC2 user data

---
 README.md          |  5 ++++-
 user-data.sh.tftpl | 41 +++++++++++++++++++++++++++++++++++++++--
 2 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index de47dc1..cac5913 100644
--- a/README.md
+++ b/README.md
@@ -10,7 +10,10 @@ device mapping. Runtime deployments should consume the AMI ID directly and
 should not independently choose a `/nix` snapshot.
 
 The image is intentionally project-agnostic. It provides NixOS, ZFS, SSM,
-flakes, and systemd ordering so the Nix daemon waits for `/nix`.
+flakes, and systemd ordering so the Nix daemon waits for `/nix`. The final AMI
+also disables the default shell-oriented `amazon-init` service and enables a
+small generic service that treats EC2 user-data as `/etc/nixos/flake.nix`, then
+runs `nixos-rebuild switch --flake /etc/nixos`.
 
 ## Publish
 
diff --git a/user-data.sh.tftpl b/user-data.sh.tftpl
index 44dce59..fba06e7 100644
--- a/user-data.sh.tftpl
+++ b/user-data.sh.tftpl
@@ -21,7 +21,7 @@ write_bootstrap_config() {
     nixosConfigurations.ami-builder-bootstrap = nixpkgs.lib.nixosSystem {
       system = "${system}";
       modules = [
-        ({ modulesPath, pkgs, ... }: {
+        ({ config, lib, modulesPath, pkgs, ... }: {
           imports = [ "$${modulesPath}/virtualisation/amazon-image.nix" ];
 
           networking.hostName = "nixos-zfs-ec2-builder";
@@ -67,7 +67,7 @@ write_final_config() {
     nixosConfigurations.nixos-zfs-ec2 = nixpkgs.lib.nixosSystem {
       system = "${system}";
       modules = [
-        ({ modulesPath, pkgs, ... }: {
+        ({ config, lib, modulesPath, pkgs, ... }: {
           imports = [ "$${modulesPath}/virtualisation/amazon-image.nix" ];
 
           networking.hostName = "nixos-zfs-ec2";
@@ -101,6 +101,43 @@ write_final_config() {
           ];
 
           services.amazon-ssm-agent.enable = true;
+          systemd.services.amazon-init.enable = lib.mkForce false;
+
+          systemd.services.nixos-user-data-flake = {
+            description = "Apply NixOS flake from EC2 user data";
+            wantedBy = [ "multi-user.target" ];
+            after = [ "fetch-ec2-metadata.service" "network-online.target" "nix-daemon.socket" ];
+            wants = [ "fetch-ec2-metadata.service" "network-online.target" "nix-daemon.socket" ];
+            path = [ pkgs.coreutils pkgs.git ];
+            serviceConfig = {
+              Type = "oneshot";
+              RemainAfterExit = true;
+            };
+            script = ''
+              set -eu
+
+              user_data=/etc/ec2-metadata/user-data
+              state_dir=/var/lib/nixos-user-data-flake
+              applied_hash_file="$state_dir/applied-sha256"
+
+              if [ ! -s "$user_data" ]; then
+                echo "No EC2 user-data flake found at $user_data" >&2
+                exit 1
+              fi
+
+              mkdir -p "$state_dir" /etc/nixos
+              new_hash="$(sha256sum "$user_data" | cut -d ' ' -f1)"
+              old_hash="$(cat "$applied_hash_file" 2>/dev/null || true)"
+
+              if [ "$new_hash" = "$old_hash" ]; then
+                exit 0
+              fi
+
+              install -m 0644 "$user_data" /etc/nixos/flake.nix
+              ${config.system.build.nixos-rebuild}/bin/nixos-rebuild switch --flake /etc/nixos
+              printf '%s\n' "$new_hash" > "$applied_hash_file"
+            '';
+          };
 
           users.users.builder-admin = {
             isNormalUser = true;