From a72da6580a32d3033e8ae7d34857712c2e6ef90b Mon Sep 17 00:00:00 2001
From: SergeyRyabinin <serzhick@gmail.com>
Date: Mon, 16 Dec 2024 16:53:39 +0000
Subject: [PATCH] SigV4A auth selection update

---
 .../domainmodels/codegeneration/Metadata.java |  4 ++
 .../codegeneration/Operation.java             |  3 ++
 .../generators/cpp/CppClientGenerator.java    | 40 +++++++++++++++++--
 3 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/tools/code-generation/generator/src/main/java/com/amazonaws/util/awsclientgenerator/domainmodels/codegeneration/Metadata.java b/tools/code-generation/generator/src/main/java/com/amazonaws/util/awsclientgenerator/domainmodels/codegeneration/Metadata.java
index 9193c5093e8..646f8d53ec4 100644
--- a/tools/code-generation/generator/src/main/java/com/amazonaws/util/awsclientgenerator/domainmodels/codegeneration/Metadata.java
+++ b/tools/code-generation/generator/src/main/java/com/amazonaws/util/awsclientgenerator/domainmodels/codegeneration/Metadata.java
@@ -7,6 +7,7 @@
 
 import lombok.Data;
 
+import java.util.List;
 import java.util.Map;
 
 @Data
@@ -43,4 +44,7 @@ public class Metadata {
     private boolean hasPreSignedUrl;
 
     private boolean awsQueryCompatible;
+
+    // Priority-ordered list of auth types present on the service model
+    private List<String> auth;
 }
diff --git a/tools/code-generation/generator/src/main/java/com/amazonaws/util/awsclientgenerator/domainmodels/codegeneration/Operation.java b/tools/code-generation/generator/src/main/java/com/amazonaws/util/awsclientgenerator/domainmodels/codegeneration/Operation.java
index 240fe611120..2cef55e84de 100644
--- a/tools/code-generation/generator/src/main/java/com/amazonaws/util/awsclientgenerator/domainmodels/codegeneration/Operation.java
+++ b/tools/code-generation/generator/src/main/java/com/amazonaws/util/awsclientgenerator/domainmodels/codegeneration/Operation.java
@@ -23,6 +23,9 @@ public class Operation {
     private boolean virtualAddressAllowed;
     private String virtualAddressMemberName;
     private String authtype;
+    // Non-empty, priority-ordered list of string auth types.
+    // This trait should only be present if its value differs from the service-level trait
+    private List<String> auth; // aws.auth#sigv4 | aws.auth#sigv4a | smithy.api#httpBearerAuth | smithy.api#noAuth
     private String signerName;
     private String authorizer;
     private boolean eventStream;
diff --git a/tools/code-generation/generator/src/main/java/com/amazonaws/util/awsclientgenerator/generators/cpp/CppClientGenerator.java b/tools/code-generation/generator/src/main/java/com/amazonaws/util/awsclientgenerator/generators/cpp/CppClientGenerator.java
index 7ab3845ae22..89f369a1c70 100644
--- a/tools/code-generation/generator/src/main/java/com/amazonaws/util/awsclientgenerator/generators/cpp/CppClientGenerator.java
+++ b/tools/code-generation/generator/src/main/java/com/amazonaws/util/awsclientgenerator/generators/cpp/CppClientGenerator.java
@@ -18,6 +18,7 @@
 import com.amazonaws.util.awsclientgenerator.generators.exceptions.SourceGenerationFailedException;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
 import org.apache.velocity.Template;
 import org.apache.velocity.VelocityContext;
@@ -208,6 +209,40 @@ protected void addRequestIdToResults(final ServiceModel serviceModel) {
                 });
     }
 
+    private static Set<String> servicesMissingMultiAuthMRAPTrait = ImmutableSet.of(
+            "S3",
+            "S3-CRT",
+            "CloudFront KeyValueStore",
+            "SESv2",
+            "EventBridge");
+
+    private void CheckAndEnableSigV4A(final ServiceModel serviceModel, VelocityContext context) {
+        List<String> c2jAuthList = serviceModel.getMetadata().getAuth();
+        String serviceId = serviceModel.getMetadata().getServiceId();
+        if (c2jAuthList != null && c2jAuthList.contains("aws.auth#sigv4a") ||
+             servicesMissingMultiAuthMRAPTrait.contains(serviceId)) {
+            context.put("multiRegionAccessPointSupported", true);
+        }
+        // todo: remove these checks later
+        if (!context.containsKey("multiRegionAccessPointSupported")) {
+            boolean hasSigV4AOperation = serviceModel.getOperations().values().stream()
+                    .anyMatch(op -> op.getAuth() != null && op.getAuth().contains("aws.auth#sigv4a"));
+
+            if (serviceModel.getEndpointRules().contains("\"sigv4a\"") || hasSigV4AOperation) {
+                throw new RuntimeException("Endpoint rules or operation reference sigv4a auth scheme but c2j model " + serviceId +
+                        " does not list aws.auth#sigv4a as a supported auth!");
+            }
+        }
+
+        if (c2jAuthList != null) {
+            boolean hasSigV4AndBearer = c2jAuthList.contains("smithy.api#httpBearerAuth") &&
+                    (c2jAuthList.contains("aws.auth#sigv4a") || c2jAuthList.contains("aws.auth#sigv4"));
+            if (!serviceModel.isUseSmithyClient() && hasSigV4AndBearer) {
+                throw new RuntimeException("SDK Clients cannot mix AWS and Bearer Credentials without enabling Smithy Identity!");
+            }
+        }
+    }
+
     protected final VelocityContext createContext(final ServiceModel serviceModel) {
         VelocityContext context = new VelocityContext();
         context.put("nl", System.lineSeparator());
@@ -216,9 +251,8 @@ protected final VelocityContext createContext(final ServiceModel serviceModel) {
         context.put("output.encoding", StandardCharsets.UTF_8.name());
         context.put("nullChar", '\0');
 
-        if (serviceModel.getEndpointRules().contains("\"sigv4a\"")) {
-            context.put("multiRegionAccessPointSupported", true);
-        }
+        CheckAndEnableSigV4A(serviceModel, context);
+
         return context;
     }