#!/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
