From 846ece14ed286a91a84e6eb42d32109d79ea6167 Mon Sep 17 00:00:00 2001
From: Xuanqi He <93849823+hehe7318@users.noreply.github.com>
Date: Wed, 13 Sep 2023 13:16:40 -0400
Subject: [PATCH] [develop] Add custom munge key config logic (#2451)

* Munge key config

* Address comments

* Modify the config of kitchen tests

* Add a null value for custom munge key secret arn for now.

* Test of kitchen tests
---
 .../kitchen.slurm-config.yml                  | 15 ++++
 .../libraries/helpers.rb                      | 18 ++---
 .../recipes/config/config_head_node.rb        |  2 +-
 .../recipes/config/config_munge_key.rb        | 18 +++++
 .../resources/munge_key_manager.rb            | 72 +++++++++++++++++++
 .../test/controls/munge_spec.rb               | 28 ++++++++
 kitchen.validate-config.yml                   |  4 ++
 7 files changed, 143 insertions(+), 14 deletions(-)
 create mode 100644 cookbooks/aws-parallelcluster-slurm/recipes/config/config_munge_key.rb
 create mode 100644 cookbooks/aws-parallelcluster-slurm/resources/munge_key_manager.rb

diff --git a/cookbooks/aws-parallelcluster-slurm/kitchen.slurm-config.yml b/cookbooks/aws-parallelcluster-slurm/kitchen.slurm-config.yml
index dabb395309..52fbe82e33 100644
--- a/cookbooks/aws-parallelcluster-slurm/kitchen.slurm-config.yml
+++ b/cookbooks/aws-parallelcluster-slurm/kitchen.slurm-config.yml
@@ -176,3 +176,18 @@ suites:
       cluster:
         node_type: ComputeFleet
         head_node_private_ip: '127.0.0.1'
+  - name: config_head_node_munge
+    run_list:
+    - recipe[aws-parallelcluster-tests::setup]
+    - recipe[aws-parallelcluster-slurm::config_munge_key]
+    verifier:
+      controls:
+        - /tag:config_munge/
+    attributes:
+      cluster:
+        node_type: HeadNode
+        scheduler: 'slurm'
+        config:
+          DevSettings:
+            SlurmSettings:
+              MungeKeySecretArn: null
diff --git a/cookbooks/aws-parallelcluster-slurm/libraries/helpers.rb b/cookbooks/aws-parallelcluster-slurm/libraries/helpers.rb
index 9dcbb9fa6e..24a2c01b40 100644
--- a/cookbooks/aws-parallelcluster-slurm/libraries/helpers.rb
+++ b/cookbooks/aws-parallelcluster-slurm/libraries/helpers.rb
@@ -66,19 +66,11 @@ def enable_munge_service
 end
 
 def setup_munge_head_node
-  # Generate munge key
-  bash 'generate_munge_key' do
-    not_if { ::File.exist?('/etc/munge/munge.key') }
-    user node['cluster']['munge']['user']
-    group node['cluster']['munge']['group']
-    cwd '/tmp'
-    code <<-HEAD_CREATE_MUNGE_KEY
-      set -e
-      # Generates munge key in /etc/munge/munge.key
-      /usr/sbin/mungekey --verbose
-      # Enforce correct permission on the key
-      chmod 0600 /etc/munge/munge.key
-    HEAD_CREATE_MUNGE_KEY
+  # Generate munge key or get it's value from secrets manager
+  munge_key_manager 'manage_munge_key' do
+    munge_key_secret_arn lazy {
+      node['cluster']['config'].dig(:DevSettings, :SlurmSettings, :MungeKeySecretArn)
+    }
   end
 
   enable_munge_service
diff --git a/cookbooks/aws-parallelcluster-slurm/recipes/config/config_head_node.rb b/cookbooks/aws-parallelcluster-slurm/recipes/config/config_head_node.rb
index 8d4ce68f0a..ecd6b08b3c 100644
--- a/cookbooks/aws-parallelcluster-slurm/recipes/config/config_head_node.rb
+++ b/cookbooks/aws-parallelcluster-slurm/recipes/config/config_head_node.rb
@@ -15,7 +15,7 @@
 # OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions and
 # limitations under the License.
 
-setup_munge_head_node unless redhat_on_docker?
+include_recipe 'aws-parallelcluster-slurm::config_munge_key'
 
 # Export /opt/slurm
 nfs_export "#{node['cluster']['slurm']['install_dir']}" do
diff --git a/cookbooks/aws-parallelcluster-slurm/recipes/config/config_munge_key.rb b/cookbooks/aws-parallelcluster-slurm/recipes/config/config_munge_key.rb
new file mode 100644
index 0000000000..bfade907e8
--- /dev/null
+++ b/cookbooks/aws-parallelcluster-slurm/recipes/config/config_munge_key.rb
@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+#
+# Cookbook:: aws-parallelcluster-slurm
+# Recipe:: config_munge_key
+#
+# Copyright:: 2013-2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance with the
+# License. A copy of the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "LICENSE.txt" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions and
+# limitations under the License.
+
+setup_munge_head_node unless redhat_on_docker?
diff --git a/cookbooks/aws-parallelcluster-slurm/resources/munge_key_manager.rb b/cookbooks/aws-parallelcluster-slurm/resources/munge_key_manager.rb
new file mode 100644
index 0000000000..2a696ba182
--- /dev/null
+++ b/cookbooks/aws-parallelcluster-slurm/resources/munge_key_manager.rb
@@ -0,0 +1,72 @@
+# frozen_string_literal: true
+
+#
+# Cookbook:: aws-parallelcluster-slurm
+# Recipe:: config_head_node
+#
+# Copyright:: 2013-2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance with the
+# License. A copy of the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "LICENSE.txt" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions and
+# limitations under the License.
+
+resource_name :munge_key_manager
+provides :munge_key_manager
+unified_mode true
+
+property :munge_key_secret_arn, String
+
+default_action :setup_munge_key
+
+action :setup_munge_key do
+  if new_resource.munge_key_secret_arn
+    # This block will fetch the munge key from Secrets Manager
+    bash 'fetch_and_decode_munge_key' do
+      user 'root'
+      group 'root'
+      cwd '/tmp'
+      code <<-FETCH_AND_DECODE
+        set -e
+        # Get encoded munge key from secrets manager
+        encoded_key=$(aws secretsmanager get-secret-value --secret-id #{new_resource.munge_key_secret_arn} --query 'SecretString' --output text --region #{node['cluster']['region']})
+        # If encoded_key doesn't have a value, error and exit
+        if [ -z "$encoded_key" ]; then
+          echo "Error fetching munge key from Secrets Manager or the key is empty"
+          exit 1
+        fi
+
+        # Decode munge key and write to /etc/munge/munge.key
+        decoded_key=$(echo $encoded_key | base64 -d)
+        if [ $? -ne 0 ]; then
+          echo "Error decoding the munge key with base64"
+          exit 1
+        fi
+
+        echo "$decoded_key" > /etc/munge/munge.key
+
+        # Set ownership on the key
+        chown #{node['cluster']['munge']['user']}:#{node['cluster']['munge']['group']} /etc/munge/munge.key
+        # Enforce correct permission on the key
+        chmod 0600 /etc/munge/munge.key
+      FETCH_AND_DECODE
+    end
+  else
+    # This block will generate a munge key if it doesn't exist
+    bash 'generate_munge_key' do
+      not_if { ::File.exist?('/etc/munge/munge.key') }
+      user node['cluster']['munge']['user']
+      group node['cluster']['munge']['group']
+      cwd '/tmp'
+      code <<-GENERATE_KEY
+        set -e
+        /usr/sbin/mungekey --verbose
+        chmod 0600 /etc/munge/munge.key
+      GENERATE_KEY
+    end
+  end
+end
diff --git a/cookbooks/aws-parallelcluster-slurm/test/controls/munge_spec.rb b/cookbooks/aws-parallelcluster-slurm/test/controls/munge_spec.rb
index c414b630f5..9fb25d8bc4 100644
--- a/cookbooks/aws-parallelcluster-slurm/test/controls/munge_spec.rb
+++ b/cookbooks/aws-parallelcluster-slurm/test/controls/munge_spec.rb
@@ -78,3 +78,31 @@
     it { should be_running }
   end
 end
+
+control 'tag:config_munge_check_munge_key_exists' do
+  title 'Check if the munge key exists'
+
+  describe file('/etc/munge/munge.key') do
+    it { should exist }
+    its('mode') { should cmp '0600' }
+    its('owner') { should eq node['cluster']['munge']['user'] }
+    its('group') { should eq node['cluster']['munge']['group'] }
+  end
+end unless os_properties.redhat_on_docker?
+
+control 'tag:config_munge_check_munge_key_content' do
+  title 'Check if the munge key content is not empty'
+
+  describe file('/etc/munge/munge.key') do
+    its('content') { should_not be_empty }
+  end
+end unless os_properties.redhat_on_docker?
+
+control 'tag:config_munge_check_munge_key_error_messages' do
+  title 'Check for error messages related to munge key'
+
+  describe file('/var/log/chef-client.log') do
+    its('content') { should_not match /Error fetching munge key/ }
+    its('content') { should_not match /Error decoding the munge key/ }
+  end
+end unless os_properties.redhat_on_docker?
diff --git a/kitchen.validate-config.yml b/kitchen.validate-config.yml
index db1058b951..98972d0e9f 100644
--- a/kitchen.validate-config.yml
+++ b/kitchen.validate-config.yml
@@ -33,6 +33,10 @@ _head_node_cluster_attributes: &_head_node_cluster_attributes
   node_type: 'HeadNode'
   ddb_table: <%= ENV['DDB_TABLE'] %>
   slurm_ddb_table: <%= ENV['DDB_TABLE'] %>
+  config:
+    DevSettings:
+      SlurmSettings:
+        MungeKeySecretArn: null
 
 _compute_node_cluster_attributes: &_compute_node_cluster_attributes
   << : *_common_cluster_attributes