From e02ffc3c2b60eb5d65b94fc9917978c1b1a9a881 Mon Sep 17 00:00:00 2001
From: Verdi March <marcverd@amazon.com>
Date: Fri, 1 Mar 2024 14:19:12 +0800
Subject: [PATCH 1/2] Upstream aws-samples/playground-persistent-cluster PR-8

---
 .../base-config/setup_mariadb_accounting.sh   | 28 +++++++++++++++++--
 1 file changed, 25 insertions(+), 3 deletions(-)

diff --git a/1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/setup_mariadb_accounting.sh b/1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/setup_mariadb_accounting.sh
index 030966fe..95496b97 100644
--- a/1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/setup_mariadb_accounting.sh
+++ b/1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/setup_mariadb_accounting.sh
@@ -1,6 +1,15 @@
 #!/bin/bash
 
-set -ex
+set -euo pipefail
+
+# https://askubuntu.com/a/1472412
+# Exclude special characters that break create_slurm_database().
+EXCLUDED_CHAR="'\"\`\\[]{}()*"
+SLURM_DB_PASSWORD=$(apg -a 1 -M SNCL -m 10 -x 10 -n 1 -E "${EXCLUDED_CHAR}")
+
+# Retain adt behavior. However, verbosity will be disabled at select places to
+# prevent credentials getting leaked to Cloudwatch logs.
+set -x
 
 SLURM_ACCOUNTING_CONFIG_FILE=/opt/slurm/etc/accounting.conf
 SLURMDB_CONFIG_FILE=/opt/slurm/etc/slurmdbd.conf
@@ -15,6 +24,7 @@ fi
 # Use expect to for the interactive shell.
 setup_mariadb() {
   echo "Running mysql_secure_installation"
+  set +x
   SECURE_MYSQL=$(expect -c "
   set timeout 10
   log_file /var/log/provision/secure_mysql.log
@@ -33,11 +43,15 @@ setup_mariadb() {
   send \"y\r\"
   expect eof
   ")
+  set -x
+  chmod 400 /var/log/provision/secure_mysql.log
 }
 
 # Create the default database for SLURM accounting
 create_slurm_database() {
+  set +x
   echo "Creating accounting database"
+  local ESCAPED_SLURM_DB_PASSWORD=$(printf '%q' "$SLURM_DB_PASSWORD")
   SETUP_MYSQL=$(expect -c "
   set timeout 15
   log_file /var/log/provision/setup_mysql.log
@@ -47,7 +61,7 @@ create_slurm_database() {
   send \"\r\"
   sleep 1
   expect \"*]>\"
-  send \"grant all on slurm_acct_db.* TO 'slurm'@'localhost' identified by 'some_pass' with grant option;\r\"
+  send \"grant all on slurm_acct_db.* TO 'slurm'@'localhost' identified by '${ESCAPED_SLURM_DB_PASSWORD}' with grant option;\r\"
   sleep 1
   expect \"*]>\"
   send \"create database slurm_acct_db;\r\"
@@ -56,13 +70,21 @@ create_slurm_database() {
   send \"exit\r\"
   expect eof
   ")
+  set -x
+  chmod 400 /var/log/provision/setup_mysql.log
 }
 
 # Setup the configuration for slurmdbd to use MariaDB.
 create_slurmdbd_config() {
-  SLURM_DB_USER=slurm SLURM_DB_PASSWORD=some_pass envsubst < "$SLURMDB_CONFIG_FILE.template" > $SLURMDB_CONFIG_FILE
+  # Do not push db credentials to Cloudwatch logs
+  echo 'BEGIN: create_slurmdbd_config()'
+  set +x
+  SLURM_DB_USER=slurm SLURM_DB_PASSWORD="$SLURM_DB_PASSWORD" envsubst < "$SLURMDB_CONFIG_FILE.template" > $SLURMDB_CONFIG_FILE
+  set -x
+
   chown slurm:slurm $SLURMDB_CONFIG_FILE
   chmod 600 $SLURMDB_CONFIG_FILE
+  echo 'END: create_slurmdbd_config()'
 }
 
 # Append the accounting settings to accounting.conf, this file is empty by default and included into

From 67e51a06d834b68bc44941c45a003220ad5a794c Mon Sep 17 00:00:00 2001
From: Verdi March <marcverd@amazon.com>
Date: Mon, 4 Mar 2024 12:10:07 +0800
Subject: [PATCH 2/2] Exclude '#'

---
 .../LifecycleScripts/base-config/setup_mariadb_accounting.sh    | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/setup_mariadb_accounting.sh b/1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/setup_mariadb_accounting.sh
index 95496b97..8dcfb622 100644
--- a/1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/setup_mariadb_accounting.sh
+++ b/1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config/setup_mariadb_accounting.sh
@@ -4,7 +4,7 @@ set -euo pipefail
 
 # https://askubuntu.com/a/1472412
 # Exclude special characters that break create_slurm_database().
-EXCLUDED_CHAR="'\"\`\\[]{}()*"
+EXCLUDED_CHAR="'\"\`\\[]{}()*#"
 SLURM_DB_PASSWORD=$(apg -a 1 -M SNCL -m 10 -x 10 -n 1 -E "${EXCLUDED_CHAR}")
 
 # Retain adt behavior. However, verbosity will be disabled at select places to