fromx support for authorizers (#2304)

Author: GarrettBeatty

Libraries/src/Amazon.Lambda.Annotations.SourceGenerator/Amazon.Lambda.Annotations.SourceGenerator.csproj

@@ -103,6 +103,14 @@
       <Generator>TextTemplatingFilePreprocessor</Generator>
       <LastGenOutput>ExecutableAssembly.cs</LastGenOutput>
     </None>
+    <None Update="Templates\AuthorizerSetupParameters.tt">
+      <Generator>TextTemplatingFilePreprocessor</Generator>
+      <LastGenOutput>AuthorizerSetupParameters.cs</LastGenOutput>
+    </None>
+    <None Update="Templates\AuthorizerInvoke.tt">
+      <Generator>TextTemplatingFilePreprocessor</Generator>
+      <LastGenOutput>AuthorizerInvoke.cs</LastGenOutput>
+    </None>
   </ItemGroup>
 
   <ItemGroup>
@@ -136,6 +144,16 @@
       <AutoGen>True</AutoGen>
       <DependentUpon>ExecutableAssembly.tt</DependentUpon>
     </Compile>
+    <Compile Update="Templates\AuthorizerSetupParameters.cs">
+      <DesignTime>True</DesignTime>
+      <AutoGen>True</AutoGen>
+      <DependentUpon>AuthorizerSetupParameters.tt</DependentUpon>
+    </Compile>
+    <Compile Update="Templates\AuthorizerInvoke.cs">
+      <DesignTime>True</DesignTime>
+      <AutoGen>True</AutoGen>
+      <DependentUpon>AuthorizerInvoke.tt</DependentUpon>
+    </Compile>
   </ItemGroup>
 
   <ItemGroup>

Libraries/src/Amazon.Lambda.Annotations.SourceGenerator/Diagnostics/AnalyzerReleases.Unshipped.md

@@ -14,3 +14,5 @@ AWSLambda0125 | AWSLambdaCSharpGenerator | Error | Duplicate Authorizer Name
 AWSLambda0127 | AWSLambdaCSharpGenerator | Error | Invalid Result TTL
 AWSLambda0128 | AWSLambdaCSharpGenerator | Error | Authorizer Payload Version Mismatch
 AWSLambda0129 | AWSLambdaCSharpGenerator | Error | Missing LambdaFunction Attribute
+AWSLambda0130 | AWSLambdaCSharpGenerator | Error | Invalid return type IAuthorizerResult
+AWSLambda0131 | AWSLambdaCSharpGenerator | Error | FromBody not supported on Authorizer functions

Libraries/src/Amazon.Lambda.Annotations.SourceGenerator/Diagnostics/DiagnosticDescriptors.cs

@@ -216,7 +216,7 @@ public static class DiagnosticDescriptors
             title: "Authorizer Payload Version Mismatch",
             messageFormat: "The authorizer '{0}' uses AuthorizerPayloadFormatVersion {1} but the endpoint uses HttpApiVersion {2}. This may cause unexpected behavior.",
             category: "AWSLambdaCSharpGenerator",
-            DiagnosticSeverity.Warning,
+            DiagnosticSeverity.Error,
             isEnabledByDefault: true);
 
         public static readonly DiagnosticDescriptor MissingLambdaFunctionAttribute = new DiagnosticDescriptor(
@@ -226,5 +226,21 @@ public static class DiagnosticDescriptors
             category: "AWSLambdaCSharpGenerator",
             DiagnosticSeverity.Error,
             isEnabledByDefault: true);
+
+        public static readonly DiagnosticDescriptor AuthorizerResultOnNonAuthorizerFunction = new DiagnosticDescriptor(
+            id: "AWSLambda0130",
+            title: $"Invalid return type IAuthorizerResult",
+            messageFormat: "IAuthorizerResult is not a valid return type for LambdaFunctions without HttpApiAuthorizer or RestApiAuthorizer attributes",
+            category: "AWSLambdaCSharpGenerator",
+            DiagnosticSeverity.Error,
+            isEnabledByDefault: true);
+
+        public static readonly DiagnosticDescriptor FromBodyNotSupportedOnAuthorizer = new DiagnosticDescriptor(
+            id: "AWSLambda0131",
+            title: "FromBody not supported on Authorizer functions",
+            messageFormat: "[FromBody] is not supported on authorizer functions. Authorizer functions only support [FromHeader], [FromQuery], and [FromRoute] parameter attributes.",
+            category: "AWSLambdaCSharpGenerator",
+            DiagnosticSeverity.Error,
+            isEnabledByDefault: true);
     }
 }

Libraries/src/Amazon.Lambda.Annotations.SourceGenerator/Extensions/ParameterListExtension.cs

