From 2745c246d24df9090e98e8af690c4e702893b8c1 Mon Sep 17 00:00:00 2001
From: jmfuchs <jessef45@gmail.com>
Date: Thu, 9 Jun 2022 17:15:27 -0400
Subject: [PATCH] Feature addition: Enable setting the STS Endpoint from the
 SAML Assertion

---
 .../amazon/redshift/plugin/SamlCredentialsProvider.java   | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/src/main/java/com/amazon/redshift/plugin/SamlCredentialsProvider.java b/src/main/java/com/amazon/redshift/plugin/SamlCredentialsProvider.java
index 3e1baf4..affbdc4 100755
--- a/src/main/java/com/amazon/redshift/plugin/SamlCredentialsProvider.java
+++ b/src/main/java/com/amazon/redshift/plugin/SamlCredentialsProvider.java
@@ -357,6 +357,14 @@ public void refresh()
                 principal = entry.getValue();
             }
 
+            // Set STS Regional Endpoint from SAML Assertion
+            List<String> attributeValues = GetSAMLAttributeValues(xPath, doc,
+                    "https://redshift.amazon.com/SAML/Attributes/StsEndpointUrl");
+            if (!attributeValues.isEmpty() && m_stsEndpoint == null)
+            {
+                m_stsEndpoint = attributeValues.get(0);
+            }
+
             AssumeRoleWithSAMLRequest samlRequest = new AssumeRoleWithSAMLRequest();
             samlRequest.setSAMLAssertion(samlAssertion);
             samlRequest.setRoleArn(roleArn);