From 7eab0aa0b7eb293c8a41a4319cd9e2045b78c234 Mon Sep 17 00:00:00 2001
From: Tristan Ross <tristan.ross@midstall.com>
Date: Sat, 4 Jan 2025 08:14:21 -0800
Subject: [PATCH] starfive visionfive2: allow uboot and opensbi source
 overrides

---
 starfive/visionfive/v2/default.nix  |  4 ++
 starfive/visionfive/v2/firmware.nix | 70 +++++++++++++++++++----------
 starfive/visionfive/v2/sd-image.nix | 10 ++---
 starfive/visionfive/v2/spl-tool.nix |  4 +-
 4 files changed, 57 insertions(+), 31 deletions(-)

diff --git a/starfive/visionfive/v2/default.nix b/starfive/visionfive/v2/default.nix
index 3807b6b..a42c606 100644
--- a/starfive/visionfive/v2/default.nix
+++ b/starfive/visionfive/v2/default.nix
@@ -5,6 +5,10 @@
   ...
 }:
 {
+  imports = [
+    ./firmware.nix
+  ];
+
   boot = {
     consoleLogLevel = lib.mkDefault 7;
     # Switch to default as soon they reach >= 6.11
diff --git a/starfive/visionfive/v2/firmware.nix b/starfive/visionfive/v2/firmware.nix
index eabeb6c..06299de 100644
--- a/starfive/visionfive/v2/firmware.nix
+++ b/starfive/visionfive/v2/firmware.nix
@@ -1,26 +1,50 @@
-{ callPackage
-, writeShellApplication
-, stdenv
-, mtdutils
-}:
-
-rec {
-  opensbi = callPackage ./opensbi.nix { };
-  uboot = callPackage ./uboot.nix { inherit opensbi; };
-  updater-flash = writeShellApplication {
-    name = "visionfive2-firmware-update-flash";
-    runtimeInputs = [ mtdutils ];
-    text = ''
-      flashcp -v ${uboot}/u-boot-spl.bin.normal.out /dev/mtd0
-      flashcp -v ${uboot}/u-boot.itb /dev/mtd2
-    '';
+{ config, pkgs, lib, ... }:
+let
+  cfg = config.hardware.visionfive2;
+in
+{
+  options = {
+    hardware.visionfive2 = {
+      opensbi.src = lib.mkOption {
+        description = "VisionFive2 OpenSBI source";
+        type = lib.types.nullOr lib.types.package;
+        default = null;
+      };
+      uboot.src = lib.mkOption {
+        description = "VisionFive2 U-boot source";
+        type = lib.types.nullOr lib.types.package;
+        default = null;
+      };
+    };
   };
-  updater-sd = writeShellApplication {
-    name = "visionfive2-firmware-update-sd";
-    runtimeInputs = [ ];
-    text = ''
-      dd if=${uboot}/u-boot-spl.bin.normal.out of=/dev/mmcblk0p1 conv=fsync
-      dd if=${uboot}/u-boot.itb of=/dev/mmcblk0p2 conv=fsync
-    '';
+
+  config = {
+    system.build = {
+      opensbi = (pkgs.callPackage ./opensbi.nix {}).overrideAttrs (f: p: {
+        src = if cfg.opensbi.src != null then cfg.opensbi.src else p.src;
+      });
+
+      uboot = (pkgs.callPackage ./uboot.nix { inherit (config.system.build) opensbi; }).overrideAttrs (f: p: {
+        src = if cfg.uboot.src != null then cfg.uboot.src else p.src;
+      });
+
+      updater-flash = pkgs.writeShellApplication {
+        name = "visionfive2-firmware-update-flash";
+        runtimeInputs = [ pkgs.mtdutils ];
+        text = ''
+          flashcp -v ${config.system.build.uboot}/u-boot-spl.bin.normal.out /dev/mtd0
+          flashcp -v ${config.system.build.uboot}/u-boot.itb /dev/mtd2
+        '';
+      };
+
+      updater-sd = pkgs.writeShellApplication {
+        name = "visionfive2-firmware-update-sd";
+        runtimeInputs = [ ];
+        text = ''
+          dd if=${config.system.build.uboot}/u-boot-spl.bin.normal.out of=/dev/mmcblk0p1 conv=fsync
+          dd if=${config.system.build.uboot}/u-boot.itb of=/dev/mmcblk0p2 conv=fsync
+        '';
+      };
+    };
   };
 }
diff --git a/starfive/visionfive/v2/sd-image.nix b/starfive/visionfive/v2/sd-image.nix
index d66bd2c..a14b87b 100644
--- a/starfive/visionfive/v2/sd-image.nix
+++ b/starfive/visionfive/v2/sd-image.nix
@@ -1,7 +1,5 @@
 { config, pkgs, modulesPath, ... }:
-
-let firmware = pkgs.callPackage ./firmware.nix { };
-in {
+{
   imports = [
     "${modulesPath}/profiles/base.nix"
     "${modulesPath}/installer/sd-card/sd-image.nix"
@@ -36,10 +34,10 @@ in {
       EOF
 
       eval $(partx $img -o START,SECTORS --nr 1 --pairs)
-      dd conv=notrunc if=${firmware.uboot}/u-boot-spl.bin.normal.out of=$img seek=$START count=$SECTORS
+      dd conv=notrunc if=${config.system.build.uboot}/u-boot-spl.bin.normal.out of=$img seek=$START count=$SECTORS
 
       eval $(partx $img -o START,SECTORS --nr 2 --pairs)
-      dd conv=notrunc if=${firmware.uboot}/u-boot.itb of=$img seek=$START count=$SECTORS
+      dd conv=notrunc if=${config.system.build.uboot}/u-boot.itb of=$img seek=$START count=$SECTORS
     '';
 
     populateRootCommands = ''
@@ -48,5 +46,5 @@ in {
     '';
   };
 
-  environment.systemPackages = [ firmware.updater-flash ];
+  environment.systemPackages = [ config.system.build.updater-flash ];
 }
diff --git a/starfive/visionfive/v2/spl-tool.nix b/starfive/visionfive/v2/spl-tool.nix
index 3949845..cd0a77b 100644
--- a/starfive/visionfive/v2/spl-tool.nix
+++ b/starfive/visionfive/v2/spl-tool.nix
@@ -1,6 +1,6 @@
 { stdenv, fetchFromGitHub }:
 
-stdenv.mkDerivation rec{
+stdenv.mkDerivation (finalAttrs: {
   pname = "spi_tool";
   version = "0x01010101";
   src = fetchFromGitHub {
@@ -15,4 +15,4 @@ stdenv.mkDerivation rec{
     mkdir -p $out/bin
     cp spl_tool $out/bin
   '';
-}
+})