feat/platform-1628: updated lambda functions

src/cacheHandler/strategy/s3.ts

@@ -66,6 +66,9 @@ export class S3Cache implements CacheStrategy {
     const baseInput: PutObjectCommandInput = {
       Bucket: this.bucketName,
       Key: `${pageKey}/${cacheKey}`,
+      Metadata: {
+        'Cache-Fragment-Key': cacheKey
+      },
       ...(data.revalidate ? { CacheControl: `max-age=${data.revalidate}` } : undefined)
     }
 

src/cdk/constructs/CloudFrontDistribution.ts

@@ -12,6 +12,7 @@ interface CloudFrontPropsDistribution {
   renderServerDomain: string
   requestEdgeFunction: cloudfront.experimental.EdgeFunction
   responseEdgeFunction: cloudfront.experimental.EdgeFunction
+  viewerResponseEdgeFunction: cloudfront.experimental.EdgeFunction
   cacheConfig: CacheConfig
   imageTTL?: number
 }
@@ -29,7 +30,15 @@ export class CloudFrontDistribution extends Construct {
   constructor(scope: Construct, id: string, props: CloudFrontPropsDistribution) {
     super(scope, id)
 
-    const { staticBucket, requestEdgeFunction, responseEdgeFunction, cacheConfig, renderServerDomain, imageTTL } = props
+    const {
+      staticBucket,
+      requestEdgeFunction,
+      responseEdgeFunction,
+      viewerResponseEdgeFunction,
+      cacheConfig,
+      renderServerDomain,
+      imageTTL
+    } = props
 
     const splitCachePolicy = new cloudfront.CachePolicy(this, 'SplitCachePolicy', {
       cachePolicyName: `${id}-SplitCachePolicy`,
@@ -86,6 +95,10 @@ export class CloudFrontDistribution extends Construct {
           {
             functionVersion: responseEdgeFunction.currentVersion,
             eventType: cloudfront.LambdaEdgeEventType.ORIGIN_RESPONSE
+          },
+          {
+            functionVersion: viewerResponseEdgeFunction.currentVersion,
+            eventType: cloudfront.LambdaEdgeEventType.VIEWER_RESPONSE
           }
         ],
         cachePolicy: splitCachePolicy

src/cdk/constructs/RoutingLambdaEdge.ts → src/cdk/constructs/OriginRequestLambdaEdge.ts

@@ -8,7 +8,7 @@ import path from 'node:path'
 import { buildLambda } from '../../build/edge'
 import { CacheConfig } from '../../types'
 
-interface RoutingLambdaEdgeProps extends cdk.StackProps {
+interface OriginRequestLambdaEdgeProps extends cdk.StackProps {
   bucketName: string
   renderServerDomain: string
   buildOutputPath: string
@@ -22,15 +22,15 @@ const NodeJSEnvironmentMapping: Record<string, lambda.Runtime> = {
   '20': lambda.Runtime.NODEJS_20_X
 }
 
-export class RoutingLambdaEdge extends Construct {
+export class OriginRequestLambdaEdge extends Construct {
   public readonly lambdaEdge: cloudfront.experimental.EdgeFunction
 
-  constructor(scope: Construct, id: string, props: RoutingLambdaEdgeProps) {
+  constructor(scope: Construct, id: string, props: OriginRequestLambdaEdgeProps) {
     const { bucketName, bucketRegion, renderServerDomain, nodejs, buildOutputPath, cacheConfig } = props
     super(scope, id)
 
     const nodeJSEnvironment = NodeJSEnvironmentMapping[nodejs ?? ''] ?? NodeJSEnvironmentMapping['20']
-    const name = 'edgeRouting'
+    const name = 'originRequest'
 
     buildLambda(name, buildOutputPath, {
       define: {
@@ -41,13 +41,13 @@ export class RoutingLambdaEdge extends Construct {
       }
     })
 
-    const logGroup = new logs.LogGroup(this, 'RoutingLambdaEdgeLogGroup', {
-      logGroupName: `/aws/lambda/${id}-edgeRouting`,
+    const logGroup = new logs.LogGroup(this, 'OriginRequestLambdaEdgeLogGroup', {
+      logGroupName: `/aws/lambda/${id}-originRequest`,
       removalPolicy: cdk.RemovalPolicy.DESTROY,
       retention: logs.RetentionDays.ONE_DAY
     })
 
-    this.lambdaEdge = new cloudfront.experimental.EdgeFunction(this, 'RoutingLambdaEdge', {
+    this.lambdaEdge = new cloudfront.experimental.EdgeFunction(this, 'OriginRequestLambdaEdge', {
       runtime: nodeJSEnvironment,
       code: lambda.Code.fromAsset(path.join(buildOutputPath, 'server-functions', name)),
       handler: 'index.handler',

src/cdk/constructs/CheckExpirationLambdaEdge.ts → src/cdk/constructs/OriginResponseLambdaEdge.ts

@@ -8,7 +8,7 @@ import path from 'node:path'
 import { buildLambda } from '../../build/edge'
 import { CacheConfig } from '../../types'
 
-interface CheckExpirationLambdaEdgeProps extends cdk.StackProps {
+interface OriginResponseLambdaEdgeProps extends cdk.StackProps {
   renderWorkerQueueUrl: string
   renderWorkerQueueArn: string
   buildOutputPath: string
@@ -22,15 +22,15 @@ const NodeJSEnvironmentMapping: Record<string, lambda.Runtime> = {
   '20': lambda.Runtime.NODEJS_20_X
 }
 
-export class CheckExpirationLambdaEdge extends Construct {
+export class OriginResponseLambdaEdge extends Construct {
   public readonly lambdaEdge: cloudfront.experimental.EdgeFunction
 
-  constructor(scope: Construct, id: string, props: CheckExpirationLambdaEdgeProps) {
+  constructor(scope: Construct, id: string, props: OriginResponseLambdaEdgeProps) {
     const { nodejs, buildOutputPath, cacheConfig, renderWorkerQueueUrl, renderWorkerQueueArn, region } = props
     super(scope, id)
 
     const nodeJSEnvironment = NodeJSEnvironmentMapping[nodejs ?? ''] ?? NodeJSEnvironmentMapping['20']
-    const name = 'checkExpiration'
+    const name = 'originResponse'
 
     buildLambda(name, buildOutputPath, {
       define: {
@@ -40,13 +40,13 @@ export class CheckExpirationLambdaEdge extends Construct {
       }
     })
 
-    const logGroup = new logs.LogGroup(this, 'CheckExpirationLambdaEdgeLogGroup', {
-      logGroupName: `/aws/lambda/${id}-checkExpiration`,
+    const logGroup = new logs.LogGroup(this, 'OriginResponseLambdaEdgeLogGroup', {
+      logGroupName: `/aws/lambda/${id}-originResponse`,
       removalPolicy: cdk.RemovalPolicy.DESTROY,
       retention: logs.RetentionDays.ONE_DAY
     })
 
-    this.lambdaEdge = new cloudfront.experimental.EdgeFunction(this, 'CheckExpirationLambdaEdge', {
+    this.lambdaEdge = new cloudfront.experimental.EdgeFunction(this, 'OriginResponseLambdaEdge', {
       runtime: nodeJSEnvironment,
       code: lambda.Code.fromAsset(path.join(buildOutputPath, 'server-functions', name)),
       handler: 'index.handler',

src/cdk/constructs/ViewerResponseLambdaEdge.ts

@@ -0,0 +1,35 @@
+import { Construct } from 'constructs'
+import * as lambda from 'aws-cdk-lib/aws-lambda'
+import * as cloudfront from 'aws-cdk-lib/aws-cloudfront'
+import path from 'node:path'
+import { buildLambda } from '../../build/edge'
+
+const NodeJSEnvironmentMapping: Record<string, lambda.Runtime> = {
+  '18': lambda.Runtime.NODEJS_18_X,
+  '20': lambda.Runtime.NODEJS_20_X
+}
+
+interface ViewerResponseLambdaEdgeProps {
+  nodejs?: string
+  buildOutputPath: string
+}
+
+export class ViewerResponseLambdaEdge extends Construct {
+  public readonly lambdaEdge: cloudfront.experimental.EdgeFunction
+
+  constructor(scope: Construct, id: string, props: ViewerResponseLambdaEdgeProps) {
+    const { nodejs, buildOutputPath } = props
+    super(scope, id)
+
+    const nodeJSEnvironment = NodeJSEnvironmentMapping[nodejs ?? ''] ?? NodeJSEnvironmentMapping['20']
+    const name = 'viewerResponse'
+
+    buildLambda(name, buildOutputPath)
+
+    this.lambdaEdge = new cloudfront.experimental.EdgeFunction(this, 'ViewerResponseLambdaEdge', {
+      runtime: nodeJSEnvironment,
+      code: lambda.Code.fromAsset(path.join(buildOutputPath, 'server-functions', name)),
+      handler: 'index.handler'
+    })
+  }
+}

src/cdk/stacks/NextCloudfrontStack.ts

@@ -1,10 +1,11 @@
 import { Stack, type StackProps } from 'aws-cdk-lib'
 import { Construct } from 'constructs'
 import * as s3 from 'aws-cdk-lib/aws-s3'
-import { RoutingLambdaEdge } from '../constructs/RoutingLambdaEdge'
+import { OriginRequestLambdaEdge } from '../constructs/OriginRequestLambdaEdge'
 import { CloudFrontDistribution } from '../constructs/CloudFrontDistribution'
 import { CacheConfig } from '../../types'
-import { CheckExpirationLambdaEdge } from '../constructs/CheckExpirationLambdaEdge'
+import { OriginResponseLambdaEdge } from '../constructs/OriginResponseLambdaEdge'
+import { ViewerResponseLambdaEdge } from '../constructs/ViewerResponseLambdaEdge'
 
 export interface NextCloudfrontStackProps extends StackProps {
   nodejs?: string
@@ -19,8 +20,9 @@ export interface NextCloudfrontStackProps extends StackProps {
 }
 
 export class NextCloudfrontStack extends Stack {
-  public readonly routingLambdaEdge: RoutingLambdaEdge
-  public readonly checkExpLambdaEdge: CheckExpirationLambdaEdge
+  public readonly originRequestLambdaEdge: OriginRequestLambdaEdge
+  public readonly originResponseLambdaEdge: OriginResponseLambdaEdge
+  public readonly viewerResponseLambdaEdge: ViewerResponseLambdaEdge
   public readonly cloudfront: CloudFrontDistribution
 
   constructor(scope: Construct, id: string, props: NextCloudfrontStackProps) {
@@ -37,7 +39,7 @@ export class NextCloudfrontStack extends Stack {
       imageTTL
     } = props
 
-    this.routingLambdaEdge = new RoutingLambdaEdge(this, `${id}-RoutingLambdaEdge`, {
+    this.originRequestLambdaEdge = new OriginRequestLambdaEdge(this, `${id}-OriginRequestLambdaEdge`, {
       nodejs,
       bucketName: staticBucketName,
       renderServerDomain,
@@ -46,7 +48,7 @@ export class NextCloudfrontStack extends Stack {
       bucketRegion: region
     })
 
-    this.checkExpLambdaEdge = new CheckExpirationLambdaEdge(this, `${id}-CheckExpirationLambdaEdge`, {
+    this.originResponseLambdaEdge = new OriginResponseLambdaEdge(this, `${id}-OriginResponseLambdaEdge`, {
       nodejs,
       renderWorkerQueueUrl,
       buildOutputPath,
@@ -55,6 +57,11 @@ export class NextCloudfrontStack extends Stack {
       region
     })
 
+    this.viewerResponseLambdaEdge = new ViewerResponseLambdaEdge(this, `${id}-ViewerResponseLambdaEdge`, {
+      nodejs,
+      buildOutputPath
+    })
+
     const staticBucket = s3.Bucket.fromBucketAttributes(this, `${id}-StaticAssetsBucket`, {
       bucketName: staticBucketName,
       region
@@ -63,13 +70,14 @@ export class NextCloudfrontStack extends Stack {
     this.cloudfront = new CloudFrontDistribution(this, `${id}-NextCloudFront`, {
       staticBucket,
       renderServerDomain,
-      requestEdgeFunction: this.routingLambdaEdge.lambdaEdge,
-      responseEdgeFunction: this.checkExpLambdaEdge.lambdaEdge,
+      requestEdgeFunction: this.originRequestLambdaEdge.lambdaEdge,
+      responseEdgeFunction: this.originResponseLambdaEdge.lambdaEdge,
+      viewerResponseEdgeFunction: this.viewerResponseLambdaEdge.lambdaEdge,
       cacheConfig,
       imageTTL
     })
 
-    staticBucket.grantRead(this.routingLambdaEdge.lambdaEdge)
-    staticBucket.grantRead(this.checkExpLambdaEdge.lambdaEdge)
+    staticBucket.grantRead(this.originRequestLambdaEdge.lambdaEdge)
+    staticBucket.grantRead(this.originResponseLambdaEdge.lambdaEdge)
   }
 }

src/lambdas/viewerResponse.ts

@@ -0,0 +1,28 @@
+import type { CloudFrontRequestCallback, Context, CloudFrontResponseEvent } from 'aws-lambda'
+
+/**
+ * Lambda@Edge viewer response handler that processes CloudFront responses.
+ * This handler extracts the cache fragment key from x-amz-meta headers and
+ * sets it as a standard Cache-Fragment-Key header while removing the original
+ * x-amz-meta header.
+ *
+ * @param {CloudFrontResponseEvent} event - The CloudFront response event object
+ * @param {Context} _context - AWS Lambda context object (unused)
+ * @param {CloudFrontRequestCallback} callback - Callback to return the modified response
+ * @returns {Promise<void>} - Returns nothing, uses callback to return response
+ */
+export const handler = async (
+  event: CloudFrontResponseEvent,
+  _context: Context,
+  callback: CloudFrontRequestCallback
+) => {
+  const response = event.Records[0].cf.response
+  const fileCacheKey = response.headers['x-amz-meta-cache-fragment-key']?.[0].value
+
+  if (fileCacheKey) {
+    response.headers['cache-fragment-key'] = [{ key: 'Cache-Fragment-Key', value: fileCacheKey }]
+    response.headers['x-amz-meta-cache-fragment-key'] = []
+  }
+
+  callback(null, response)
+}