Skip to content

Commit

Permalink
[pam] rework pam_oar_adpot wip
Browse files Browse the repository at this point in the history
  • Loading branch information
npf committed Jan 10, 2025
1 parent fee6553 commit 6ee6162
Showing 1 changed file with 32 additions and 25 deletions.
57 changes: 32 additions & 25 deletions sources/core/tools/oarsh-legacy/pam_oar_adopt
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -9,47 +9,55 @@
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:
# - the connecting user is oar or root, we fail silently (since we are in 'sufficient' mode)
# - 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
Expand All @@ -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
}
Expand All @@ -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
Expand Down

0 comments on commit 6ee6162

Please sign in to comment.