From ec5b6be824e06cd51d27d4fe210e3405315d903e Mon Sep 17 00:00:00 2001
From: Artturin <Artturin@artturin.com>
Date: Fri, 15 Sep 2023 07:45:10 +0300
Subject: [PATCH] nixos/lib/make-btrfs-fs: copy improvements from

https://git.sr.ht/~c00w/nixpkgs/tree/sdimagebtrfs/item/nixos/lib/make-btrfs-fs.nix

Co-authored-by: Colin L Rice <colin@daedrum.net>
---
 nixos/lib/make-btrfs-fs.nix | 40 +++++++++++++++++++++++++++----------
 1 file changed, 30 insertions(+), 10 deletions(-)

diff --git a/nixos/lib/make-btrfs-fs.nix b/nixos/lib/make-btrfs-fs.nix
index fceeea8d5755b16..270a8dfd7bcce94 100644
--- a/nixos/lib/make-btrfs-fs.nix
+++ b/nixos/lib/make-btrfs-fs.nix
@@ -7,6 +7,8 @@
 , lib
 # List of derivations to be included
 , storePaths
+# Whether or not to compress the resulting image with zstd
+, compressImage ? false, zstd
 # Shell commands to populate the ./files directory.
 # All files in that directory are copied to the root of the FS.
 , populateImageCommands ? ""
@@ -19,27 +21,45 @@ let
   sdClosureInfo = pkgs.buildPackages.closureInfo { rootPaths = storePaths; };
 in
 pkgs.stdenv.mkDerivation {
-  name = "btrfs-fs.img";
+  name = "btrfs-fs.img${lib.optionalString compressImage ".zst"}";
 
-  nativeBuildInputs = [ btrfs-progs ];
+  nativeBuildInputs = [ btrfs-progs ] ++ lib.optional compressImage zstd;
 
   buildCommand =
     ''
+      ${if compressImage then "img=temp.img" else "img=$out"}
+
       set -x
       (
           mkdir -p ./files
           ${populateImageCommands}
       )
 
-      mkdir -p ./files/nix/store
-      cp ${sdClosureInfo}/registration ./files/nix-path-registration
+      mkdir -p ./rootImage/nix/store
+
+      xargs -I % cp -a --reflink=auto % -t ./rootImage/nix/store/ < ${sdClosureInfo}/store-paths
+      (
+        GLOBIGNORE=".:.."
+        shopt -u dotglob
+
+        for f in ./files/*; do
+            cp -a --reflink=auto -t ./rootImage/ "$f"
+        done
+      )
+
+      cp ${sdClosureInfo}/registration ./rootImage/nix-path-registration
+
+      touch $img
+      mkfs.btrfs -L ${volumeLabel} -U ${uuid} -r ./rootImage --shrink $img
 
-      # Add the closures of the top-level store objects.
-      for p in $(cat ${sdClosureInfo}/store-paths); do
-        echo cp -r $p "./files/nix/store"
-      done
+      if ! fsck.btrfs $img; then
+        echo "--- Fsck failed for BTRFS image ---"
+        return 1
+      fi
 
-      touch $out
-      mkfs.btrfs -L ${volumeLabel} -U ${uuid} -r ./files --shrink $out
+      if [ ${builtins.toString compressImage} ]; then
+        echo "Compressing image"
+        zstd -v --no-progress ./$img -o $out
+      fi
     '';
 }