From 6ee61628af80de4e02b6ebd2541b5b85c7f36a77 Mon Sep 17 00:00:00 2001
From: Pierre Neyron <pierre.neyron@imag.fr>
Date: Wed, 8 Jan 2025 17:44:13 +0100
Subject: [PATCH] [pam] rework pam_oar_adpot wip

---
 sources/core/tools/oarsh-legacy/pam_oar_adopt | 57 +++++++++++--------
 1 file changed, 32 insertions(+), 25 deletions(-)

diff --git a/sources/core/tools/oarsh-legacy/pam_oar_adopt b/sources/core/tools/oarsh-legacy/pam_oar_adopt
index b0e89425..9ef59435 100755
--- a/sources/core/tools/oarsh-legacy/pam_oar_adopt
+++ b/sources/core/tools/oarsh-legacy/pam_oar_adopt
@@ -1,4 +1,4 @@
-#!/bin/sh
+#!/bin/bash
 #
 # pam_oar_adopt is a PAM module that adopts processes launched under ssh
 # connections made by users. The processes will be moved inside the correct
@@ -9,31 +9,34 @@
 set -eu
 
 CGROUP_MOUNT_POINT="/dev/oar_cgroups_links"
-OAR_CPUSETS_BASE="${CGROUP_MOUNT_POINT}/cpuset/oar"
+OAR_CPUSETS_BASE="$CGROUP_MOUNT_POINT/cpuset/oar"
+USER_UID_MIN=1000
 
-get_user_cgroups() {
-	ls -d ${OAR_CPUSETS_BASE}/${1}_* 2>/dev/null | awk -F / '{ ORS=" "; print $NF }'
+
+get_oar_cpusets_of_user() {
+    readarray OAR_CPUSETS < <( cd "$OAR_CPUSETS_BASE" && ls -d "$1"_* 2>/dev/null )
+    OAR_CPUSET=${OAR_CPUSETS[0]}
 }
 
 pam_account() {
-	if [ -z "${PAM_USER+x}" ]; then
+	if [ -z "$PAM_USER" ]; then
 		echo "Please launch this module via PAM"
 		exit 1
 	fi
 
 	# We exit if the pam service is su, we don't want to have the error
 	# message when using su.
-	if [ "${PAM_SERVICE}" = "su-l" ]; then
+	if [ "$PAM_SERVICE" = "su-l" ]; then
 		exit 0
 	fi
 
 	# Exit if the user id is inferior than 1000 (system user), indeed there is
 	# no need to do OAR cgroups machinery in that case.
-	if [ $(getent passwd "${PAM_USER}" | awk -F: '{ print $3 }') -lt 1000 ]; then
+	if [ "$(getent passwd "$PAM_USER" | cut -d: -f3)" -lt "$USER_UID_MIN" ]; then
 		exit 0
 	fi
 
-	get_vars $PAM_USER
+	get_vars "$PAM_USER"
 	test_pam_activation
 
 	# Four cases:
@@ -41,15 +44,20 @@ pam_account() {
 	# - the user has no cgroups (= no jobs) on node
 	# - the user has more than one cgroup or one but without all cores
 	# - the user has one cgroup with all cores
-	if [ ${PAM_USER} = "oar" ] || [ ${PAM_USER} = "root" ] || [ ${PAM_USER} = "vagrant" ]; then
+	if [ "$PAM_USER" = "oar" ] || [ "$PAM_USER" = "root" ]; then
+		exit 1
+	elif [ -z "$OAR_CPUSET" ]; then
+		echo "No running job for user $PAM_USER on this node." >&2
 		exit 1
-	elif [ -z "${USER_CGROUPS+x}" ]; then
-		echo "No running job for user ${PAM_USER} on this node." >&2
+	elif [ "${#OAR_CPUSETS}" -ne 1 ]; then
+		cat << EOF >&2
+Cannot connect to node using 'ssh' because you appear to have more than one job on the node.
+Make sure to only have one job on the node, or use 'oarsh' to connect to a specific job.
+EOF
 		exit 1
-	elif [ $(echo "${USER_CGROUPS}" | awk '{ print NF}') -ne 1 ] ||
-	     [ $(cat ${OAR_CPUSETS_BASE}/$(echo -n ${USER_CGROUPS})/cpuset.cpus) != $ALL_CPUSETS ]; then
+    elif [ "$(< "$OAR_CPUSETS_BASE/$OAR_CPUSET"/cpuset.cpus)" != "$ALL_CPUSET_CPUS" ]; then
 		cat << EOF >&2
-Cannot connect to node using 'ssh' because not all its CPU cores are assigned to the job which reserves it.
+Cannot connect to node using 'ssh' because not all its compute resources (e.g. CPU cores or threads) are assigned to the job which reserves it.
 Reserve the whole node, or use 'oarsh' instead.
 EOF
 		exit 1
@@ -59,34 +67,33 @@ EOF
 }
 
 pam_session() {
-	if [ -z "${PAM_TYPE+x}" ]; then
+	if [ -z "$PAM_TYPE" ]; then
 		echo "Please launch this module via PAM"
 		exit 1
 	fi
 
 	# Exit if not a login
-	if [ "${PAM_TYPE}" != "open_session" ]; then
+	if [ "$PAM_TYPE" != "open_session" ]; then
 		exit 0
 	fi
 
-	G5K_USER=${PAM_RUSER:-$PAM_USER}
-	get_vars $G5K_USER
+	get_vars "${PAM_RUSER:-$PAM_USER}"
 
 	# We could not find a running OAR job for this user on this node. It probably means that
 	# the user connecting is either root or oar (for example because of oarsh).
 	# We do nothing in that case.
-	if [ -z "${USER_CGROUPS}" ]; then
+	if [ -z "$OAR_CPUSET" ]; then
 		exit 0
 	fi
 
 	# To have job's environment variables, we create a symkink to the already
 	# created (by oarsh) environment file. pam_env while then load it.
-	ln -fs /var/lib/oar/$(echo -n ${USER_CGROUPS}).env /var/lib/oar/pam.env
+	ln -fs "/var/lib/oar/$OAR_CPUSET.env" /var/lib/oar/pam.env
 
 	PIDS="$(ps -o ppid= $$)"
 	for pid in $PIDS; do
-		for cgroup in $CGROUP_MOUNT_POINT/*; do
-			echo $pid > "${cgroup}/oar/$(echo -n ${USER_CGROUPS})/tasks"
+		for cg in "$CGROUP_MOUNT_POINT"/*; do
+			echo "$pid" > "$cg/oar/$OAR_CPUSET/tasks"
 		done
 	done
 }
@@ -101,11 +108,11 @@ test_pam_activation() {
 }
 
 get_vars() {
-	USER_CGROUPS=$(get_user_cgroups $1)
-	ALL_CPUSETS=$(cat ${OAR_CPUSETS_BASE}/cpuset.cpus 2> /dev/null || true)
+	get_oar_cpusets_of_user "$1"
+	ALL_CPUSET_CPUS=$(< ${OAR_CPUSETS_BASE}/cpuset.cpus)
 }
 
-[ $# -eq 0 ] && echo "Please provide mode" && exit 1
+[ $# -eq 0 ] && echo "Please provide PAM mode" && exit 1
 
 while getopts ":as" opt; do
 	case $opt in