From b6ef9ffdfd54e91cce180c306852489f5ebc3098 Mon Sep 17 00:00:00 2001
From: Maximilian Bosch <maximilian@mbosch.me>
Date: Wed, 10 Apr 2024 10:04:43 +0200
Subject: [PATCH] nixos/udev: compress firmware with zstd if possible

Closes #267442

    $ nix path-info -Sh /nix/store/qj1dm7wfw5m3mxf1gn3fdm0az9y1h5ny-linux-firmware-20240312-xz
    /nix/store/qj1dm7wfw5m3mxf1gn3fdm0az9y1h5ny-linux-firmware-20240312-xz	440.3M
    $ nix path-info -Sh /nix/store/c3szcjxb3g990dbiz7llwmkaf0bi98j2-linux-firmware-20240312-zstd
    /nix/store/c3szcjxb3g990dbiz7llwmkaf0bi98j2-linux-firmware-20240312-zstd	460.6M

This is an increase of 4.4%, but OTOH zstd has a significantly higher
decompression speed[1].

[1] https://gregoryszorc.com/blog/2017/03/07/better-compression-with-zstandard/
---
 .../manual/release-notes/rl-2405.section.md   |  4 ++
 nixos/modules/services/hardware/udev.nix      | 14 ++++--
 .../kernel/compress-firmware-xz.nix           | 29 -------------
 .../kernel/compress-firmware.nix              | 43 +++++++++++++++++++
 pkgs/top-level/all-packages.nix               |  4 +-
 5 files changed, 60 insertions(+), 34 deletions(-)
 delete mode 100644 pkgs/build-support/kernel/compress-firmware-xz.nix
 create mode 100644 pkgs/build-support/kernel/compress-firmware.nix

diff --git a/nixos/doc/manual/release-notes/rl-2405.section.md b/nixos/doc/manual/release-notes/rl-2405.section.md
index a3a9980bc58f..6b45c9717083 100644
--- a/nixos/doc/manual/release-notes/rl-2405.section.md
+++ b/nixos/doc/manual/release-notes/rl-2405.section.md
@@ -53,6 +53,10 @@ Use `services.pipewire.extraConfig` or `services.pipewire.configPackages` for Pi
 - A new option `system.etc.overlay.enable` was added. If enabled, `/etc` is
   mounted via an overlayfs instead of being created by a custom perl script.
 
+- For each supporting version of the Linux kernel firmware blobs and kernel modules
+  are compressed with zstd. For firmware blobs this means an increase of 4.4% in size, however
+  a significantly higher decompression speed.
+
 - NixOS AMIs are now uploaded regularly to a new AWS Account.
   Instructions on how to use them can be found on <https://nixos.github.io/amis>.
   We are working on integration the data into the NixOS homepage.
diff --git a/nixos/modules/services/hardware/udev.nix b/nixos/modules/services/hardware/udev.nix
index 3db661644281..62603d20e2d3 100644
--- a/nixos/modules/services/hardware/udev.nix
+++ b/nixos/modules/services/hardware/udev.nix
@@ -167,10 +167,16 @@ let
       mv etc/udev/hwdb.bin $out
     '';
 
-  compressFirmware = firmware: if (config.boot.kernelPackages.kernelAtLeast "5.3" && (firmware.compressFirmware or true)) then
-    pkgs.compressFirmwareXz firmware
-  else
-    id firmware;
+  compressFirmware = firmware:
+    let
+      inherit (config.boot.kernelPackages) kernelAtLeast;
+    in
+      if ! (firmware.compressFirmware or true) then
+        firmware
+      else
+        if kernelAtLeast "5.19" then pkgs.compressFirmwareZstd firmware
+        else if kernelAtLeast "5.3" then pkgs.compressFirmwareXz firmware
+        else firmware;
 
   # Udev has a 512-character limit for ENV{PATH}, so create a symlink
   # tree to work around this.
