Merge pull request #159 from TerraConstructs/fix/edge-function-associations-and-function-url

fix: render Distribution functionAssociations; grant lambda:InvokeFunction for public Function URLs

Author: so0k

integ/aws/edge/Makefile

@@ -21,6 +21,10 @@ distribution-policies: ## Test Distribution Policies
 	go test -v -timeout 45m ./... -run ^TestDistributionPolicies$
 .PHONY: distribution-policies
 
+distribution-function: ## Test Distribution Function association
+	go test -v -timeout 45m ./... -run ^TestDistributionFunction$
+.PHONY: distribution-function
+
 service-with-http-namespace: ## Test CloudMap Service with HTTP Namespace
 	go test -v -timeout 30m ./... -run ^TestServiceWithHttpNamespace$
 .PHONY: service-with-http-namespace

integ/aws/edge/README.md

@@ -15,6 +15,8 @@ Test Targets:
   kvs-jwt-verify             Test Edge function for KVS JWT verify
   multi-zone-acm-pub-cert    Test Multi Zone ACM Public Certificate
   distribution-policies      Test Distribution Policies
+  distribution-function      Test Distribution Function association
+  service-with-http-namespace Test CloudMap Service with HTTP Namespace
 
 Other Targets:
   help                       Print out every target with a description

integ/aws/edge/apps/distribution-function.ts

@@ -0,0 +1,75 @@
+// https://github.com/aws/aws-cdk/blob/7926560f0a150d8fd39d0775df5259621b8068ae/packages/@aws-cdk-testing/framework-integ/test/aws-cloudfront/test/integ.distribution-function.ts
+import { cloudfrontDistribution } from "@cdktn/provider-aws";
+import { App, LocalBackend } from "cdktn";
+import { aws } from "../../../../src";
+
+const environmentName = process.env.ENVIRONMENT_NAME ?? "test";
+const region = process.env.AWS_REGION ?? "us-east-1";
+const outdir = process.env.OUT_DIR ?? "cdktf.out";
+const stackName = process.env.STACK_NAME ?? "distribution-function";
+
+// https://github.com/aws/aws-cdk/blob/17b12f2aa7a2b519a6e802bf79d3099f2fcd7851/packages/@aws-cdk-testing/framework-integ/test/aws-cloudfront/test/test-origin.ts
+/** Used for testing common Origin functionality */
+class TestOrigin extends aws.edge.OriginBase {
+  constructor(domainName: string, props: aws.edge.OriginProps = {}) {
+    super(domainName, props);
+  }
+  protected renderCustomOriginConfig():
+    | cloudfrontDistribution.CloudfrontDistributionOriginCustomOriginConfig
+    | undefined {
+    return {
+      httpPort: 80,
+      httpsPort: 443,
+      originProtocolPolicy: aws.edge.OriginProtocolPolicy.HTTPS_ONLY,
+      originSslProtocols: [aws.edge.OriginSslPolicy.TLS_V1_2],
+    };
+  }
+}
+
+const app = new App({
+  outdir,
+});
+const stack = new aws.AwsStack(app, stackName, {
+  gridUUID: "g12345678-1234",
+  environmentName,
+  providerConfig: {
+    region,
+  },
+});
+
+new LocalBackend(stack, {
+  path: `${stackName}.tfstate`,
+});
+
+// Viewer-request function that stamps a marker header on the request so the
+// association's effect is directly observable (via TestFunction and, once
+// deployed, on the actual viewer response echoed back by the origin).
+const cfFunction = new aws.edge.Function(stack, "Function", {
+  nameSuffix: "distribution-function",
+  code: aws.edge.FunctionCode.fromInline(
+    `function handler(event) {
+  var request = event.request;
+  request.headers['x-distribution-function'] = { value: 'true' };
+  return request;
+}`,
+  ),
+  registerOutputs: true,
+  outputName: "function",
+});
+
+new aws.edge.Distribution(stack, "Dist", {
+  defaultBehavior: {
+    origin: new TestOrigin("www.example.com"),
+    cachePolicy: aws.edge.ManagedCachePolicy.CACHING_DISABLED,
+    functionAssociations: [
+      {
+        function: cfFunction,
+        eventType: aws.edge.FunctionEventType.VIEWER_REQUEST,
+      },
+    ],
+  },
+  registerOutputs: true,
+  outputName: "distribution",
+});
+
+app.synth();

integ/aws/edge/edge_test.go

