From baa3956ed53c2e4a9015198af6a963289f16dd0f Mon Sep 17 00:00:00 2001
From: Jack O'Sullivan <jackos1998@gmail.com>
Date: Sat, 12 Feb 2022 23:44:24 +0000
Subject: [PATCH] Add tmproot defaults and ignores (and fix early binds)

---
 modules/tmproot.nix | 86 +++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 79 insertions(+), 7 deletions(-)

diff --git a/modules/tmproot.nix b/modules/tmproot.nix
index 2b69897..f8a6e16 100644
--- a/modules/tmproot.nix
+++ b/modules/tmproot.nix
@@ -1,6 +1,7 @@
-{ lib, pkgs, inputs, config, ... }:
+{ lib, pkgs, inputs, config, utils, ... }:
   let
-    inherit (lib) concatStringsSep mkIf mkDefault mkMerge mkVMOverride;
+    inherit (builtins) elem;
+    inherit (lib) concatStringsSep concatMap concatMapStringsSep mkIf mkDefault mkMerge mkForce mkVMOverride;
     inherit (lib.my) mkOpt mkBoolOpt mkVMOverride' dummyOption;
 
     cfg = config.my.tmproot;
@@ -54,14 +55,12 @@
   in {
     imports = [ inputs.impermanence.nixosModule ];
 
-    options = {
-      my.tmproot = with lib.types; {
+    options = with lib.types; {
+      my.tmproot = {
         enable = mkBoolOpt true;
         persistDir = mkOpt str "/persist";
         size = mkOpt str "2G";
-        ignoreUnsaved = mkOpt (listOf str) [
-          "/tmp"
-        ];
+        ignoreUnsaved = mkOpt (listOf str) [];
       };
 
       # Forward declare options that won't exist until the VM module is actually imported
@@ -77,16 +76,77 @@
             assertion = config.fileSystems ? "${cfg.persistDir}";
             message = "The 'fileSystems' option does not specify your persistence file system (${cfg.persistDir}).";
           }
+          {
+            # I mean you probably _could_, but if you're doing tmproot... come on
+            assertion = !config.users.mutableUsers;
+            message = "users.mutableUsers is incompatible with tmproot";
+          }
+        ];
+
+        my.tmproot.ignoreUnsaved = [
+          "/tmp"
+
+          # setup-etc.pl will create this for us
+          "/etc/NIXOS"
+
+          # Once mutableUsers is disabled, we should be all clear here
+          "/etc/passwd"
+          "/etc/group"
+          "/etc/shadow"
+          "/etc/subuid"
+          "/etc/subgid"
+
+          # Lock file for /etc/{passwd,shadow}
+          "/etc/.pwd.lock"
+
+          # systemd last updated? I presume they'll get updated on boot...
+          "/etc/.updated"
+          "/var/.updated"
+
+          # Specifies obsolete files that should be deleted on activation - we'll never have those!
+          "/etc/.clean"
         ];
 
         environment.systemPackages = [
           (pkgs.writeScriptBin "tmproot-unsaved" showUnsaved)
         ];
 
+        # Catch non-existent source directories that are needed for boot (see `pathsNeededForBoot` in
+        # nixos/lib/util.nix). We do this by monkey-patching the `waitDevice` function that would otherwise hang.
+        boot.initrd.postDeviceCommands =
+          ''
+            ensurePersistSource() {
+              [ -e "/mnt-root$1" ] && return
+              echo "Persistent source directory $1 does not exist, creating..."
+              install -dm "$2" "/mnt-root$1" || fail
+            }
+
+            _waitDevice() {
+              local device="$1"
+
+              ${concatMapStringsSep " || \\\n  " (d:
+                let
+                  sourceDir = "${d.persistentStoragePath}${d.directory}";
+                in
+                  ''([ "$device" = "/mnt-root${sourceDir}" ] && ensurePersistSource "${sourceDir}" "${d.mode}")'')
+                config.environment.persistence."${cfg.persistDir}".directories}
+
+              waitDevice "$@"
+            }
+
+            type waitDevice > /dev/null || (echo "waitDevice is missing!"; fail)
+            alias waitDevice=_waitDevice
+          '';
+
         environment.persistence."${cfg.persistDir}" = {
           hideMounts = mkDefault true;
           directories = [
             "/var/log"
+            # In theory we'd include only the files needed individually (i.e. the {U,G}ID map files that track deleted
+            # users and groups), but `update-users-groups.pl` actually deletes the original files for "atomic update".
+            # Also the script runs before impermanence does.
+            "/var/lib/nixos"
+            "/var/lib/systemd"
           ];
           files = [
             "/etc/machine-id"
@@ -99,7 +159,19 @@
           diskImage = "./.vms/${config.system.name}-persist.qcow2";
         };
       }
+      (mkIf config.services.openssh.enable {
+        environment.persistence."${cfg.persistDir}".files =
+          concatMap (k: [ k.path "${k.path}.pub" ]) config.services.openssh.hostKeys;
+      })
+      (mkIf config.networking.resolvconf.enable {
+        my.tmproot.ignoreUnsaved = [ "/etc/resolv.conf" ];
+      })
+      (mkIf config.security.doas.enable {
+        my.tmproot.ignoreUnsaved = [ "/etc/doas.conf" ];
+      })
       (mkIf config.my.boot.isDevVM {
+        my.tmproot.ignoreUnsaved = [ "/nix" ];
+
         fileSystems = mkVMOverride {
           "/" = mkVMOverride' rootDef;
           # Hijack the "root" device for persistence in the VM