Skip to content

Commit ff70582

Browse files
JordonPhillipsJordonPhillips
committed
Code generate non-http auth
1 parent f759355 commit ff70582

File tree

16 files changed

+506
-549
lines changed

16 files changed

+506
-549
lines changed

‎codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsAuthIntegration.java

Lines changed: 19 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -7,18 +7,20 @@
77
import static software.amazon.smithy.python.aws.codegen.AwsConfiguration.REGION;
88

99
import java.util.List;
10+
import java.util.Set;
1011
import software.amazon.smithy.aws.traits.auth.SigV4Trait;
1112
import software.amazon.smithy.codegen.core.Symbol;
13+
import software.amazon.smithy.model.shapes.ServiceShape;
1214
import software.amazon.smithy.model.shapes.ShapeId;
1315
import software.amazon.smithy.python.codegen.ApplicationProtocol;
1416
import software.amazon.smithy.python.codegen.CodegenUtils;
1517
import software.amazon.smithy.python.codegen.ConfigProperty;
16-
import software.amazon.smithy.python.codegen.DerivedProperty;
1718
import software.amazon.smithy.python.codegen.GenerationContext;
1819
import software.amazon.smithy.python.codegen.SmithyPythonDependency;
1920
import software.amazon.smithy.python.codegen.integrations.AuthScheme;
2021
import software.amazon.smithy.python.codegen.integrations.PythonIntegration;
2122
import software.amazon.smithy.python.codegen.integrations.RuntimeClientPlugin;
23+
import software.amazon.smithy.python.codegen.writer.PythonWriter;
2224
import software.amazon.smithy.utils.SmithyInternalApi;
2325

