#!/bin/bash
# Build the content of a Unified Kernel Image addon in ./tree: the files of
# an initrd built by dracut that are not in both base initrds (virt and
# generic), so the addon works on top of any of them
set -euo pipefail

usage() {
	cat >&2 <<EOF
Usage: $0 --name NAME --kver KVER [OPTIONS]

  --name NAME             addon name
  --kver KVER             kernel version
  --dracut-modules LIST   dracut modules (default: systemd systemd-initrd)
  --drivers LIST          kernel drivers to add
  --keep REGEX            keep only the files matching REGEX
  --exclude REGEX         drop the files matching REGEX
  --subtract-addon FILE   drop the files of this addon too
  --nvidia-family FAMILY  keep the NVIDIA firmware of one GPU family
  --kernel-modules        the addon carries kernel modules
EOF
	exit 1
}

name=
kver=
dracut_modules='systemd systemd-initrd'
drivers=
keep='.'
exclude='^$'
subtract_addon=
nvidia_family=
with_modules=
while [ $# -gt 0 ]; do
	case "$1" in
	--name) name=$2; shift 2 ;;
	--kver) kver=$2; shift 2 ;;
	--dracut-modules) dracut_modules=$2; shift 2 ;;
	--drivers) drivers=$2; shift 2 ;;
	--keep) keep=$2; shift 2 ;;
	--exclude) exclude=$2; shift 2 ;;
	--subtract-addon) subtract_addon=$2; shift 2 ;;
	--nvidia-family) nvidia_family=$2; shift 2 ;;
	--kernel-modules) with_modules=1; shift ;;
	*) usage ;;
	esac
done
[ -n "$name" ] && [ -n "$kver" ] || usage

moddir=/usr/lib/modules/$kver
fwdir=/usr/lib/firmware

# The files present in both base initrds, plus the ones of the addon to
# subtract
base_files() {
	local f
	for f in virt generic; do
		mkdir base-$f
		(cd base-$f && lsinitrd --unpack "$moddir/uki-initrd-$f.img")
		(cd base-$f && find . -mindepth 1 | sort) > base-$f.list
	done
	comm -12 base-virt.list base-generic.list > base.list
	if [ -n "$subtract_addon" ]; then
		objcopy -O binary --only-section=.initrd "$subtract_addon" subtract.cpio
		cpio -it --quiet < subtract.cpio | sed 's|^|./|' | sort -u - base.list -o base.list
	fi
}

# Build the full initrd and list the files of the addon
initrd_files() {
	local d
	mkdir -p empty.conf.d
	dracut \
		--conf /dev/null \
		--confdir ./empty.conf.d \
		--kver "$kver" \
		--no-hostonly \
		--no-hostonly-cmdline \
		--reproducible \
		--no-compress \
		--no-early-microcode \
		--modules "$dracut_modules" \
		${drivers:+--add-drivers "$drivers"} \
		--force \
		./initrd.img > dracut.log 2>&1 || { cat dracut.log; exit 1; }
	cat dracut.log
	if grep -q 'dracut-install: ERROR' dracut.log; then
		exit 1
	fi
	mkdir full
	(cd full && cpio -id --quiet < ../initrd.img)
	(cd full && find . -mindepth 1 | sort) > full.list
	for d in $drivers; do
		if ! grep -qE "/$d\.ko(\.[a-z]+)?$" full.list; then
			echo "Driver $d not found"
			exit 1
		fi
	done
	comm -23 full.list base.list | grep -E "$keep" | grep -vE "$exclude" > addon.list
}

