#!/bin/bash
# $Id$
# This is the shell program used with the oarsh script.
# It adds its own process in the cpuset OAR_CPUSET and launches
# the shell or the script of the user OAR_JOB_USER.

DEFAULT_SHELL=/bin/bash
XAUTH_LOCATION="%%XAUTHCMDPATH%%"
CPUSET_MOUNT_POINT="/dev/oar_cgroups_links"
OARDIR="%%OARDIR%%"
OAR_RUNTIME_DIRECTORY="/tmp/oar_runtime"

export PATH=$OARDIR/oardodo:$PATH

source %%OARCONFDIR%%/oar.conf

# File generated by job resource manager on all job nodes
JOBENVFILE="$OAR_RUNTIME_DIRECTORY/${OAR_CPUSET##*/}.env"

OLDUMASK=$(umask)
umask 0022

# $1 = Name of the cpuset
# $2,$3,... = PIDs to add
# Add PIDS into the cpuset $1
function add_process_to_cpuset() {
    [ "$1" = "undef" ] && return
    ######################
    # CPUSET assignation #
    ###########################################################################
    CPUSETNAME=$1
    if [ ! -w /dev/cpuset/$CPUSETNAME/tasks ]; then
        echo "oarsh: Cannot find cpuset file : /dev/cpuset/$CPUSETNAME/tasks" 1>&2
        exit 61
    fi
    shift

    # Add each processes to the right OAR cgroups
    for p in $@; do
        if [ -r $CPUSET_MOUNT_POINT/cpuset/$CPUSETNAME/tasks ]; then
            # Feed at least cpuset (security)
            $OARDIR/oardodo/oardodo sh -c "echo $p > $CPUSET_MOUNT_POINT/cpuset/$CPUSETNAME/tasks" || return 1
            for c in $CPUSET_MOUNT_POINT/*; do
                if [ "$c" != "$CPUSET_MOUNT_POINT/cpuset" -a -r "$c/$CPUSETNAME/tasks" ]; then
                    $OARDIR/oardodo/oardodo sh -c "echo $p > $c/$CPUSETNAME/tasks" || return 1
                fi
            done
        else
            # Old behaviour without cgroups (just cpuset)
            $OARDIR/oardodo/oardodo sh -c "echo $p > /dev/cpuset/$CPUSETNAME/tasks" || return 1
        fi
    done
    ###########################################################################
}

if [ "$OAR_JOB_USER" = "" ]
then
    if [ "$SSH_CLIENT" != ""  ] && [ "$OAR_KEY" != "1" ]
    then
        echo "oarsh: The OAR_KEY environment variable is not defined and this seems to be a oar user connection." 1>&2
        exit 65
    fi
    # It must be oar
    if [ "$OAR_CPUSET" != "" ]
    then
        add_process_to_cpuset $OAR_CPUSET $$ $PPID || exit 62
        [ -r "$JOBENVFILE" ] && source $JOBENVFILE
    fi
    $OARDIR/oardodo/oardodo renice 0 $$ $PPID > /dev/null 2>&1
    export SHELL=$DEFAULT_SHELL
    
    umask $OLDUMASK
    exec $DEFAULT_SHELL "$@"
    echo "oarsh: exec failed" 1>&2
    exit 66
else
    if [ "$OAR_CPUSET" = "" ]
    then
        echo "oarsh: OAR_CPUSET variable is empty; Is your sshd right configured with 'AcceptEnv OAR_CPUSET OAR_JOB_USER' on all computing nodes?" 1>&2
        exit 63
    fi
    add_process_to_cpuset $OAR_CPUSET $$ $PPID || exit 62
    
    #Manage display
    if [ -n "$DISPLAY" ]
    then
        if [ -x "$XAUTH_LOCATION" ]
        then
            $XAUTH_LOCATION -q extract - ${DISPLAY#localhost} | OARDO_BECOME_USER=${OAR_JOB_USER} $OARDIR/oardodo/oardodo $XAUTH_LOCATION merge -
            [ "${OAR_JOB_USER}" != "$OAR_JOB_USER" ] && OARDO_BECOME_USER=${OAR_JOB_USER} $OARDIR/oardodo/oardodo bash --noprofile --norc -c 'chmod 660 $HOME/.Xauthority'
        fi
    fi
    #Change tty owner
    TTY=$(tty) && test -e $TTY && $OARDIR/oardodo/oardodo chown $OAR_JOB_USER:oar $TTY && $OARDIR/oardodo/oardodo chmod 660 $TTY
    $OARDIR/oardodo/oardodo renice 0 $$ $PPID > /dev/null 2>&1
    [ -r "$JOBENVFILE" ] && source $JOBENVFILE
    if [ "$1" = "" ]
    then
        # Simulate initial login
        export OARDO_BECOME_USER=$OAR_JOB_USER
        umask $OLDUMASK
        exec $OARDIR/oardodo/oardodo
        #exec oardodo su - $OAR_JOB_USER
        echo "oarsh: exec failed" 1>&2
        exit 66
    else
        export OARDO_BECOME_USER=$OAR_JOB_USER
        export OARDO_USE_USER_SHELL=1
        umask $OLDUMASK
        exec $OARDIR/oardodo/oardodo "$@"
        echo "oarsh: exec failed" 1>&2
        exit 66
    fi
fi

echo "oarsh: Really bad error" 1>&2
exit 67

