From 1957539b5c0a7ac25ce6dae50038f724ec65603b Mon Sep 17 00:00:00 2001
From: "Mark A. Grondona" <mark.grondona@gmail.com>
Date: Fri, 13 Dec 2024 08:10:42 -0800
Subject: [PATCH] shell: support expansion of select env vars as mustache
 templates

Problem: It would be useful to allow expansion of mustache templates
in envionment variables in the job shell. However, expanding templates
in the entire jobspec environment dictionary would be inefficient,
since values with templates would be by far the uncommon case. Also,
environment variables may legitimately contain values that look like
mustache templates but should not be expanded.

Add a new shell builtin plugin for expanding select environment
variables in a `env-expand` shell options dictionary. This allows
the jobspec to opt-in to mustache expansion for a select set of
environment variables, solving the issues noted above.
---
 src/shell/Makefile.am  |  3 +-
 src/shell/builtins.c   |  2 ++
 src/shell/env-expand.c | 62 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 src/shell/env-expand.c

diff --git a/src/shell/Makefile.am b/src/shell/Makefile.am
index 57c51d26cbd6..1382d6244ff1 100644
--- a/src/shell/Makefile.am
+++ b/src/shell/Makefile.am
@@ -99,7 +99,8 @@ flux_shell_SOURCES = \
 	signal.c \
 	files.c \
 	hwloc.c \
-	rexec.c
+	rexec.c \
+	env-expand.c
 
 if HAVE_INOTIFY
 flux_shell_SOURCES += oom.c
diff --git a/src/shell/builtins.c b/src/shell/builtins.c
index 38baa5ecb969..4af5a20f18f7 100644
--- a/src/shell/builtins.c
+++ b/src/shell/builtins.c
@@ -56,6 +56,7 @@ extern struct shell_builtin builtin_signal;
 extern struct shell_builtin builtin_oom;
 extern struct shell_builtin builtin_hwloc;
 extern struct shell_builtin builtin_rexec;
+extern struct shell_builtin builtin_env_expand;
 
 static struct shell_builtin * builtins [] = {
     &builtin_tmpdir,
@@ -86,6 +87,7 @@ static struct shell_builtin * builtins [] = {
 #endif
     &builtin_hwloc,
     &builtin_rexec,
+    &builtin_env_expand,
     &builtin_list_end,
 };
 
diff --git a/src/shell/env-expand.c b/src/shell/env-expand.c
new file mode 100644
index 000000000000..b0b2199118e6
--- /dev/null
+++ b/src/shell/env-expand.c
@@ -0,0 +1,62 @@
+/************************************************************\
+ * Copyright 2024 Lawrence Livermore National Security, LLC
+ * (c.f. AUTHORS, NOTICE.LLNS, COPYING)
+ *
+ * This file is part of the Flux resource manager framework.
+ * For details, see https://github.com/flux-framework.
+ *
+ * SPDX-License-Identifier: LGPL-3.0
+\************************************************************/
+
+#define FLUX_SHELL_PLUGIN_NAME "env-expand"
+
+#if HAVE_CONFIG_H
+#include "config.h"
+#endif
+
+#include <jansson.h>
+
+#include <flux/shell.h>
+
+#include "builtins.h"
+
+static int env_expand (flux_plugin_t *p,
+                       const char *topic,
+                       flux_plugin_arg_t *args,
+                       void *data)
+{
+    json_t *to_expand = NULL;
+    json_t *value;
+    const char *key;
+    flux_shell_t *shell = flux_plugin_get_shell (p);
+
+    if (!(shell = flux_plugin_get_shell (p)))
+        return shell_log_errno ("unable to get shell handle");
+    if (flux_shell_getopt_unpack (shell, "env-expand", "o", &to_expand) != 1)
+        return 0;
+    json_object_foreach (to_expand, key, value) {
+        const char *s = json_string_value (value);
+        char *result;
+        if (s == NULL) {
+            shell_log_error ("invalid value for env var %s", key);
+            continue;
+        }
+        if (!(result = flux_shell_mustache_render (shell, s))) {
+            shell_log_errno ("failed to expand env var %s=%s", key, s);
+            continue;
+        }
+        if (flux_shell_setenvf (shell, 1, key, "%s", result) < 0)
+            shell_log_errno ("failed to set %s=%s", key, result);
+        free (result);
+    }
+    return 0;
+}
+
+struct shell_builtin builtin_env_expand = {
+    .name = FLUX_SHELL_PLUGIN_NAME,
+    .init = env_expand,
+};
+
+/*
+ * vi:tabstop=4 shiftwidth=4 expandtab
+ */