#
# Copyright 2023 Google LLC
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# version 2 as published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#

EAPI=7

FEATURES="xattr"
MODULES_OPTIONAL_USE="driver"
inherit flag-o-matic cos-linux-mod linux-mod multilib readme.gentoo-r1 \
	systemd toolchain-funcs unpacker user-info

NV_URI="https://download.nvidia.com/XFree86"

DESCRIPTION="NVIDIA Accelerated Graphics Driver"
HOMEPAGE="https://www.nvidia.com/download/index.aspx"
SRC_URI="${NV_URI}/NVIDIA-kernel-module-source/NVIDIA-kernel-module-source-${PV}.tar.xz
${NV_URI}/Linux-x86_64/${PV}/NVIDIA-Linux-x86_64-${PV}.run"

S="${WORKDIR}"

LICENSE="NVIDIA-r2 BSD BSD-2 GPL-2 MIT openssl"

KEYWORDS="*"
SLOT="0/${PV%%.*}"
IUSE="+driver +clang +llvm_ias +xattr -install-nvidia-usr"

REQUIRED_USE="
	llvm_ias? ( clang )
"

DEPEND=""
RDEPEND="${DEPEND}"
BDEPEND="dev-libs/openssl"

QA_PREBUILT="lib/firmware/* opt/nvidia/*"

pkg_setup() {
	MODULE_NAMES="
		nvidia(nvidia/${PV}:kernel-module-source:kernel-module-source/kernel-open)
		nvidia-drm(nvidia/${PV}:kernel-module-source:kernel-module-source/kernel-open)
		nvidia-modeset(nvidia/${PV}:kernel-module-source:kernel-module-source/kernel-open)
		nvidia-peermem(nvidia/${PV}:kernel-module-source:kernel-module-source/kernel-open)
		nvidia-uvm(nvidia/${PV}:kernel-module-source:kernel-module-source/kernel-open)"
	KBUILD_OUTPUT="${KERNEL_DIR}/build"
	linux-mod_pkg_setup
	BUILD_PARAMS="CC=${CC} CXX=${CBUILD}-gcc LD=${LD} KERNEL_SOURCES=${KV_DIR} KBUILD_OUTPUT=${KBUILD_OUTPUT} NV_VERBOSE=1 IGNORE_CC_MISMATCH=yes SYSSRC=${KV_DIR} SYSOUT=${KV_OUT_DIR}"
	BUILD_TARGETS="modules"
}

src_prepare() {
	mv NVIDIA-kernel-module-source-${PV} kernel-module-source || die
	eapply "${FILESDIR}"/nvidia-drivers-525.60.13-utils.patch
	eapply "${FILESDIR}"/nvidia-drivers-525.60.13-ld.patch
	use clang || cros_use_gcc
	default
}

src_compile() {
	tc-export AR CC CXX LD OBJCOPY OBJDUMP

	if use llvm_ias; then
		export LLVM_IAS=1
	fi
	cros_allow_gnu_build_tools
	BUILD_FIXES="CFLAGS_MODULE='-Wno-error=strict-prototypes'" linux-mod_src_compile
}

sign_module() {
	local module_file="$1"

	"${KBUILD_OUTPUT}"/scripts/sign-file \
		sha256 \
		"${KBUILD_OUTPUT}"/certs/signing_key.pem \
		"${KBUILD_OUTPUT}"/certs/signing_key.x509 \
		"${module_file}"
}

sign_firmware() {
	local firmware_file="$1"

	# produce raw signature
	openssl dgst -sign "${KBUILD_OUTPUT}"/certs/signing_key.pem -keyform PEM -sha256 -out "${firmware_file}.sign" -binary "${firmware_file}"
	# construct prefix
	local -r subjectKeyId=$(openssl x509 -in "${KBUILD_OUTPUT}"/certs/signing_key.pem -noout -ext subjectKeyIdentifier | xargs | sed 's/://g')
	local prefixBytes="030204${subjectKeyId: -8}0180"
	# assemble ima signature
	local -r detachedSignature=$(cat "${firmware_file}.sign" | xxd -p | tr -d '\n')
	setfattr -n security.ima -v "0x${prefixBytes}${detachedSignature}" "${firmware_file}"
}

