Find and use AuthSchemeFactories in dynamic client

Author: mtdowling

client/dynamic-client/build.gradle.kts

@@ -10,6 +10,7 @@ extra["moduleName"] = "software.amazon.smithy.java.dynamicclient"
 dependencies {
     api(project(":dynamic-schemas"))
     api(project(":client:client-core"))
+    implementation(project(":logging"))
 
     testImplementation(project(":client:client-rulesengine"))
     testImplementation(project(":aws:client:aws-client-restjson"))

client/dynamic-client/src/main/java/software/amazon/smithy/java/dynamicclient/plugins/SimpleAuthDetectionPlugin.java

@@ -5,20 +5,37 @@
 
 package software.amazon.smithy.java.dynamicclient.plugins;
 
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.ServiceLoader;
 import software.amazon.smithy.java.client.core.AutoClientPlugin;
 import software.amazon.smithy.java.client.core.ClientConfig;
+import software.amazon.smithy.java.client.core.auth.scheme.AuthSchemeFactory;
 import software.amazon.smithy.java.client.core.auth.scheme.AuthSchemeResolver;
 import software.amazon.smithy.java.dynamicclient.settings.ModelSetting;
 import software.amazon.smithy.java.dynamicclient.settings.ServiceIdSetting;
+import software.amazon.smithy.java.logging.InternalLogger;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.knowledge.ServiceIndex;
 import software.amazon.smithy.model.shapes.ShapeId;
 
 /**
  * A plugin used to detect if built-in auth schemes can be applied to a client automatically.
  */
+@SuppressWarnings("rawtypes")
 public final class SimpleAuthDetectionPlugin implements AutoClientPlugin {
 
+    private static final InternalLogger LOGGER = InternalLogger.getLogger(SimpleAuthDetectionPlugin.class);
+
+    private static final Map<ShapeId, AuthSchemeFactory> AUTH_SCHEME_FACTORIES = new HashMap<>();
+    static {
+        for (var factory : ServiceLoader.load(AuthSchemeFactory.class,
+                SimpleAuthDetectionPlugin.class.getClassLoader())) {
+            AUTH_SCHEME_FACTORIES.put(factory.schemeId(), factory);
+        }
+    }
+
     public static final SimpleAuthDetectionPlugin INSTANCE = new SimpleAuthDetectionPlugin();
 
     @Override
@@ -40,14 +57,35 @@ public void configureClient(ClientConfig.Builder config) {
         }
     }
 
+    @SuppressWarnings("unchecked")
     private void injectAuthSchemeResolver(ClientConfig.Builder config, Model model, ShapeId service) {
         var index = ServiceIndex.of(model);
         var potentialAuthSchemes = index.getEffectiveAuthSchemes(service);
         if (potentialAuthSchemes.isEmpty()) {
             config.authSchemeResolver(AuthSchemeResolver.NO_AUTH);
-        } else {
-            // TODO: Add similar behavior as done in ClientInterfaceGenerator to register AuthSchemeFactories.
-            config.authSchemeResolver(AuthSchemeResolver.DEFAULT);
+            return;
+        }
+
+        // Make a set of the auth schemes explicitly configured since we don't want to overwrite them.
+        var existingSchemeIds = new HashSet<>(config.supportedAuthSchemes().size());
+        for (var scheme : config.supportedAuthSchemes()) {
+            existingSchemeIds.add(scheme.schemeId());
         }
+
+        for (var entry : potentialAuthSchemes.entrySet()) {
+            var id = entry.getKey();
+            if (existingSchemeIds.contains(id)) {
+                continue;
+            }
+            var factory = AUTH_SCHEME_FACTORIES.get(id);
+            if (factory != null) {
+                config.putSupportedAuthSchemes(factory.createAuthScheme(entry.getValue()));
+            } else {
+                LOGGER.warn("Could not find software.amazon.smithy.java.client.core.auth.scheme.AuthSchemeFactory "
+                        + "implementation for auth scheme {}", id);
+            }
+        }
+
+        config.authSchemeResolver(AuthSchemeResolver.DEFAULT);
     }
 }

client/dynamic-client/src/test/java/software/amazon/smithy/java/dynamicclient/plugins/SimpleAuthDetectionPluginTest.java

@@ -0,0 +1,90 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.java.dynamicclient.plugins;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.is;
+
+import org.junit.jupiter.api.Test;
+import software.amazon.smithy.aws.traits.auth.SigV4Trait;
+import software.amazon.smithy.java.client.core.auth.scheme.AuthSchemeResolver;
+import software.amazon.smithy.java.client.core.endpoint.EndpointResolver;
+import software.amazon.smithy.java.dynamicclient.DynamicClient;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.shapes.ShapeId;
+
+class SimpleAuthDetectionPluginTest {
+
+    @Test
+    void registersAuthSchemeFactoriesForServiceAuthTraits() {
+        var model = Model.assembler()
+                .addUnparsedModel("test.smithy", """
+                        $version: "2"
+                        namespace smithy.example
+
+                        use aws.auth#sigv4
+                        use aws.protocols#awsJson1_0
+
+                        @awsJson1_0
+                        @sigv4(name: "testservice")
+                        service AuthService {
+                            operations: [DoThing]
+                        }
+
+                        operation DoThing {
+                            input := {}
+                            output := {}
+                        }
+                        """)
+                .discoverModels()
+                .assemble()
+                .unwrap();
+
+        var client = DynamicClient.builder()
+                .serviceId(ShapeId.from("smithy.example#AuthService"))
+                .model(model)
+                .endpointResolver(EndpointResolver.staticEndpoint("https://example.com"))
+                .build();
+
+        var authSchemes = client.config().supportedAuthSchemes();
+        var hasSigV4 = authSchemes.stream().anyMatch(s -> s.schemeId().equals(SigV4Trait.ID));
+        assertThat("Expected SigV4 auth scheme to be registered", hasSigV4, is(true));
+        assertThat(client.config().authSchemeResolver(), is(AuthSchemeResolver.DEFAULT));
+    }
+
+    @Test
+    void setsNoAuthResolverWhenNoAuthTraits() {
+        var model = Model.assembler()
+                .addUnparsedModel("test.smithy", """
+                        $version: "2"
+                        namespace smithy.example
+
+                        use aws.protocols#awsJson1_0
+
+                        @awsJson1_0
+                        service NoAuthService {
+                            operations: [DoThing]
+                        }
+
+                        operation DoThing {
+                            input := {}
+                            output := {}
+                        }
+                        """)
+                .discoverModels()
+                .assemble()
+                .unwrap();
+
+        var client = DynamicClient.builder()
+                .serviceId(ShapeId.from("smithy.example#NoAuthService"))
+                .model(model)
+                .endpointResolver(EndpointResolver.staticEndpoint("https://example.com"))
+                .build();
+
+        assertThat(client.config().authSchemeResolver(), equalTo(AuthSchemeResolver.NO_AUTH));
+    }
+}