@@ -4,10 +4,12 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
+	cftypes "github.com/aws/aws-sdk-go-v2/service/cloudfront/types"
 	"github.com/aws/aws-sdk-go-v2/service/servicediscovery/types"
 	"github.com/stretchr/testify/require"
 	"github.com/terraconstructs/base/integ"
@@ -58,6 +60,14 @@ func TestDistributionPolicies(t *testing.T) {
 		})
 }
 
+// Run the apps/distribution-function.ts integration test
+// ref: https://github.com/TerraConstructs/base/issues/50
+// ref: https://github.com/TerraConstructs/base/issues/99
+func TestDistributionFunction(t *testing.T) {
+	envVars := executors.EnvMap(os.Environ())
+	runEdgeIntegrationTest(t, "distribution-function", "us-east-1", envVars, validateDistributionFunction)
+}
+
 // Test the apps/service-with-http-namespace.ts app
 // ref: https://github.com/aws/aws-cdk/blob/v2.233.0/packages/@aws-cdk-testing/framework-integ/test/aws-servicediscovery/test/integ.service-with-http-namespace.lit.ts
 func TestServiceWithHttpNamespace(t *testing.T) {
@@ -187,6 +197,73 @@ func validateURLRewriteFunction(t *testing.T, workingDir string, _awsRegion stri
 	}
 }
 
+// validateDistributionFunction verifies the aws.edge.Distribution's
+// defaultBehavior functionAssociations wiring from apps/distribution-function.ts:
+// it waits for the distribution to deploy, then confirms the associated
+// viewer-request CloudFront Function actually runs by exercising the
+// TestFunction API and asserting the marker header it stamps on the request.
+func validateDistributionFunction(t *testing.T, workingDir string, awsRegion string) {
+	// Load the Terraform Options saved by the earlier deploy_terraform stage
+	terraformOptions := test_structure.LoadTerraformOptions(t, workingDir)
+
+	distributionId := util.LoadOutputAttribute(t, terraformOptions, "distribution", "id")
+	util.WaitForDistributionDeployed(t, awsRegion, distributionId, 10, 10*time.Second)
+
+	functionName := util.LoadOutputAttribute(t, terraformOptions, "function", "name")
+
+	// Assert the deployed distribution config actually carries the
+	// viewer-request FunctionAssociation for the function -- this is the
+	// direct regression check for #99/#50 (a dropped association would
+	// still let the distribution deploy and the TestFunction call below
+	// would still succeed, since TestFunction invokes the function by name
+	// independent of any distribution).
+	dist, err := util.GetDistributionE(t, awsRegion, distributionId)
+	require.NoError(t, err)
+	functionAssociations := dist.DistributionConfig.DefaultCacheBehavior.FunctionAssociations
+	require.NotNil(t, functionAssociations)
+	require.EqualValues(t, 1, aws.ToInt32(functionAssociations.Quantity))
+	require.Len(t, functionAssociations.Items, 1)
+	require.Equal(t, cftypes.EventTypeViewerRequest, functionAssociations.Items[0].EventType)
+	functionArn := aws.ToString(functionAssociations.Items[0].FunctionARN)
+	require.NotEmpty(t, functionArn)
+	require.True(t, strings.HasSuffix(functionArn, functionName),
+		"expected FunctionARN %q to end with function name %q", functionArn, functionName)
+
+	functionStage := "LIVE"
+	testEvent := &util.CloudFrontFunctionEvent{
+		Version: "1.0",
+		Context: util.Context{
+			DistributionDomainName: "d111111abcdef8.cloudfront.net",
+			DistributionID:         distributionId,
+			EventType:              "viewer-request",
+			RequestID:              "test-request-id",
+		},
+		Viewer: util.Viewer{
+			IP: "1.2.3.4",
+		},
+		Request: &util.Request{
+			Method:      "GET",
+			URI:         "/",
+			Querystring: util.ValueObject{},
+			Headers: util.ValueObject{
+				"host": util.ValueEntry{Value: "d111111abcdef8.cloudfront.net"},
+			},
+		},
+	}
+	util.TestCloudFrontFunctionWithCustomValidation(t, functionName, functionStage, *testEvent,
+		func(r *util.CloudFrontTestFunctionResult) error {
+			if r.Output == nil {
+				return fmt.Errorf("got nil Output response")
+			}
+			return integ.AssertE(r.Output, []integ.Assertion{
+				{
+					Path:           "request.headers.\"x-distribution-function\".value",
+					ExpectedRegexp: strPtr("^true$"),
+				},
+			})
+		})
+}
+
 // validateJwtVerifyFunction with testevents
 func validateJwtVerifyFunction(t *testing.T, workingDir string, _awsRegion string) {
 	// Load the Terraform Options saved by the earlier deploy_terraform stage

src/aws/compute/function-base.ts

@@ -372,6 +372,7 @@ export abstract class LambdaFunctionBase
       sourceArn: permission.sourceArn ?? sourceArn,
       principalOrgId: permission.organizationId ?? principalOrgID,
       functionUrlAuthType: permission.functionUrlAuthType,
+      invokedViaFunctionUrl: permission.invokedViaFunctionUrl,
     });
   }
 

