From ab223837c44846cf129a56d360983ddaa6b35790 Mon Sep 17 00:00:00 2001
From: "Timothy J. Aveni" <tja@berkeley.edu>
Date: Tue, 30 Jun 2026 19:10:56 -0700
Subject: [PATCH] Extract shared AMI NixOS modules

---
 user-data.sh.tftpl | 96 ++++++++++++++++------------------------------
 1 file changed, 34 insertions(+), 62 deletions(-)

diff --git a/user-data.sh.tftpl b/user-data.sh.tftpl
index 25346d5..58b9679 100644
--- a/user-data.sh.tftpl
+++ b/user-data.sh.tftpl
@@ -9,7 +9,7 @@ ready_marker="$state_dir/ami-ready"
 
 mkdir -p "$state_dir"
 
-write_bootstrap_config() {
+write_config() {
   mkdir -p /etc/nixos
   cat > /etc/nixos/flake.nix <<'EOF'
 {
@@ -17,14 +17,11 @@ write_bootstrap_config() {
     nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
   };
 
-  outputs = { nixpkgs, ... }: {
-    nixosConfigurations.ami-builder-bootstrap = nixpkgs.lib.nixosSystem {
-      system = "${system}";
-      modules = [
-        ({ config, lib, modulesPath, pkgs, ... }: {
+  outputs = { nixpkgs, ... }:
+    let
+      baseModule = { lib, modulesPath, pkgs, ... }: {
           imports = [ "$${modulesPath}/virtualisation/amazon-image.nix" ];
 
-          networking.hostName = "nixos-zfs-ec2-builder";
           networking.hostId = "6e69787a";
 
           nix.settings.experimental-features = [ "nix-command" "flakes" ];
@@ -38,11 +35,18 @@ write_bootstrap_config() {
           ];
 
           environment.systemPackages = [
+            pkgs.coreutils
             pkgs.git
             pkgs.rsync
+            pkgs.systemd
             pkgs.zfs
           ];
 
+          programs.bash.shellAliases = {
+            nixrb = "sudo nix --refresh flake update --flake /etc/nixos && sudo nixos-rebuild switch --flake /etc/nixos#nixos-zfs-ec2";
+            nixfu = "sudo nix --refresh flake update --flake /etc/nixos";
+          };
+
           services.amazon-ssm-agent.enable = true;
 
           users.users.builder-admin = {
@@ -53,42 +57,11 @@ write_bootstrap_config() {
           security.sudo.wheelNeedsPassword = false;
 
           system.stateVersion = "25.05";
-        })
-      ];
-    };
-  };
-}
-EOF
-}
+        };
 
-write_final_config() {
-  mkdir -p /etc/nixos
-  cat > /etc/nixos/flake.nix <<'EOF'
-{
-  inputs = {
-    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
-  };
-
-  outputs = { nixpkgs, ... }: {
-    nixosConfigurations.nixos-zfs-ec2 = nixpkgs.lib.nixosSystem {
-      system = "${system}";
-      modules = [
-        ({ config, lib, modulesPath, pkgs, ... }: {
-          imports = [ "$${modulesPath}/virtualisation/amazon-image.nix" ];
-
-          networking.hostName = "nixos-zfs-ec2";
-          networking.hostId = "6e69787a";
-
-          nix.settings.experimental-features = [ "nix-command" "flakes" ];
-          boot.supportedFilesystems = [ "zfs" ];
+      zfsNixModule = { lib, ... }: {
           boot.zfs.extraPools = [ "${nix_pool}" ];
-          boot.zfs.forceImportRoot = false;
-          swapDevices = [
-            {
-              device = "/swapfile";
-              size = 4096;
-            }
-          ];
+          systemd.services.amazon-init.enable = lib.mkForce false;
 
           fileSystems."/nix" = {
             device = "${nix_pool}/nix";
@@ -105,16 +78,9 @@ write_final_config() {
             after = [ "nix.mount" ];
             requires = [ "nix.mount" ];
           };
+        };
 
-          environment.systemPackages = [
-            pkgs.git
-            pkgs.rsync
-            pkgs.zfs
-          ];
-
-          services.amazon-ssm-agent.enable = true;
-          systemd.services.amazon-init.enable = lib.mkForce false;
-
+      userDataApplyModule = { config, lib, pkgs, ... }: {
           systemd.services.nixos-user-data-flake = {
             description = "Apply NixOS flake from EC2 user data";
             wantedBy = [ "multi-user.target" ];
@@ -166,25 +132,31 @@ RUNNER_EOF
                 "$${pkgs.runtimeShell}" "$runner"
             '';
           };
+        };
 
-          users.users.builder-admin = {
-            isNormalUser = true;
-            extraGroups = [ "wheel" ];
-            openssh.authorizedKeys.keys = ${ssh_keys};
-          };
-          security.sudo.wheelNeedsPassword = false;
-
-          system.stateVersion = "25.05";
-        })
+      mkHost = hostName: extraModules: nixpkgs.lib.nixosSystem {
+        system = "${system}";
+        modules = [
+          baseModule
+          ({ ... }: {
+            networking.hostName = hostName;
+          })
+        ] ++ extraModules;
+      };
+    in
+    {
+      nixosConfigurations.ami-builder-bootstrap = mkHost "nixos-zfs-ec2-builder" [];
+      nixosConfigurations.nixos-zfs-ec2 = mkHost "nixos-zfs-ec2" [
+        zfsNixModule
+        userDataApplyModule
       ];
     };
-  };
 }
 EOF
 }
 
 if [ ! -e "$bootstrap_marker" ]; then
-  write_bootstrap_config
+  write_config
   nixos-rebuild boot --flake /etc/nixos#ami-builder-bootstrap
   touch "$bootstrap_marker"
   systemctl reboot
@@ -224,7 +196,7 @@ if ! mountpoint -q /mnt/nixos-zfs-nix; then
   mount -t zfs "$nix_pool/nix" /mnt/nixos-zfs-nix
 fi
 
-write_final_config
+write_config
 nixos-rebuild boot --flake /etc/nixos#nixos-zfs-ec2
 
 rsync -aHAX --numeric-ids \