#!/bin/bash

module_init() {
	local i url sigfile file hash
	declare -a apts
	[ "$SYS_DISTRIBUTION" = "ubuntu" ] || [ "$SYS_DISTRIBUTION" = "debian" ] || return 0
	# Forcefully add docker repo
	local vers="$SYS_VERSION"
	# XXX HACK - currently 12 is still testing, so we get "n/a"
	# and then, there is no release for it yet at nvidia. So use debian 11 repo.
	local codename="$(lsb_release -cs)"
	if [ "$vers" = "n/a" ]; then
		vers=11
		codename="buster"
	fi
	if [ "$SYS_DISTRIBUTION" = "debian" ]; then
		# Brute force down to a valid version
		while (( vers > 10 )) && ! curl -sSfL "https://nvidia.github.io/nvidia-docker/${SYS_DISTRIBUTION}${vers}/nvidia-docker.list" > /dev/null; do
			(( vers-- ))
		done
	fi
	apts=(
		"https://download.docker.com/linux/${SYS_DISTRIBUTION}/gpg deb [arch=$(dpkg --print-architecture)] https://download.docker.com/linux/$SYS_DISTRIBUTION $codename stable"
		"https://nvidia.github.io/nvidia-docker/gpgkey https://nvidia.github.io/nvidia-docker/${SYS_DISTRIBUTION}${vers}/nvidia-docker.list"
	)
	for i in "${apts[@]}"; do
		# First part is GPG signing key URL
		url="${i%% *}"
		hash="$( echo "$i" | md5sum | cut -c1-10 )"
		sigfile="/etc/apt/trusted.gpg.d/docker-${hash}"
		[ -s "$sigfile" ] && continue
		curl -fsSL "$url" > "$sigfile" \
			|| perror "Could not download docker gpg key from $url"
		if grep -qF -- '---BEGIN' "$sigfile"; then
			mv "$sigfile" "${sigfile}.asc"
			sigfile="${sigfile}.asc"
		else
			mv "$sigfile" "${sigfile}.gpg"
			sigfile="${sigfile}.gpg"
		fi
		# Cut away first part (URL to GPG)
		url="${i#* }"
		file="/etc/apt/sources.list.d/docker-${hash}.list"
		[ -s "$file" ] && continue
		if [[ "${url}" == http* ]]; then
			# Start with http, assume this is the URL for a sources.list file
			download "$url" "$file"
		else
			# Otherwise, assume it's a line for a sources.list
			echo "$url" > "$file"
		fi
	done
	apt-get update
}

build() {
	local service
	for service in docker containerd; do
		systemctl disable "${service}.service" || perror "Could not disable $service"
	done
	systemctl enable "docker.socket" || perror "Could not enable docker.socket activation"
	# Plugin binary
	download_untar "https://github.com/ad-freiburg/docker-no-trivial-root/releases/download/v0.1.0/docker-no-trivial-root_x86_64.tar.bz2" \
		"$MODULE_WORK_DIR/src"
	mkdir -p "$MODULE_BUILD_DIR/usr/sbin"
	mv "$MODULE_WORK_DIR/src/docker-no-trivial-root_x86_64/docker-no-trivial-root" \
		"$MODULE_BUILD_DIR/usr/sbin/docker-no-trivial-root" \
		|| perror "Cannot move docker-no-trivial-root"
	chmod +x "$MODULE_BUILD_DIR/usr/sbin/docker-no-trivial-root"
	chown 0:0 "$MODULE_BUILD_DIR/usr/sbin/docker-no-trivial-root"
	# Patch systemd service
	mkdir -p "$MODULE_BUILD_DIR/etc/systemd/system"
	sed -r 's/^(ExecStart=.*dockerd) (.*)$/\1 --authorization-plugin=no-trivial-root \2/' \
		"/lib/systemd/system/docker.service" > "$MODULE_BUILD_DIR/etc/systemd/system/docker.service" \
		|| perror "Could not patch docker.service"
	# That weird range stuff
	local item
	for item in subuid subgid; do
		awk -F: 'BEGIN {
				max=0
				found=0
			} {
				if ($1=="dockremap")
					found=1
				if ($2>max)
					max=($2)
				print $0
			} END {
				if (!found)
					print "dockremap:"max+65536":65536"
			}' "/etc/${item}" > "${MODULE_BUILD_DIR}/etc/${item}" \
				|| perror "Could not patch /etc/$item"
	done
	# Workaround for https://github.com/NVIDIA/nvidia-docker/issues/1399
	mkdir -p "${MODULE_BUILD_DIR}/etc/nvidia-container-runtime"
	sed -r 's#^ldconfig\s*=.*$#ldconfig = "/sbin/ldconfig"#' "/etc/nvidia-container-runtime/config.toml" \
		> "${MODULE_BUILD_DIR}/etc/nvidia-container-runtime/config.toml" \
		|| perror "Could not patch nvidia-container config.toml"
}