# nouveau declares the firmware of all the NVIDIA GPUs.  Every firmware
# directory belongs to exactly one GPU family, and the addon keeps the one
# of its family.  Of the GSP-RM firmware, nouveau uses the newest version
nvidia_firmware() {
	local -A families=(
		[legacy]='nvidia/(gk|gm|gp|gv)[0-9]'
		[turing]='nvidia/(tu[0-9]|ga100/)'
		[ampere]='nvidia/(ga10[2-7]|ad[0-9])'
		[blackwell]='nvidia/(gh|gb)[0-9]'
	)
	local chip family n
	for chip in $(sed -n 's|^\./usr/lib/firmware/nvidia/\([^/]*\).*|\1|p' full.list | sort -u); do
		n=0
		for family in "${families[@]}"; do
			if echo "nvidia/$chip/" | grep -qE "^$family"; then
				n=$((n + 1))
			fi
		done
		if [ $n -ne 1 ]; then
			echo "nvidia/$chip belongs to $n GPU families"
			exit 1
		fi
	done
	family=${families[$nvidia_family]:-}
	if [ -z "$family" ]; then
		echo "Unknown NVIDIA GPU family $nvidia_family"
		exit 1
	fi
	{ grep -vE '^\./usr/lib/firmware/nvidia/' addon.list || :
	  grep -E "^\./usr/lib/firmware/$family" addon.list || :
	} > addon.family.list
	{ grep -E '/gsp/[^/]*-[0-9][0-9.]*[0-9]\.bin' addon.family.list || :; } \
		| sed -E 's|^(.*)/gsp/[^/]*-([0-9][0-9.]*[0-9])\.bin.*$|\1 \2|' \
		| sort -k1,1 -k2,2V \
		| awk '{ newest[$1] = $2 } END { for (d in newest) print d, newest[d] }' > gsp.newest
	awk -v list=gsp.newest '
		BEGIN { while ((getline l < list) > 0) { split(l, a, " "); newest[a[1]] = a[2] } }
		match($0, /\/gsp\/[^\/]*-[0-9][0-9.]*[0-9]\.bin/) {
			d = substr($0, 1, RSTART - 1)
			v = substr($0, RSTART, RLENGTH); sub(/.*-/, "", v); sub(/\.bin$/, "", v)
			if (v != newest[d]) next
		}
		{ print }' addon.family.list > addon.list
}

# Copy the files of the addon.  The firmware links of a GPU family can point
# to the files of another one
copy_files() {
	local links l t src
	mkdir -p tree
	(cd full && cpio -pdm --quiet ../tree) < addon.list
	[ -n "$nvidia_family" ] || return 0
	while links=$(find tree$fwdir -xtype l) && [ -n "$links" ]; do
		for l in $links; do
			t=$(realpath -m -s "$(dirname "$l")/$(readlink "$l")")
			src="full${t#"$PWD"/tree}"
			test -e "$src" -o -L "$src"
			mkdir -p "$(dirname "$t")"
			cp -P "$src" "$t"
		done
	done
}

# Report the firmware declared by the drivers but not available, and fail
# when most of it is missing (for example, after a firmware package split)
driver_firmware() {
	local d m fw f found declared missing total=0
	for d in $drivers; do
		m=$(find "$moddir/kernel" -name "$d.ko*" | head -n1)
		declared=0
		missing=0
		for fw in $(modinfo -F firmware "$m"); do
			declared=$((declared + 1))
			found=
			for f in $fw $fw.xz $fw.zst; do
				if [ -e "$fwdir/$f" ]; then
					found=1
				fi
			done
			if [ -z "$found" ]; then
				missing=$((missing + 1))
				echo "$d: missing firmware $fw"
			fi
		done
		echo "$d: $declared firmware files declared, $missing missing"
		if [ $((missing * 2)) -gt $declared ]; then
			exit 1
		fi
		total=$((total + declared))
	done
	if [ $total -gt 0 ] && ! grep -q '^\./usr/lib/firmware/' addon.list; then
		echo "No firmware in the addon"
		exit 1
	fi
}

# The modules load on top of any of the base initrds, and the addon is marked
# for uki-addon-depmod.service
module_checks() {
	local f
	if ! grep -q '\.ko' addon.list; then
		echo "No kernel modules in the addon"
		exit 1
	fi
	mkdir -p "tree$moddir/uki-addons"
	touch "tree$moddir/uki-addons/$name"
	zcat "$moddir/symvers.gz" > Module.symvers
	for f in virt generic; do
		rm -rf merged
		cp -a base-$f merged
		cp -a tree/. merged/
		depmod -b merged -e -E Module.symvers "$kver" 2> depmod.log || { cat depmod.log; exit 1; }
		if grep 'unknown symbol' depmod.log; then
			exit 1
		fi
	done
	rm -rf merged
}

base_files
initrd_files
if [ -n "$nvidia_family" ]; then
	nvidia_firmware
fi
copy_files
if [ -n "$drivers" ]; then
	driver_firmware
fi
if [ -n "$with_modules" ]; then
	module_checks
fi
du -sh tree$fwdir 2>/dev/null || :
wc -l addon.list