install_nvidia_usr_components() {

	local libdir=$(get_libdir)
	local skip_types=(
		SYSTEMD_UNIT_SYMLINK
		NVIDIA_MODPROBE.\*
		MANPAGE
		INTERNAL_UTILITY.\*
		XMODULE_SHARED_LIB
		UTILITY_BIN_SYMLINK
	)

	# mimic nvidia-installer by reading .manifest to install files
	# 0:file 1:perms 2:type 3+:subtype/arguments -:module
	local m into
	while IFS=' ' read -ra m; do
		! [[ ${#m[@]} -ge 2 && ${m[-1]} =~ MODULE: ]] ||
		[[ " ${m[2]}" =~ ^(\ ${skip_types[*]/%/|\\} )$ ]] ||
		[[ ${m[0]} =~ nvngx.dll ]] && continue

		case ${m[2]} in
			GBM_BACKEND_LIB_SYMLINK) m[4]=../${m[4]};; # missing ../
			VDPAU_SYMLINK) m[4]=vdpau/; m[5]=${m[5]#vdpau/};; # .so to vdpau/
			GLX_MODULE_SYMLINK) m[4]=../${m[4]};; # missing ../
		esac

		if [[ ${m[3]} == COMPAT32 ]]; then
			continue
		elif [[ ${m[2]} =~ GBM_BACKEND_LIB_SYMLINK ]]; then
			into=/opt/nvidia/${PV}/${libdir}/gbm
		elif [[ ${m[2]} =~ _BINARY ]]; then
			into=/opt/nvidia/${PV}/bin
		elif [[ ${m[2]} =~ _LIB$|_SYMLINK$ ]]; then
			into=/opt/nvidia/${PV}/${libdir}
		else
			continue
		fi
		[[ ${m[3]: -2} == ?/ ]] && into+=/${m[3]%/}
		[[ ${m[4]: -2} == ?/ ]] && into+=/${m[4]%/}

		if [[ ${m[2]} =~ _SYMLINK$ ]]; then
			[[ ${m[4]: -1} == / ]] && m[4]=${m[5]}
			dosym ${m[4]} ${into}/${m[0]}
			continue
		fi
		[[ ${m[0]} =~ ^libnvidia-ngx.so|^libnvidia-egl-gbm.so ]] &&
			dosym ${m[0]} ${into}/${m[0]%.so*}.so.1 # soname not in .manifest

		printf -v m[1] %o $((m[1] | 0200)) # 444->644
		insopts -m${m[1]}
		insinto ${into}
		doins ${m[0]}
	done < .manifest || die

	dosym nvidia-installer /opt/nvidia/${PV}/bin/nvidia-uninstall
	insopts -m4755
	insinto /opt/nvidia/${PV}/bin/
	doins nvidia-modprobe

	# Add Nvidia library path to /etc/ld.so.conf
	echo "/opt/nvidia/${PV}/lib64" > nvidia.conf
	insinto /etc/ld.so.conf.d
	doins nvidia.conf
}

src_install() {

	if use install-nvidia-usr; then
		install_nvidia_usr_components
	fi

	if use module_sign ; then
		TMP_DIR=$(mktemp -d)
		DRIVERS_DIR="${TMP_DIR}/drivers"
		FIRMWARE_DIR="${TMP_DIR}/firmware/nvidia/${PV}"
		mkdir -p "${DRIVERS_DIR}"
		mkdir -p "${FIRMWARE_DIR}"
		MODULE_OBJECTS=(nvidia.ko nvidia-drm.ko nvidia-modeset.ko nvidia-peermem.ko nvidia-uvm.ko)
		pushd kernel-module-source/kernel-open || die
		for MODULE_OBJECT in "${MODULE_OBJECTS[@]}"; do
			# Sign the nvidia modules.
			sign_module "${MODULE_OBJECT}" || die "signing failed"
			# copy the object file with extended attributes.
			cp  --preserve=all "${MODULE_OBJECT}" "${DRIVERS_DIR}"/
		done
		popd || die

		FIRMWARE_OBJECTS=(gsp_tu10x.bin gsp_ad10x.bin)
		pushd firmware || die
		for FIRMWARE_OBJECT in "${FIRMWARE_OBJECTS[@]}"; do
			# Sign GSP firmware modules.
			sign_firmware "${FIRMWARE_OBJECT}" || die "signing failed"
			# copy the object file with extended attributes.
			cp --preserve=all "${FIRMWARE_OBJECT}" "${FIRMWARE_DIR}"/
		done
		popd || die
		# Nvidia drivers are stored as build artifacts and not installed in image.
		local source_dir=opt/google/drivers
		dodir ${source_dir}
		# tar the files with extended attributes.
		tar --xattrs --xattrs-include=* -czf "${D}/${source_dir}/nvidia-drivers-${PV}.tgz" -C "${TMP_DIR}"/ .  || die
		rm -rf "${TMP_DIR}"

	fi

	# install X509 ima certificate
	insinto /etc/ima
	newins "${KBUILD_OUTPUT}"/certs/signing_key.x509 pubkey.x509
}