From 66e80a49ff110fe31c647895341f7bf171aa1a70 Mon Sep 17 00:00:00 2001
From: Matthieu Hautreux <matthieu.hautreux@cea.fr>
Date: Mon, 22 Oct 2012 12:39:34 -0700
Subject: [PATCH] Add task_pre_launch_priv call in task plugin API

This privileged call is executed just after the fork of each
forked task by user root before becoming the user. It enables
the task plugin to perform actions as a privileged user in
every task context.
---
 doc/html/taskplugins.shtml                | 16 +++++++++++++++-
 doc/man/man5/slurm.conf.5                 |  3 +++
 src/plugins/task/affinity/task_affinity.c |  9 +++++++++
 src/plugins/task/cgroup/task_cgroup.c     |  9 +++++++++
 src/plugins/task/none/task_none.c         | 11 +++++++++++
 src/slurmd/common/task_plugin.c           | 22 ++++++++++++++++++++++
 src/slurmd/common/task_plugin.h           |  7 +++++++
 src/slurmd/slurmstepd/mgr.c               |  7 +++++--
 8 files changed, 81 insertions(+), 3 deletions(-)

diff --git a/doc/html/taskplugins.shtml b/doc/html/taskplugins.shtml
index e727c02db10..d34b400635c 100644
--- a/doc/html/taskplugins.shtml
+++ b/doc/html/taskplugins.shtml
@@ -5,7 +5,7 @@
 <h2> Overview</h2>
 <p> This document describes SLURM task management plugins and the API
 that defines them. It is intended as a resource to programmers wishing
-to write their own SLURM scheduler plugins. This is version 1 of the API.</p>
+to write their own SLURM scheduler plugins. This is version 2 of the API.</p>
 
 <p>SLURM task management plugins are SLURM plugins that implement the
 SLURM task management API described herein. They would typically be
@@ -146,6 +146,20 @@ data structure definition.</p>
 On failure, the plugin should return SLURM_ERROR and set the errno to an
 appropriate value to indicate the reason for failure.</p>
 
+<p class="commandline">int task_pre_launch_priv (slurmd_job_t *job);</p>
+<p style="margin-left:.2in"><b>Description</b>: task_pre_launch_priv() is called
+by each forked task just after the fork. Note that no particular task related
+information is available in the job structure at that time.
+Executed by the <b>slurmstepd</b> program as user root.</p>
+<p style="margin-left:.2in"><b>Arguments</b>:
+<span class="commandline">job</span>&nbsp;&nbsp;&nbsp;(input)
+pointer to the job to be initiated.
+See <b>src/slurmd/slurmstepd/slurmstepd_job.h</b> for the
+data structure definition.</p>
+<p style="margin-left:.2in"><b>Returns</b>: SLURM_SUCCESS if successful.
+On failure, the plugin should return SLURM_ERROR and set the errno to an
+appropriate value to indicate the reason for failure.</p>
+
 <p class="commandline">int task_pre_launch (slurmd_job_t *job);</p>
 <p style="margin-left:.2in"><b>Description</b>: task_pre_launch() is called
 prior to exec of application task.
diff --git a/doc/man/man5/slurm.conf.5 b/doc/man/man5/slurm.conf.5
index 5c8f225e989..8f892fb7283 100644
--- a/doc/man/man5/slurm.conf.5
+++ b/doc/man/man5/slurm.conf.5
@@ -2115,6 +2115,9 @@ Will clear environment variables for the task being spawned.
 .TP
 The order of task prolog/epilog execution is as follows:
 .TP
+\fB1. pre_launch_priv()\fR
+Function in TaskPlugin
+.TP
 \fB1. pre_launch()\fR
 Function in TaskPlugin
 .TP
diff --git a/src/plugins/task/affinity/task_affinity.c b/src/plugins/task/affinity/task_affinity.c
index ee90ec14edb..4e980147c86 100644
--- a/src/plugins/task/affinity/task_affinity.c
+++ b/src/plugins/task/affinity/task_affinity.c
@@ -438,6 +438,15 @@ extern int task_pre_launch (slurmd_job_t *job)
 	return rc;
 }
 