@@ -12,7 +12,7 @@ public static bool HasConvertibleParameter(this IList<ParameterModel> parameters
             return parameters.Any(p =>
             {
                 // All request types are forwarded to lambda method if specified, there is no parameter conversion required.
-                if (TypeFullNames.Requests.Contains(p.Type.FullName))
+                if (TypeFullNames.ApiGatewayRequests.Contains(p.Type.FullName))
                 {
                     return false;
                 }

Libraries/src/Amazon.Lambda.Annotations.SourceGenerator/Models/Attributes/AttributeModelBuilder.cs

@@ -90,6 +90,24 @@ public static AttributeModel Build(AttributeData att, GeneratorExecutionContext
                     Type = TypeModelBuilder.Build(att.AttributeClass, context)
                 };
             }
+            else if (att.AttributeClass.Equals(context.Compilation.GetTypeByMetadataName(TypeFullNames.HttpApiAuthorizerAttribute), SymbolEqualityComparer.Default))
+            {
+                var data = HttpApiAuthorizerAttributeBuilder.Build(att);
+                model = new AttributeModel<HttpApiAuthorizerAttribute>
+                {
+                    Data = data,
+                    Type = TypeModelBuilder.Build(att.AttributeClass, context)
+                };
+            }
+            else if (att.AttributeClass.Equals(context.Compilation.GetTypeByMetadataName(TypeFullNames.RestApiAuthorizerAttribute), SymbolEqualityComparer.Default))
+            {
+                var data = RestApiAuthorizerAttributeBuilder.Build(att);
+                model = new AttributeModel<RestApiAuthorizerAttribute>
+                {
+                    Data = data,
+                    Type = TypeModelBuilder.Build(att.AttributeClass, context)
+                };
+            }
             else
             {
                 model = new AttributeModel

Libraries/src/Amazon.Lambda.Annotations.SourceGenerator/Models/EventType.cs

@@ -10,6 +10,7 @@ public enum EventType
         S3,
         SQS,
         DynamoDB,
-        Schedule
+        Schedule,
+        Authorizer
     }
 }

Libraries/src/Amazon.Lambda.Annotations.SourceGenerator/Models/EventTypeBuilder.cs

@@ -26,6 +26,11 @@ public static HashSet<EventType> Build(IMethodSymbol lambdaMethodSymbol,
                 {
                     events.Add(EventType.SQS);
                 }
+                else if (attribute.AttributeClass.ToDisplayString() == TypeFullNames.HttpApiAuthorizerAttribute
+                    || attribute.AttributeClass.ToDisplayString() == TypeFullNames.RestApiAuthorizerAttribute)
+                {
+                    events.Add(EventType.Authorizer);
+                }
             }
 
             return events;

Libraries/src/Amazon.Lambda.Annotations.SourceGenerator/Models/GeneratedMethodModelBuilder.cs

@@ -55,7 +55,7 @@ private static IList<string> BuildUsings(LambdaMethodModel lambdaMethodModel,
 
             namespaces.Add("Amazon.Lambda.Core");
 
-            if(lambdaMethodModel.ReturnsIHttpResults)
+            if(lambdaMethodModel.ReturnsIHttpResults || lambdaMethodModel.ReturnsIAuthorizerResult)
             {
                 namespaces.Add("Amazon.Lambda.Annotations.APIGateway");
             }
@@ -79,6 +79,26 @@ private static TypeModel BuildResponseType(IMethodSymbol lambdaMethodSymbol,
                 return TypeModelBuilder.Build(typeStream, context);
             }
 
+            // For authorizer functions returning IAuthorizerResult, the generated handler returns Stream
+            // (the IAuthorizerResult.Serialize method produces the JSON stream)
+            if (lambdaMethodModel.ReturnsIAuthorizerResult)
+            {
+                var typeStream = context.Compilation.GetTypeByMetadataName(TypeFullNames.Stream);
+                if (lambdaMethodModel.ReturnsGenericTask)
+                {
+                    var genericTask = task.Construct(typeStream);
+                    return TypeModelBuilder.Build(genericTask, context);
+                }
+                return TypeModelBuilder.Build(typeStream, context);
+            }
+
+            // For authorizer functions that return raw API Gateway types (backwards compatibility),
+            // pass through the return type as-is
+            if (lambdaMethodSymbol.HasAttribute(context, TypeFullNames.HttpApiAuthorizerAttribute) ||
+                lambdaMethodSymbol.HasAttribute(context, TypeFullNames.RestApiAuthorizerAttribute))
+            {
+                return lambdaMethodModel.ReturnType;
+            }
 
             if (lambdaMethodSymbol.HasAttribute(context, TypeFullNames.RestApiAttribute))
             {
@@ -143,6 +163,20 @@ private static HttpApiVersion GetHttpApiVersion(IMethodSymbol lambdaMethodSymbol
             return (HttpApiVersion)versionArgument.Value;
         }
 
+        private static AuthorizerPayloadFormatVersion GetAuthorizerPayloadFormatVersion(AttributeData authorizerAttribute)
+        {
+            var versionArg = authorizerAttribute.NamedArguments
+                .FirstOrDefault(arg => arg.Key == "AuthorizerPayloadFormatVersion").Value;
+
+            if (versionArg.Type == null || versionArg.Value == null)
+            {
+                // Default is V2
+                return AuthorizerPayloadFormatVersion.V2;
+            }
+
+            return (AuthorizerPayloadFormatVersion)(int)versionArg.Value;
+        }
+
         private static IList<ParameterModel> BuildParameters(IMethodSymbol lambdaMethodSymbol,
             LambdaMethodModel lambdaMethodModel, GeneratorExecutionContext context)
         {
@@ -158,7 +192,48 @@ private static IList<ParameterModel> BuildParameters(IMethodSymbol lambdaMethodS
                 Documentation = "The ILambdaContext that provides methods for logging and describing the Lambda environment."
             };
 
-            if (lambdaMethodSymbol.HasAttribute(context, TypeFullNames.RestApiAttribute))
+            if (lambdaMethodSymbol.HasAttribute(context, TypeFullNames.HttpApiAuthorizerAttribute))
+            {
+                // For HTTP API authorizer functions, the generated handler accepts the authorizer request type
+                var authorizerAttribute = lambdaMethodSymbol.GetAttributeData(context, TypeFullNames.HttpApiAuthorizerAttribute);
+                var payloadVersion = GetAuthorizerPayloadFormatVersion(authorizerAttribute);
+
+                string requestTypeName;
+                if (payloadVersion == AuthorizerPayloadFormatVersion.V2)
+                {
+                    requestTypeName = TypeFullNames.APIGatewayCustomAuthorizerV2Request;
+                }
+                else
+                {
+                    requestTypeName = TypeFullNames.APIGatewayCustomAuthorizerRequest;
+                }
+
+                var symbol = context.Compilation.GetTypeByMetadataName(requestTypeName);
+                var type = TypeModelBuilder.Build(symbol, context);
+                var requestParameter = new ParameterModel
+                {
+                    Name = "__request__",
+                    Type = type,
+                    Documentation = "The API Gateway authorizer request object that will be processed by the Lambda function handler."
+                };
+                parameters.Add(requestParameter);
+                parameters.Add(contextParameter);
+            }
+            else if (lambdaMethodSymbol.HasAttribute(context, TypeFullNames.RestApiAuthorizerAttribute))
+            {
+                // For REST API authorizer functions, always use APIGatewayCustomAuthorizerRequest
+                var symbol = context.Compilation.GetTypeByMetadataName(TypeFullNames.APIGatewayCustomAuthorizerRequest);
+                var type = TypeModelBuilder.Build(symbol, context);
+                var requestParameter = new ParameterModel
+                {
+                    Name = "__request__",
+                    Type = type,
+                    Documentation = "The API Gateway authorizer request object that will be processed by the Lambda function handler."
+                };
+                parameters.Add(requestParameter);
+                parameters.Add(contextParameter);
+            }
+            else if (lambdaMethodSymbol.HasAttribute(context, TypeFullNames.RestApiAttribute))
             {
                 var symbol = context.Compilation.GetTypeByMetadataName(TypeFullNames.APIGatewayProxyRequest);
                 var type = TypeModelBuilder.Build(symbol, context);

Libraries/src/Amazon.Lambda.Annotations.SourceGenerator/Models/LambdaMethodModel.cs

@@ -64,6 +64,31 @@ public bool ReturnsIHttpResults
             }
         }
 
+        /// <summary>
+        /// Returns true if the Lambda function returns either IAuthorizerResult or Task&lt;IAuthorizerResult&gt;
+        /// </summary>
+        public bool ReturnsIAuthorizerResult
+        {
+            get
+            {
+                if (ReturnsVoid)
+                {
+                    return false;
+                }
+
+                if (ReturnType.FullName == TypeFullNames.IAuthorizerResult)
+                {
+                    return true;
+                }
+                if (ReturnsGenericTask && ReturnType.TypeArguments.Count == 1 && ReturnType.TypeArguments[0].FullName == TypeFullNames.IAuthorizerResult)
+                {
+                    return true;
+                }
+
+                return false;
+            }
+        }
+
         /// <summary>
         /// Returns true if the Lambda function returns either void, Task, SQSBatchResponse or Task<SQSBatchResponse>
         /// </summary>

Libraries/src/Amazon.Lambda.Annotations.SourceGenerator/Templates/APIGatewaySetupParameters.cs

@@ -45,7 +45,7 @@ public virtual string TransformText()
                 }
 
                 // Pass the same request parameter for Request Type that comes from the generated method.
-                if (TypeFullNames.Requests.Contains(p.Type.FullName))
+                if (TypeFullNames.ApiGatewayRequests.Contains(p.Type.FullName))
                 {
                     return "__request__";
                 }
@@ -83,7 +83,7 @@ public virtual string TransformText()
 
         foreach (var parameter in _model.LambdaMethod.Parameters)
         {
-            if (parameter.Type.FullName == TypeFullNames.ILambdaContext || TypeFullNames.Requests.Contains(parameter.Type.FullName))
+            if (parameter.Type.FullName == TypeFullNames.ILambdaContext || TypeFullNames.ApiGatewayRequests.Contains(parameter.Type.FullName))
             {
                 // No action required for ILambdaContext and RequestType, they are passed from the generated method parameter directly to the original method.
             }