path: root/epistemia/systems/linux-initrd.scm

;;; SPDX-License-Identifier: GPL-3.0-or-later
;; AI disclosure: LLM slop, but it works. This was made before I got more familiar with Guile Scheme and took many retries to get working.
;; TODO: rewrite this later by deleting this and rewriting by reading the main initrd and my previous pre-guix initrd script. could maybe be deduplicated a lot via inherits.
(define-module (epistemia systems linux-initrd)
  #:use-module (guix gexp)
  #:use-module (guix utils)
  #:use-module ((guix store) #:select (%store-prefix))
  #:use-module ((guix derivations) #:select (derivation->output-path))
  #:use-module (guix modules)
  #:use-module (gnu packages compression)
  #:use-module (gnu packages disk)
  #:use-module (gnu packages guile)
  #:use-module (gnu packages linux)
  #:use-module (gnu packages file-systems)
  #:use-module (gnu system file-systems)
  #:use-module (gnu system mapped-devices)
  #:use-module (gnu system keyboard)
  #:use-module (gnu system linux-initrd)
  #:use-module (ice-9 match)
  #:use-module (ice-9 regex)
  #:use-module (ice-9 vlist)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-26)
  #:export (epistemia-zfs-initrd))

(define (flat-linux-module-directory* linux extra-packages modules)
  "Return a flat directory containing the Linux kernel modules listed in
MODULES."
  (define imported-modules
    (source-module-closure '((gnu build linux-modules)
                             (guix build utils))))

  (define build-exp
    (with-imported-modules imported-modules
      (with-extensions (list guile-zlib guile-zstd)
        #~(begin
            (use-modules (gnu build linux-modules)
                         (guix build utils)
                         (rnrs io ports)
                         (srfi srfi-1)
                         (srfi srfi-26)
                         (ice-9 match))

            (define module-dirs
              (map (lambda (pkg) (string-append pkg "/lib/modules"))
                   (cons #$linux '#$extra-packages)))

            (define builtin-modules
              (match (find-files (string-append #$linux "/lib/modules")
                                 (lambda (file stat)
                                   (string=? (basename file) "modules.builtin")))
                ((file . _)
                 (call-with-input-file file
                   (lambda (port)
                     (map file-name->module-name
                          (string-tokenize (get-string-all port))))))
                (_ '())))

            (define modules-to-lookup
              (lset-difference string=? '#$modules builtin-modules))

            (define (lookup-in-dirs name)
              (let loop ((dirs module-dirs))
                (match dirs
                  ((dir . rest)
                   (let ((candidates 
                          (find-files dir 
                                      (lambda (file stat)
                                        (let ((base (basename file)))
                                          (or (string=? base (string-append name ".ko"))
                                              (string=? base (string-append name ".ko.gz"))
                                              (string=? base (string-append name ".ko.zst"))))))))
                     (match candidates
                       ((first . _) first)
                       (() (loop rest)))))
                  (() #f))))

            (define modules
              (let ((found-modules (map lookup-in-dirs modules-to-lookup)))
                (append (filter identity found-modules)
                        (recursive-module-dependencies
                         (filter identity found-modules)
                         #:lookup-module lookup-in-dirs))))

            (define (maybe-uncompress file)
              (cond
               ((string-contains file ".ko.gz")
                (invoke #+(file-append gzip "/bin/gunzip") file))
               ((string-contains file ".ko.zst")
                (invoke #+(file-append zstd "/bin/zstd") "-d" file))))

            (mkdir #$output)
            (for-each (lambda (module)
                        (when module
                          (let ((out-module
                                 (string-append #$output "/"
                                                (basename module))))
                            (format #t "copying '~a'...~%" module)
                            (copy-file module out-module)
                            (maybe-uncompress out-module))))
                      (delete-duplicates modules))

            (write-module-name-database #$output)))))

  (computed-file "linux-modules-combined" build-exp))

(define* (zfs-raw-initrd file-systems
                         #:key
                         (linux linux-libre)
                         (zfs-package zfs)
                         (linux-modules '())
                         (pre-mount #t)
                         (mapped-devices '())
                         (keyboard-layout #f)
                         (helper-packages '())
                         qemu-networking?
                         volatile-root?
                         (on-error 'debug)
                         (monkey-patch? #t)
                         #:allow-other-keys)
  
  (define device-mapping-commands
    (map (lambda (md)
           (let* ((source  (mapped-device-source md))
                  (targets (mapped-device-targets md))
                  (type    (mapped-device-type md))
                  (open    (mapped-device-kind-open type)))
             (apply open source targets
                    (mapped-device-arguments md))))
         mapped-devices))

  (define file-system-scan-commands
    (let ((file-system-types (map file-system-type file-systems)))
      (if (member "btrfs" file-system-types)
          #~((system* (string-append #$btrfs-progs/static "/bin/btrfs")
                      "device" "scan"))
          #~())))

  (define zfs-import-commands
    #~(begin
        (format #t "Epistemia: Importing ZFS pools...~%")
        (false-if-exception 
         (system* (string-append #$zfs-package "/sbin/zpool")
                  "import" "-a" "-N" "-f" "-d" "/dev"))
        (format #t "Epistemia: Imported ZFS pools:~%")
        (false-if-exception
         (system* (string-append #$zfs-package "/sbin/zpool")
                  "status"))
        (sleep 2)
        #t))

  (define kodir
    (flat-linux-module-directory* linux 
				  (list (gexp-input zfs-package "module")) 
				  linux-modules))

  (expression->initrd
   (with-imported-modules (source-module-closure
                           '((gnu build linux-boot)
                             (guix build utils)
                             (guix build bournish)
                             (gnu system file-systems)
                             (gnu build file-systems)))
     #~(begin
         (use-modules (gnu build linux-boot)
                      (gnu system file-systems)
                      ((guix build utils) #:hide (delete))
                      (guix build bournish)
                      (srfi srfi-1)
                      (srfi srfi-13)
                      (srfi srfi-26))
         
         ;; We must patch BOTH modules because linux-boot likely already
         ;; imported the function.
         #$(if monkey-patch?
               #~(let* ((fs-mod   (resolve-module '(gnu build file-systems)))
                        (boot-mod (resolve-module '(gnu build linux-boot)))
                        (orig-canon (module-ref fs-mod 'canonicalize-device-spec))
                        (new-canon
                         (lambda (spec)
                           ;; ZFS datasets in Guix config are either strings or file-system-labels.
                           ;; If it's a label record, the first field is the label string.
                           (let ((device (cond
                                          ((string? spec) spec)
                                          ((and (struct? spec) 
                                                (string? (struct-ref spec 0)))
                                           (struct-ref spec 0))
                                          (else #f))))
                             (if (and device
                                      (or (string-contains device "zroot/")
                                          (string-contains device "zfs")))
                                 (begin
                                   (format #t "Epistemia: ZFS bypass for ~s~%" device)
                                   device)
                                 (orig-canon spec))))))
                   
                   ;; 1. Overwrite the definition source
                   (module-set! fs-mod 'canonicalize-device-spec new-canon)
                   
                   ;; 2. Overwrite the consumer's binding (Crucial for boot-system)
                   (module-set! boot-mod 'canonicalize-device-spec new-canon)
                   
                   (format #t "Epistemia: ZFS Monkey Patch Applied.~%"))
               #~#t)

         (with-output-to-port (%make-void-port "w")
           (lambda ()
             (set-path-environment-variable "PATH" '("bin" "sbin")
                                            '#$helper-packages)))

         (parameterize ((current-warning-port (%make-void-port "w")))
           (boot-system #:mounts
                        (map spec->file-system
                             '#$(map file-system->spec file-systems))
                        #:pre-mount (lambda ()
                                      (and #$pre-mount
                                           #$@device-mapping-commands
                                           #$@file-system-scan-commands
                                           #$zfs-import-commands))
                        #:linux-modules '#$linux-modules
                        #:linux-module-directory '#$kodir
                        #:keymap-file #+(and=> keyboard-layout
                                               keyboard-layout->console-keymap)
                        #:qemu-guest-networking? #$qemu-networking?
                        #:volatile-root? '#$volatile-root?
                        #:on-error '#$on-error))))
   #:name "zfs-initrd"))

(define* (epistemia-zfs-initrd file-systems
                               #:key
                               (linux linux-libre)
                               (zfs-package zfs)
                               (linux-modules '())
                               (mapped-devices '())
                               (keyboard-layout #f)
                               qemu-networking?
                               volatile-root?
                               (extra-modules '())
                               (on-error 'debug)
                               (monkey-patch? #t)
                               #:allow-other-keys)
  
  (define linux-modules*
    (cons "zfs"
          `(,@linux-modules
            ,@(file-system-modules file-systems)
            ,@(if volatile-root? '("overlay") '())
            ,@extra-modules)))

  (define helper-packages
    (cons zfs-package
          (append (file-system-packages file-systems
                                        #:volatile-root? volatile-root?)
                  (if keyboard-layout
                      (list loadkeys-static)
                      '()))))

  (zfs-raw-initrd file-systems
                  #:linux linux
                  #:zfs-package zfs-package
                  #:linux-modules linux-modules*
                  #:mapped-devices mapped-devices
                  #:helper-packages helper-packages
                  #:keyboard-layout keyboard-layout
                  #:qemu-networking? qemu-networking?
                  #:volatile-root? volatile-root?
                  #:on-error on-error
                  #:monkey-patch? monkey-patch?))