+/*
+ * task_pre_launch_priv() is called prior to exec of application task.
+ * in privileged mode, just after slurm_spank_task_init_privileged
+ */
+extern int task_pre_launch_priv (slurmd_job_t *job)
+{
+	return SLURM_SUCCESS;
+}
+
 /*
  * task_term() is called after termination of application task.
  *	It is preceded by --task-epilog (from srun command line)
diff --git a/src/plugins/task/cgroup/task_cgroup.c b/src/plugins/task/cgroup/task_cgroup.c
index b7a4300727c..63a3889f4e6 100644
--- a/src/plugins/task/cgroup/task_cgroup.c
+++ b/src/plugins/task/cgroup/task_cgroup.c
@@ -269,6 +269,15 @@ extern int task_pre_launch (slurmd_job_t *job)
 	return SLURM_SUCCESS;
 }
 
+/*
+ * task_pre_launch_priv() is called prior to exec of application task.
+ * in privileged mode, just after slurm_spank_task_init_privileged
+ */
+extern int task_pre_launch_priv (slurmd_job_t *job)
+{
+	return SLURM_SUCCESS;
+}
+
 /*
  * task_term() is called after termination of application task.
  *	It is preceded by --task-epilog (from srun command line)
diff --git a/src/plugins/task/none/task_none.c b/src/plugins/task/none/task_none.c
index e2f3a035e53..ec5a5042f88 100644
--- a/src/plugins/task/none/task_none.c
+++ b/src/plugins/task/none/task_none.c
@@ -181,6 +181,17 @@ extern int task_pre_launch (slurmd_job_t *job)
 	return SLURM_SUCCESS;
 }
 
+/*
+ * task_pre_launch_priv() is called prior to exec of application task.
+ * in privileged mode, just after slurm_spank_task_init_privileged
+ */
+extern int task_pre_launch_priv (slurmd_job_t *job)
+{
+	debug("task_pre_launch_priv: %u.%u",
+		job->jobid, job->stepid);
+	return SLURM_SUCCESS;
+}
+
 /*
  * task_term() is called after termination of application task.
  *	It is preceded by --task-epilog (from srun command line)
diff --git a/src/slurmd/common/task_plugin.c b/src/slurmd/common/task_plugin.c
index df648d96d3a..e71356f7313 100644
--- a/src/slurmd/common/task_plugin.c
+++ b/src/slurmd/common/task_plugin.c
@@ -63,6 +63,7 @@ typedef struct slurmd_task_ops {
 	int	(*slurmd_release_resources) (uint32_t job_id);
 
 	int	(*pre_setuid)		    (slurmd_job_t *job);
+	int	(*pre_launch_priv)	    (slurmd_job_t *job);
 	int	(*pre_launch)		    (slurmd_job_t *job);
 	int	(*post_term)		    (slurmd_job_t *job);
 	int	(*post_step)		    (slurmd_job_t *job);
@@ -79,6 +80,7 @@ static const char *syms[] = {
 	"task_slurmd_resume_job",
 	"task_slurmd_release_resources",
 	"task_pre_setuid",
+	"task_pre_launch_priv",
 	"task_pre_launch",
 	"task_post_term",
 	"task_post_step",
@@ -335,6 +337,26 @@ extern int pre_setuid(slurmd_job_t *job)
 	return (rc);
 }
 
+/*
+ * Note in privileged mode that a task launch is about to occur.
+ *
+ * RET - slurm error code
+ */
+extern int pre_launch_priv(slurmd_job_t *job)
+{
+	int i, rc = SLURM_SUCCESS;
+
+	if (slurmd_task_init())
+		return SLURM_ERROR;
+
+	slurm_mutex_lock( &g_task_context_lock );
+	for (i = 0; ((i < g_task_context_num) && (rc == SLURM_SUCCESS)); i++)
+		rc = (*(g_task_context[i]->ops.pre_launch_priv))(job);
+	slurm_mutex_unlock( &g_task_context_lock );
+
+	return (rc);
+}
+
 /*
  * Note that a task launch is about to occur.
  *
diff --git a/src/slurmd/common/task_plugin.h b/src/slurmd/common/task_plugin.h
index 4419c8ef9bf..730017735a8 100644
--- a/src/slurmd/common/task_plugin.h
+++ b/src/slurmd/common/task_plugin.h
@@ -116,6 +116,13 @@ extern int slurmd_release_resources(uint32_t job_id);
  */
 extern int pre_setuid(slurmd_job_t *job);
 
+/*
+ * Note in privileged mode that a task launch is about to occur.
+ *
+ * RET - slurm error code
+ */
+extern int pre_launch_priv(slurmd_job_t *job);
+
 /*
  * Note that a task launch is about to occur.
  *
diff --git a/src/slurmd/slurmstepd/mgr.c b/src/slurmd/slurmstepd/mgr.c
index 1fb304733f9..366ddee5a57 100644
--- a/src/slurmd/slurmstepd/mgr.c
+++ b/src/slurmd/slurmstepd/mgr.c
@@ -1112,7 +1112,7 @@ fail1:
 }
 
 static int
-_spank_task_privileged(slurmd_job_t *job, int taskid, struct priv_state *sp)
+_pre_task_privileged(slurmd_job_t *job, int taskid, struct priv_state *sp)
 {
 	if (_reclaim_privileges(sp) < 0)
 		return SLURM_ERROR;
@@ -1120,6 +1120,9 @@ _spank_task_privileged(slurmd_job_t *job, int taskid, struct priv_state *sp)
 	if (spank_task_privileged (job, taskid) < 0)
 		return error("spank_task_init_privileged failed");
 
+	if (pre_launch_priv(job) < 0)
+		return error("pre_launch_priv failed");
+
 	return(_drop_privileges (job, true, sp));
 }
 
@@ -1422,7 +1425,7 @@ _fork_all_tasks(slurmd_job_t *job, bool *io_initialized)
 			 *  Reclaim privileges and call any plugin hooks
 			 *   that may require elevated privs
 			 */
-			if (_spank_task_privileged(job, i, &sprivs) < 0)
+			if (_pre_task_privileged(job, i, &sprivs) < 0)
 				exit(1);
 
  			if (_become_user(job, &sprivs) < 0) {
-- 
GitLab