2426
/**
@@ -38,7 +40,7 @@ public List<RuntimeClientPlugin> getClientPlugins(GenerationContext context) {
3840
.name("aws_credentials_identity_resolver")
3941
.documentation("Resolves AWS Credentials. Required for operations that use Sigv4 Auth.")
4042
.type(Symbol.builder()
41-
.name("IdentityResolver[AWSCredentialsIdentity, IdentityProperties]")
43+
.name("IdentityResolver[AWSCredentialsIdentity, AWSIdentityProperties]")
4244
.addReference(Symbol.builder()
4345
.addDependency(SmithyPythonDependency.SMITHY_CORE)
4446
.name("IdentityResolver")
@@ -51,8 +53,8 @@ public List<RuntimeClientPlugin> getClientPlugins(GenerationContext context) {
5153
.build())
5254
.addReference(Symbol.builder()
5355
.addDependency(SmithyPythonDependency.SMITHY_CORE)
54-
.name("IdentityProperties")
55-
.namespace("smithy_core.interfaces.identity", ".")
56+
.name("AWSIdentityProperties")
57+
.namespace("smithy_aws_core.identity", ".")
5658
.build())
5759
.build())
5860
// TODO: Initialize with the provider chain?
@@ -69,30 +71,26 @@ public void customize(GenerationContext context) {
6971
return;
7072
}
7173
var trait = context.settings().service(context.model()).expectTrait(SigV4Trait.class);
72-
var params = CodegenUtils.getHttpAuthParamsSymbol(context.settings());
7374
var resolver = CodegenUtils.getHttpAuthSchemeResolverSymbol(context.settings());
7475

7576
// Add a function that generates the http auth option for api key auth.
7677
// This needs to be generated because there's modeled parameters that
7778
// must be accounted for.
7879
context.writerDelegator().useFileWriter(resolver.getDefinitionFile(), resolver.getNamespace(), writer -> {
7980
writer.addDependency(SmithyPythonDependency.SMITHY_HTTP);
80-
writer.addImport("smithy_http.aio.interfaces.auth", "HTTPAuthOption");
81+
writer.addImport("smithy_core.interfaces.auth", "AuthOption", "_AuthOption");
82+
writer.addImports("smithy_core.auth", Set.of("AuthOption", "AuthParams"));
8183
writer.pushState();
8284

8385
writer.write("""
84-
def $1L(auth_params: $2T) -> HTTPAuthOption | None:
85-
return HTTPAuthOption(
86-
scheme_id=$3S,
87-
identity_properties={},
88-
signer_properties={
89-
"service": $4S,
90-
"region": auth_params.region
91-
}
86+
def $1L(auth_params: AuthParams[Any, Any]) -> _AuthOption | None:
87+
return AuthOption(
88+
scheme_id=$2S,
89+
identity_properties={}, # type: ignore
90+
signer_properties={} # type: ignore
9291
)
9392
""",
9493
SIGV4_OPTION_GENERATOR_NAME,
95-
params,
9694
SigV4Trait.ID.toString(),
9795
trait.getName());
9896
writer.popState();
@@ -119,17 +117,6 @@ public ApplicationProtocol getApplicationProtocol() {
119117
return ApplicationProtocol.createDefaultHttpApplicationProtocol();
120118
}
121119

122-
@Override
123-
public List<DerivedProperty> getAuthProperties() {
124-
return List.of(
125-
DerivedProperty.builder()
126-
.name("region")
127-
.source(DerivedProperty.Source.CONFIG)
128-
.type(Symbol.builder().name("str").build())
129-
.sourcePropertyName("region")
130-
.build());
131-
}
132-
133120
@Override
134121
public Symbol getAuthOptionGenerator(GenerationContext context) {
135122
var resolver = CodegenUtils.getHttpAuthSchemeResolverSymbol(context.settings());
@@ -148,5 +135,11 @@ public Symbol getAuthSchemeSymbol(GenerationContext context) {
148135
.addDependency(AwsPythonDependency.SMITHY_AWS_CORE)
149136
.build();
150137
}
138+
139+
@Override
140+
public void initializeScheme(GenerationContext context, PythonWriter writer, ServiceShape service) {
141+
var trait = service.expectTrait(SigV4Trait.class);
142+
writer.write("$T(service=$S)", getAuthSchemeSymbol(context), trait.getName());
143+
}
151144
}
152145
}

‎codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java

Lines changed: 45 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -300,8 +300,10 @@ def _classify_error(
300300
writer.consumer(w -> context.protocolGenerator().wrapInputStream(context, w)),
301301
writer.consumer(w -> context.protocolGenerator().wrapOutputStream(context, w)));
302302
}
303+
303304
writer.addStdlibImport("typing", "Any");
304305
writer.addStdlibImport("asyncio", "iscoroutine");
306+
writer.addImport("smithy_core.auth", "AuthParams");
305307
writer.write(
306308
"""
307309
async def _execute_operation[Input: SerializeableShape, Output: DeserializeableShape](
@@ -466,7 +468,6 @@ await sleep(retry_token.retry_delay)
466468
try:
467469
# Step 7a: Invoke read_before_attempt
468470
interceptor.read_before_attempt(context)
469-
470471
""",
471472
pluginSymbol,
472473
transportRequest,
@@ -476,53 +477,60 @@ await sleep(retry_token.retry_delay)
476477

477478
boolean supportsAuth = !ServiceIndex.of(model).getAuthSchemes(service).isEmpty();
478479
writer.pushState(new ResolveIdentitySection());
479-
if (context.applicationProtocol().isHttpProtocol() && supportsAuth) {
480-
writer.pushState(new InitializeHttpAuthParametersSection());
481-
writer.write("""
482-
# Step 7b: Invoke service_auth_scheme_resolver.resolve_auth_scheme
483-
auth_parameters: $1T = $1T(
484-
operation=operation.schema.id.name,
485-
${2C|}
486-
)
487-
488-
""",
489-
CodegenUtils.getHttpAuthParamsSymbol(context.settings()),
490-
writer.consumer(this::initializeHttpAuthParameters));
491-
writer.popState();
480+
if (supportsAuth) {
481+
// TODO: delete InitializeHttpAuthParametersSection
492482

493483
writer.addDependency(SmithyPythonDependency.SMITHY_CORE);
494-
writer.addDependency(SmithyPythonDependency.SMITHY_HTTP);
495484
writer.addImport("smithy_core.interfaces.identity", "Identity");
496-
writer.addImports("smithy_http.aio.interfaces.auth", Set.of("HTTPSigner", "HTTPAuthOption"));
485+
writer.addImport("smithy_core.interfaces.auth", "AuthOption");
486+
writer.addImport("smithy_core.aio.interfaces.auth", "Signer");
487+
writer.addImport("smithy_core.shapes", "ShapeID");
497488
writer.addStdlibImport("typing", "Any");
498489
writer.write("""
499-
auth_options = config.http_auth_scheme_resolver.resolve_auth_scheme(
490+
auth_parameters = AuthParams(
491+
protocol_id=ShapeID($1S),
492+
operation=operation,
493+
context=context.properties,
494+
)
495+
auth_options = config.auth_scheme_resolver.resolve_auth_scheme(
500496
auth_parameters=auth_parameters
501497
)
502-
auth_option: HTTPAuthOption | None = None
498+
499+
auth_option: AuthOption | None = None
503500
for option in auth_options:
504-
if option.scheme_id in config.http_auth_schemes:
501+
if option.scheme_id in config.auth_schemes:
505502
auth_option = option
506503
break
507504
508-
signer: HTTPSigner[Any, Any] | None = None
505+
signer: Signer[$2T, Any, Any] | None = None
509506
identity: Identity | None = None
507+
auth_scheme: Any = None
510508
511509
if auth_option:
512-
auth_scheme = config.http_auth_schemes[auth_option.scheme_id]
510+
auth_scheme = config.auth_schemes[auth_option.scheme_id]
511+
context.properties["auth_scheme"] = auth_scheme
513512
514513
# Step 7c: Invoke auth_scheme.identity_resolver
515-
identity_resolver = auth_scheme.identity_resolver(config=config)
514+
identity_resolver = auth_scheme.identity_resolver(context=context.properties)
515+
context.properties["identity_resolver"] = identity_resolver
516516
517517
# Step 7d: Invoke auth_scheme.signer
518-
signer = auth_scheme.signer
518+
signer = auth_scheme.signer()
519+
520+
# TODO: merge from auth_option
521+
identity_properties = auth_scheme.identity_properties(
522+
context=context.properties
523+
)
524+
context.properties["identity_properties"] = identity_properties
519525
520526
# Step 7e: Invoke identity_resolver.get_identity
521527
identity = await identity_resolver.get_identity(
522-
identity_properties=auth_option.identity_properties
528+
identity_properties=identity_properties
523529
)
524530
525-
""");
531+
""",
532+
context.protocolGenerator().getProtocol(),
533+
transportRequest);
526534
}
527535
writer.popState();
528536

@@ -580,45 +588,29 @@ await sleep(retry_token.retry_delay)
580588
""");
581589

582590
writer.pushState(new SignRequestSection());
583-
if (context.applicationProtocol().isHttpProtocol() && supportsAuth) {
591+
if (supportsAuth) {
584592
writer.addStdlibImport("re");
585593
writer.addStdlibImport("typing", "Any");
586594
writer.addImport("smithy_core.interfaces.identity", "Identity");
587595
writer.addImport("smithy_core.types", "PropertyKey");
588596
writer.write("""
589597
# Step 7i: sign the request
590598
if auth_option and signer:
591-
logger.debug("HTTP request to sign: %s", context.transport_request)
592-
logger.debug(
593-
"Signer properties: %s",
594-
auth_option.signer_properties
595-
)
599+
signer_properties = auth_scheme.signer_properties(context=context.properties)
600+
context.properties["signer_properties"] = signer_properties
601+
602+
logger.debug("Request to sign: %s", context.transport_request)
603+
logger.debug("Signer properties: %s", signer_properties)
604+
596605
context = replace(
597606
context,
598-
transport_request= await signer.sign(
599-
http_request=context.transport_request,
607+
transport_request = await signer.sign(
608+
request=context.transport_request,
600609
identity=identity,
601-
signing_properties=auth_option.signer_properties,
610+
properties=signer_properties,
602611
)
603612
)
604613
logger.debug("Signed HTTP request: %s", context.transport_request)
605-
606-
# TODO - Move this to separate resolution/population function
607-
fields = context.transport_request.fields
608-
auth_value = fields["Authorization"].as_string() # type: ignore
609-
signature = re.split("Signature=", auth_value)[-1] # type: ignore
610-
context.properties["signature"] = signature.encode('utf-8')
611-
612-
identity_key: PropertyKey[Identity | None] = PropertyKey(
613-
key="identity",
614-
value_type=Identity | None # type: ignore
615-
)
616-
sp_key: PropertyKey[dict[str, Any]] = PropertyKey(
617-
key="signer_properties",
618-
value_type=dict[str, Any] # type: ignore
619-
)
620-
context.properties[identity_key] = identity
621-
context.properties[sp_key] = auth_option.signer_properties
622614
""");
623615
}
624616
writer.popState();
@@ -788,28 +780,6 @@ private boolean hasEventStream() {
788780
return false;
789781
}
790782

791-
private void initializeHttpAuthParameters(PythonWriter writer) {
792-
var derived = new LinkedHashSet<DerivedProperty>();
793-
for (PythonIntegration integration : context.integrations()) {
794-
for (RuntimeClientPlugin plugin : integration.getClientPlugins(context)) {
795-
if (plugin.matchesService(model, service)
796-
&& plugin.getAuthScheme().isPresent()
797-
&& plugin.getAuthScheme().get().getApplicationProtocol().isHttpProtocol()) {
798-
derived.addAll(plugin.getAuthScheme().get().getAuthProperties());
799-
}
800-
}
801-
}
802-
803-
for (DerivedProperty property : derived) {
804-
var source = property.source().scopeLocation();
805-
if (property.initializationFunction().isPresent()) {
806-
writer.write("$L=$T($L),", property.name(), property.initializationFunction().get(), source);
807-
} else if (property.sourcePropertyName().isPresent()) {
808-
writer.write("$L=$L.$L,", property.name(), source, property.sourcePropertyName().get());
809-
}
810-
}
811-
}
812-
813783
private void writeDefaultPlugins(PythonWriter writer, Collection<SymbolReference> plugins) {
814784
for (SymbolReference plugin : plugins) {
815785
writer.write("$T,", plugin);
@@ -873,8 +843,8 @@ private void writeSharedOperationInit(PythonWriter writer, OperationShape operat
873843
.orElse("The operation's input.");
874844

875845
writer.write("""
876-
$L
877-
""",docs);
846+
$L
847+
""", docs);
878848
writer.write("");
879849
writer.write(":param input: $L", inputDocs);
880850
writer.write("");

‎codegen/core/src/main/java/software/amazon/smithy/python/codegen/CodegenUtils.java

Lines changed: 0 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -139,20 +139,6 @@ public static Symbol getUnknownApiError(PythonSettings settings) {
139139
.build();
140140
}
141141

142-
/**
143-
* Gets the symbol for the http auth params.
144-
*
145-
* @param settings The client settings, used to account for module configuration.
146-
* @return Returns the http auth params symbol.
147-
*/
148-
public static Symbol getHttpAuthParamsSymbol(PythonSettings settings) {
149-
return Symbol.builder()
150-
.name("HTTPAuthParams")
151-
.namespace(String.format("%s.auth", settings.moduleName()), ".")
152-
.definitionFile(String.format("./src/%s/auth.py", settings.moduleName()))
153-
.build();
154-
}
155-
156142
/**
157143
* Gets the symbol for the http auth scheme resolver.
158144
*

0 commit comments

Comments
 (0)