src/aws/compute/function-permission.ts

@@ -88,4 +88,18 @@ export interface Permission {
    * @default - No functionUrlAuthType
    */
   readonly functionUrlAuthType?: FunctionUrlAuthType;
+
+  /**
+   * Restricts this permission to only apply to invocations that go through a
+   * Lambda Function URL (i.e. adds the `lambda:InvokedViaFunctionUrl`
+   * condition key to the generated resource policy statement).
+   *
+   * This is used, for example, to scope the `lambda:InvokeFunction`
+   * permission that is required (in addition to `lambda:InvokeFunctionUrl`)
+   * for a public (`FunctionUrlAuthType.NONE`) function URL, mirroring the
+   * `FunctionURLInvokeAllowPublicAccess` statement the AWS Console adds.
+   *
+   * @default - no lambda:InvokedViaFunctionUrl condition is added to the statement
+   */
+  readonly invokedViaFunctionUrl?: boolean;
 }

src/aws/compute/function-url.ts

@@ -260,12 +260,27 @@ export class FunctionUrl extends AwsConstructBase implements IFunctionUrl {
     this.functionArn = this.resource.functionArn;
     this.function = props.function;
 
-    if (props.authType === FunctionUrlAuthType.NONE) {
+    if (this.authType === FunctionUrlAuthType.NONE) {
       props.function.addPermission("invoke-function-url", {
         principal: new iam.AnyPrincipal(),
         action: "lambda:InvokeFunctionUrl",
         functionUrlAuthType: props.authType,
       });
+      // A public (authType NONE) Function URL also requires a standalone
+      // lambda:InvokeFunction grant - lambda:InvokeFunctionUrl alone is not
+      // sufficient and unauthenticated callers otherwise receive a 403.
+      // This mirrors the "FunctionURLInvokeAllowPublicAccess" statement the AWS
+      // Console/CLI add automatically, which is scoped down using the
+      // lambda:InvokedViaFunctionUrl condition key so this permission only
+      // applies to invocations made through the function URL (not direct
+      // lambda:InvokeFunction calls). The `invoked_via_function_url`
+      // argument on aws_lambda_permission maps 1:1 to that condition key.
+      // See: https://docs.aws.amazon.com/lambda/latest/dg/urls-auth.html
+      props.function.addPermission("invoke-function-url-via-invoke", {
+        principal: new iam.AnyPrincipal(),
+        action: "lambda:InvokeFunction",
+        invokedViaFunctionUrl: true,
+      });
     }
     this.functionUrlOutputs = {
       url: this.url,

src/aws/edge/distribution.ts

@@ -5,9 +5,16 @@ import {
   dataAwsCloudfrontOriginRequestPolicy,
   dataAwsCloudfrontResponseHeadersPolicy,
 } from "@cdktn/provider-aws";
-import { IResolvable, Token, Lazy } from "cdktn";
+import { Annotations, IResolvable, Token, Lazy } from "cdktn";
 import { Construct } from "constructs";
-import { ICertificate, IOrigin, FunctionAssociation } from ".";
+import {
+  ICertificate,
+  IOrigin,
+  FunctionAssociation,
+  FunctionEventType,
+} from ".";
+// aliased to avoid shadowing the global `Function` constructor
+import { Function as CloudFrontFunction } from "./function";
 import { Duration } from "../../duration";
 import { ArnFormat } from "../arn";
 import {
@@ -243,6 +250,7 @@ export class Distribution extends AwsConstructBase implements IDistribution {
 
   private readonly errorResponses: ErrorResponse[];
   private readonly certificate?: ICertificate;
+  private readonly warnedUnpublishedFunctions = new Set<string>();
 
   constructor(scope: Construct, name: string, props: DistributionProps) {
     super(scope, name, props);
@@ -297,11 +305,30 @@ export class Distribution extends AwsConstructBase implements IDistribution {
               ),
             ),
         }),
-        defaultCacheBehavior: this._renderDefaultCacheBehavior({
-          pathPattern: "*", // ignored for Default Cache Behavior
-          targetOriginId: defaultOriginId,
-          ...props.defaultBehavior,
-        }),
+        defaultCacheBehavior: {
+          ...this._renderDefaultCacheBehavior({
+            pathPattern: "*", // ignored for Default Cache Behavior
+            targetOriginId: defaultOriginId,
+            ...props.defaultBehavior,
+            // rendered lazily below so associations pushed onto a caller-held
+            // array *after* construction are still picked up at synth time
+            functionAssociations: undefined,
+          }),
+          functionAssociation: Lazy.anyValue(
+            {
+              produce: () =>
+                this.renderFunctionAssociations(
+                  props.defaultBehavior.functionAssociations,
+                )?.map((fa) =>
+                  // Lazy producers need additional xxxToTerraform wrap
+                  cloudfrontDistribution.cloudfrontDistributionDefaultCacheBehaviorFunctionAssociationToTerraform(
+                    fa,
+                  ),
+                ),
+            },
+            { omitEmptyArray: true },
+          ),
+        },
         orderedCacheBehavior: Lazy.anyValue(
           {
             produce: () =>
@@ -480,9 +507,67 @@ export class Distribution extends AwsConstructBase implements IDistribution {
       smoothStreaming: props.smoothStreaming,
       viewerProtocolPolicy:
         props.viewerProtocolPolicy ?? ViewerProtocolPolicy.ALLOW_ALL,
+      functionAssociation: this.renderFunctionAssociations(
+        props.functionAssociations,
+      ),
     };
   }
 
+  /**
+   * Renders the `functionAssociation` blocks for a cache behavior from the
+   * given `FunctionAssociation`s.
+   *
+   * CloudFront allows at most one function association per `FunctionEventType`
+   * for each cache behavior.
+   *
+   * @internal
+   */
+  private renderFunctionAssociations(
+    functionAssociations?: FunctionAssociation[],
+  ):
+    | cloudfrontDistribution.CloudfrontDistributionDefaultCacheBehaviorFunctionAssociation[]
+    | undefined {
+    if (!functionAssociations || functionAssociations.length === 0) {
+      return undefined;
+    }
+    const eventTypes = new Set<FunctionEventType>();
+    for (const fa of functionAssociations) {
+      if (eventTypes.has(fa.eventType)) {
+        throw new Error(
+          `Only one function association is allowed per event type, got multiple for event type ${fa.eventType}`,
+        );
+      }
+      eventTypes.add(fa.eventType);
+      // Only locally-created `Function`s are verifiable here - imported/general
+      // `IFunction` implementations may or may not be published, so leave them
+      // alone. CloudFront only allows LIVE-stage (published) functions to be
+      // associated with a distribution's cache behaviors.
+      // Lazy producers resolve more than once per synth (prepareStack +
+      // final render), so dedupe to avoid stacking identical warnings on
+      // this node's metadata.
+      if (
+        CloudFrontFunction.isFunction(fa.function) &&
+        !fa.function._autoPublish &&
+        !fa.skipPublishCheck &&
+        !this.warnedUnpublishedFunctions.has(fa.function.node.path)
+      ) {
+        this.warnedUnpublishedFunctions.add(fa.function.node.path);
+        // TODO(https://github.com/TerraConstructs/base/issues/161): switch to
+        // Annotations.addWarningV2()/acknowledgeWarning() once the Annotations
+        // facade lands, using the id prefix below as the warning's stable id.
+        Annotations.of(this).addWarning(
+          `[terraconstructs/aws-edge:unpublishedFunctionAssociation] Function '${fa.function.node.path}' is associated with a cache behavior but was created with autoPublish: false; ` +
+            "CloudFront only allows LIVE-stage functions in cache behaviors, so this will fail at apply time unless the function is published out of band. " +
+            "Set skipPublishCheck: true on the association to acknowledge.",
+        );
+      }
+    }
+    return functionAssociations.map((fa) => ({
+      eventType: fa.eventType,
+      functionArn: fa.function.functionArn,
+    }));
+  }
+
   private renderRestrictions(geoRestriction?: GeoRestriction) {
     return geoRestriction
       ? {