#!/bin/bash
# Helper for setting up podman for AMD GPU use in a rootless container
#
# Author: Christian Kastner <ckk@kvr.at>
# License: MIT
set -eu

function usage() {
	cat >&2 <<-EOF

	Verifies that a given user can use an AMD GPU in a rootless podman container.

	If USER isn't specified, then the invoking user will be checked.

	Synopsis:
	  $0 -h

	  $0 [-u USER]

	Options:
	  -h     Show this help

	Examples:

	  \$ $0

	  \$ $0 -u someuser

	EOF
	exit 0
}

while getopts "hu:" OPTNAME
do
	case $OPTNAME in
	h)  usage;;
	u)  userNAME="$OPTARG";;
	?)  usage;;
	esac
done
shift $((OPTIND - 1))

userNAME="${userNAME:-$(whoami)}"
renderGID="$(getent group render | cut -d: -f3)"
# By policy
videoGID=44

echo Checks
echo ======

packages_missing=
for pkgname in \
	podman \
	autopkgtest \
	buildah \
	catatonit \
	uidmap \
	netavark \
	aardvark-dns \
	slirp4netns
do
	dpkg -l $pkgname 2> /dev/null | grep -qE '^ii' || packages_missing+=" $pkgname"
done

if [ -n "$packages_missing" ]
then
	echo "[TODO] Key packages not installed:$packages_missing"
else
	echo "  [OK] Key packages are installed"
fi

has_cache=N
for pkgname in approx apt-cacher apt-cacher-ng
do
	dpkg -l $pkgname 2> /dev/null | grep -qE '^ii' && has_cache=Y
done

if [ "$has_cache" = N ]
then
	cat <<-EOF
	 [OPT] No local APT cache detected. While not strictly necessary, it is
	       strongly suggested that you install one of the approx, apt-cacher,
	       or 'apt-cacher-ng' packages.
	EOF
else
	echo "  [OK] Local APT cache detected, make sure to use it"
fi

if ! [ -c /dev/kfd ]
then
	echo "[TODO] /dev/kfd is not present. Has the amdgpu module been loaded?"
else
	echo "  [OK] /dev/kfd is present"
fi

if [ -z "$renderGID" ]
then
	cat <<-EOF
	[TODO] Group 'render' does not exist on this system. Are you sure that you
	are on the right system? This group should have been autmatically created
	by the udev package."
	EOF
else
	echo "  [OK] Group 'render' is present"
fi

if ! groups "$userNAME" | grep -q '\brender\b'
then
	echo "[TODO] User '$userNAME' is not in group 'render'."
	echo "       You can fix this with: sudo gpasswd -a $userNAME render"
else
	echo "  [OK] User '$userNAME' is in group 'render'"
fi

if ! groups "$userNAME" | grep -q '\bvideo\b'
then
	echo "[TODO] User '$userNAME' is not in group 'video'."
	echo "       You can fix this with: sudo gpasswd -a $userNAME video"
else
	echo "  [OK] User '$userNAME' is in group 'video'"
fi

if [ "$(cat /proc/sys/kernel/unprivileged_userns_clone)" != "1" ]
then
	echo "[TODO] unprivileged_userns_clone is not enabled."
	echo "       You can fix this with: sudo echo 1 > /proc/sys/kernel/unprivileged_userns_clone"
else
	echo "  [OK] unprivileged_userns_clone is enabled"
fi

# Assuming user=foo-user, renderGID=123, videoGID=44, we expect an /etc/subgid
# with these entries:
#
#     foo-user:44:1
#     foo-user:123:1
#     foo-user:nnnnnnnn:6553m
#
# nnnnnnnn:6553m is just a large range of subordinate GIDs that should have
# been allocated automatically when the user was created. The grep pattern is
# just a heuristic.

if ! grep -q "$userNAME:$renderGID:1" /etc/subgid
then
	echo "[TODO] /etc/subgid is missing a subordinate GID mapping for user '$userNAME' group 'render'."
	echo "       You can fix this by adding the folowing line to /etc/subgid:"
	echo "           $userNAME:$renderGID:1"
else
	echo "  [OK] /etc/subgid contains a subordinate GID mapping for user '$userNAME' group 'render'"
fi

if ! grep -q "$userNAME:$videoGID:1" /etc/subgid
then
	echo "[TODO] /etc/subgid is missing a subordinate GID mapping for user '$userNAME' group 'video'."
	echo "       You can fix this by adding the folowing line to /etc/subgid:"
	echo "           $userNAME:$videoGID:1"
else
	echo "  [OK] /etc/subgid contains a subordinate GID mapping for user '$userNAME' group 'video'"
fi

if ! grep -q -E "$userNAME:[0-9]{6,}:6553[4-6]" /etc/subgid
then
	maxID=$(cut -d: -f2 /etc/subgid | sort -n | tail -n 1)
	maxRange=$(grep :"$maxID": /etc/subgid | cut -d: -f3)
	newID=$(("$maxID" + "$maxRange"))
	echo "[TODO] /etc/subgid is missing a large subordinate GID range."
	echo "       You can fix this by adding the following line to /etc/subgid:"
	echo "           $userNAME:$newID:65536"
else
	echo "  [OK] /etc/subgid contains a large subordinate GID range"
fi

if ! grep -q -E "$userNAME:[0-9]{6,}:6553[4-6]" /etc/subuid
then
	maxID=$(cut -d: -f2 /etc/subuid | sort -n | tail -n 1)
	maxRange=$(grep :"$maxID": /etc/subuid | cut -d: -f3)
	newID=$(("$maxID" + "$maxRange"))
	echo "[TODO] /etc/subuid is missing a large subordinate UID range."
	echo "       You can fix this by adding the following line to /etc/subuid:"
	echo "           $userNAME:$newID:65536"
else
	echo "  [OK] /etc/subuid contains a large subordinate UID range"
fi
