feat: added cache fragment header to response

Author: RomanBobrovskiy

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/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

@@ -5,6 +5,7 @@ import { RoutingLambdaEdge } from '../constructs/RoutingLambdaEdge'
 import { CloudFrontDistribution } from '../constructs/CloudFrontDistribution'
 import { CacheConfig } from '../../types'
 import { CheckExpirationLambdaEdge } from '../constructs/CheckExpirationLambdaEdge'
+import { ViewerResponseLambdaEdge } from '../constructs/ViewerResponseLambdaEdge'
 
 export interface NextCloudfrontStackProps extends StackProps {
   nodejs?: string
@@ -21,6 +22,7 @@ export interface NextCloudfrontStackProps extends StackProps {
 export class NextCloudfrontStack extends Stack {
   public readonly routingLambdaEdge: RoutingLambdaEdge
   public readonly checkExpLambdaEdge: CheckExpirationLambdaEdge
+  public readonly viewerResponseLambdaEdge: ViewerResponseLambdaEdge
   public readonly cloudfront: CloudFrontDistribution
 
   constructor(scope: Construct, id: string, props: NextCloudfrontStackProps) {
@@ -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
@@ -65,6 +72,7 @@ export class NextCloudfrontStack extends Stack {
       renderServerDomain,
       requestEdgeFunction: this.routingLambdaEdge.lambdaEdge,
       responseEdgeFunction: this.checkExpLambdaEdge.lambdaEdge,
+      viewerResponseEdgeFunction: this.viewerResponseLambdaEdge.lambdaEdge,
       cacheConfig,
       imageTTL
     })

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)
+}