diff --git a/pkgs/build-support/kernel/compress-firmware-xz.nix b/pkgs/build-support/kernel/compress-firmware-xz.nix
deleted file mode 100644
index cb9ce7a71338..000000000000
--- a/pkgs/build-support/kernel/compress-firmware-xz.nix
+++ /dev/null
@@ -1,29 +0,0 @@
-{ runCommand, lib }:
-
-firmware:
-
-let
-  args = {
-    allowedRequisites = [];
-  } // lib.optionalAttrs (firmware ? meta) { inherit (firmware) meta; };
-in
-
-runCommand "${firmware.name}-xz" args ''
-  mkdir -p $out/lib
-  (cd ${firmware} && find lib/firmware -type d -print0) |
-      (cd $out && xargs -0 mkdir -v --)
-  (cd ${firmware} && find lib/firmware -type f -print0) |
-      (cd $out && xargs -0rtP "$NIX_BUILD_CORES" -n1 \
-          sh -c 'xz -9c -T1 -C crc32 --lzma2=dict=2MiB "${firmware}/$1" > "$1.xz"' --)
-  (cd ${firmware} && find lib/firmware -type l) | while read link; do
-      target="$(readlink "${firmware}/$link")"
-      if [ -f "${firmware}/$link" ]; then
-        ln -vs -- "''${target/^${firmware}/$out}.xz" "$out/$link.xz"
-      else
-        ln -vs -- "''${target/^${firmware}/$out}" "$out/$link"
-      fi
-  done
-
-  echo "Checking for broken symlinks:"
-  find -L $out -type l -print -execdir false -- '{}' '+'
-''
diff --git a/pkgs/build-support/kernel/compress-firmware.nix b/pkgs/build-support/kernel/compress-firmware.nix
new file mode 100644
index 000000000000..0949036d5127
--- /dev/null
+++ b/pkgs/build-support/kernel/compress-firmware.nix
@@ -0,0 +1,43 @@
+{ runCommand, lib, type ? "zstd", zstd }:
+
+firmware:
+
+let
+  compressor = {
+    xz = {
+      ext = "xz";
+      nativeBuildInputs = [ ];
+      cmd = file: target: ''xz -9c -T1 -C crc32 --lzma2=dict=2MiB "${file}" > "${target}"'';
+    };
+    zstd = {
+      ext = "zst";
+      nativeBuildInputs = [ zstd ];
+      cmd = file: target: ''zstd -T1 -19 --long --check -f "${file}" -o "${target}"'';
+    };
+  }.${type} or (throw "Unsupported compressor type for firmware.");
+
+  args = {
+    allowedRequisites = [];
+    inherit (compressor) nativeBuildInputs;
+  } // lib.optionalAttrs (firmware ? meta) { inherit (firmware) meta; };
+in
+
+runCommand "${firmware.name}-${type}" args ''
+  mkdir -p $out/lib
+  (cd ${firmware} && find lib/firmware -type d -print0) |
+      (cd $out && xargs -0 mkdir -v --)
+  (cd ${firmware} && find lib/firmware -type f -print0) |
+      (cd $out && xargs -0rtP "$NIX_BUILD_CORES" -n1 \
+          sh -c '${compressor.cmd "${firmware}/$1" "$1.${compressor.ext}"}' --)
+  (cd ${firmware} && find lib/firmware -type l) | while read link; do
+      target="$(readlink "${firmware}/$link")"
+      if [ -f "${firmware}/$link" ]; then
+        ln -vs -- "''${target/^${firmware}/$out}.${compressor.ext}" "$out/$link.${compressor.ext}"
+      else
+        ln -vs -- "''${target/^${firmware}/$out}" "$out/$link"
+      fi
+  done
+
+  echo "Checking for broken symlinks:"
+  find -L $out -type l -print -execdir false -- '{}' '+'
+''
diff --git a/pkgs/top-level/all-packages.nix b/pkgs/top-level/all-packages.nix
index e13755dc7f04..7b8c7a59dde7 100644
--- a/pkgs/top-level/all-packages.nix
+++ b/pkgs/top-level/all-packages.nix
@@ -1314,7 +1314,9 @@ with pkgs;
 
   makeBinaryWrapper = callPackage ../build-support/setup-hooks/make-binary-wrapper { };
 
-  compressFirmwareXz = callPackage ../build-support/kernel/compress-firmware-xz.nix { };
+  compressFirmwareXz = callPackage ../build-support/kernel/compress-firmware.nix { type = "xz"; };
+
+  compressFirmwareZstd = callPackage ../build-support/kernel/compress-firmware.nix { type = "zstd"; };
 
   makeModulesClosure = { kernel, firmware, rootModules, allowMissing ? false }:
     callPackage ../build-support/kernel/modules-closure.nix {