From 4b168dd355104465937ad3dde185c4975f1ffcde Mon Sep 17 00:00:00 2001
From: Xuanqi He <xuanqihe@amazon.com>
Date: Mon, 19 May 2025 15:45:12 -0400
Subject: [PATCH] Fix an issue where users cannot SSH into LoginNodes with
 LoginNode-specific keys when different keys are specified for HeadNode and
 LoginNodes.

---
 CHANGELOG.md                                  |  3 ++
 .../recipes/config/cluster_user.rb            | 34 +++++++++++++++++++
 .../attributes/cluster.rb                     |  1 +
 3 files changed, 38 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 58197c287..ffab7016a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,9 @@ This file is used to list changes made in each version of the AWS ParallelCluste
 **CHANGES**
 - Ubuntu 20.04 is no longer supported.
 
+**BUG FIXES**
+- Fix an issue where users cannot SSH into LoginNodes with LoginNode-specific keys when different keys are specified for HeadNode and LoginNodes.
+
 3.13.1
 ------
 
diff --git a/cookbooks/aws-parallelcluster-platform/recipes/config/cluster_user.rb b/cookbooks/aws-parallelcluster-platform/recipes/config/cluster_user.rb
index 7cd35e402..bc322a2a2 100644
--- a/cookbooks/aws-parallelcluster-platform/recipes/config/cluster_user.rb
+++ b/cookbooks/aws-parallelcluster-platform/recipes/config/cluster_user.rb
@@ -89,6 +89,40 @@
     shell '/bin/bash'
   end
 
+  directory node['cluster']['login_authorized_keys_dir'] do
+    owner 'root'
+    group 'root'
+    mode '0755'
+  end
+
+  directory "#{node['cluster']['login_authorized_keys_dir']}/#{node['cluster']['cluster_user']}" do
+    owner node['cluster']['cluster_user']
+    group node['cluster']['cluster_user']
+    mode '0700'
+  end
+
+  bash 'populate_login_node_local_key' do
+    code <<-PERMS
+      set -e
+      cp #{node['cluster']['shared_dir_login_nodes']}/authorized_keys \
+         #{node['cluster']['login_authorized_keys_dir']}/#{node['cluster']['cluster_user']}/authorized_keys
+      chown #{node['cluster']['cluster_user']}:#{node['cluster']['cluster_user']} \
+         #{node['cluster']['login_authorized_keys_dir']}/#{node['cluster']['cluster_user']}/authorized_keys
+      chmod 0600 #{node['cluster']['login_authorized_keys_dir']}/#{node['cluster']['cluster_user']}/authorized_keys
+    PERMS
+    not_if { ::File.exist?("#{node['cluster']['login_authorized_keys_dir']}/#{node['cluster']['cluster_user']}/authorized_keys") }
+  end
+
+  bash 'patch_sshd_config_for_login_nodes' do
+    code <<-CONF
+      set -e
+      AUTH_DIR="#{node['cluster']['login_authorized_keys_dir']}/#{node['cluster']['cluster_user']}"
+      LINE='AuthorizedKeysFile /etc/ssh/login_nodes_authorized_keys.d/%u/authorized_keys .ssh/authorized_keys'
+      grep -q "${AUTH_DIR}/authorized_keys" /etc/ssh/sshd_config || echo "${LINE}" >> /etc/ssh/sshd_config
+    CONF
+  end
+
+  # keep the existing copy into /home for backward compatibility
   bash "copy_auth_file" do
     code <<-PERMS
       set -e
diff --git a/cookbooks/aws-parallelcluster-shared/attributes/cluster.rb b/cookbooks/aws-parallelcluster-shared/attributes/cluster.rb
index def13a134..7feabc8be 100644
--- a/cookbooks/aws-parallelcluster-shared/attributes/cluster.rb
+++ b/cookbooks/aws-parallelcluster-shared/attributes/cluster.rb
@@ -23,6 +23,7 @@
 default['cluster']['previous_cluster_config_path'] = "#{node['cluster']['shared_dir']}/previous-cluster-config.yaml"
 default['cluster']['login_cluster_config_path'] = "#{node['cluster']['shared_dir_login_nodes']}/cluster-config.yaml"
 default['cluster']['login_previous_cluster_config_path'] = "#{node['cluster']['shared_dir_login_nodes']}/previous-cluster-config.yaml"
+default['cluster']['login_authorized_keys_dir'] = '/etc/ssh/login_nodes_authorized_keys.d'
 default['cluster']['change_set_path'] = "#{node['cluster']['shared_dir']}/change-set.json"
 default['cluster']['instance_types_data_path'] = "#{node['cluster']['shared_dir']}/instance-types-data.json"
 default['cluster']['previous_instance_types_data_path'] = "#{node['cluster']['shared_dir']}/previous-instance-types-data.json"