Handle IncludeBody for lambda triggers correctly

Author: gehan

README.md

@@ -108,6 +108,9 @@ distribution:
             ttl: 10
             lambda@edge:
               viewer-request: arn:aws:lambda:us-east-1:123:function:myFunc:version # lambda ARN including version
+              response-request: # can also send object to not include body
+                arn: lambda-arn
+                includeBody: false
 ```
 
 #### Private S3 Content

__tests__/__snapshots__/lambda-at-edge.test.js.snap

@@ -50,17 +50,17 @@ Object {
               },
               Object {
                 "EventType": "origin-request",
-                "IncludeBody": true,
+                "IncludeBody": false,
                 "LambdaFunctionARN": "arn:aws:lambda:us-east-1:123:function:originRequestFunction",
               },
               Object {
                 "EventType": "origin-response",
-                "IncludeBody": true,
+                "IncludeBody": false,
                 "LambdaFunctionARN": "arn:aws:lambda:us-east-1:123:function:originResponseFunction",
               },
               Object {
                 "EventType": "viewer-response",
-                "IncludeBody": true,
+                "IncludeBody": false,
                 "LambdaFunctionARN": "arn:aws:lambda:us-east-1:123:function:viewerResponseFunction",
               },
             ],

__tests__/lambda-at-edge.test.js

@@ -24,8 +24,14 @@ describe('Input origin as a custom url', () => {
             '/some/path': {
               ttl: 10,
               'lambda@edge': {
-                'viewer-request': 'arn:aws:lambda:us-east-1:123:function:viewerRequestFunction',
-                'origin-request': 'arn:aws:lambda:us-east-1:123:function:originRequestFunction',
+                'viewer-request': {
+                  arn: 'arn:aws:lambda:us-east-1:123:function:viewerRequestFunction',
+                  includeBody: true
+                },
+                'origin-request': {
+                  arn: 'arn:aws:lambda:us-east-1:123:function:originRequestFunction',
+                  includeBody: false
+                },
                 'origin-response': 'arn:aws:lambda:us-east-1:123:function:originResponseFunction',
                 'viewer-response': 'arn:aws:lambda:us-east-1:123:function:viewerResponseFunction'
               }
@@ -48,17 +54,17 @@ describe('Input origin as a custom url', () => {
           {
             EventType: 'origin-request',
             LambdaFunctionARN: 'arn:aws:lambda:us-east-1:123:function:originRequestFunction',
-            IncludeBody: true
+            IncludeBody: false
           },
           {
             EventType: 'origin-response',
             LambdaFunctionARN: 'arn:aws:lambda:us-east-1:123:function:originResponseFunction',
-            IncludeBody: true
+            IncludeBody: false
           },
           {
             EventType: 'viewer-response',
             LambdaFunctionARN: 'arn:aws:lambda:us-east-1:123:function:viewerResponseFunction',
-            IncludeBody: true
+            IncludeBody: false
           }
         ]
       }
@@ -92,4 +98,31 @@ describe('Input origin as a custom url', () => {
       )
     }
   })
+
+  it('throws error when includeBody given for event types other than request', async () => {
+    expect.assertions(1)
+
+    try {
+      await component.default({
+        origins: [
+          {
+            url: 'https://exampleorigin.com',
+            pathPatterns: {
+              '/some/path': {
+                ttl: 10,
+                'lambda@edge': {
+                  'viewer-response': {
+                    arn: 'arn:aws:lambda:us-east-1:123:function:viewerRequestFunction',
+                    includeBody: true
+                  }
+                }
+              }
+            }
+          }
+        ]
+      })
+    } catch (err) {
+      expect(err.message).toEqual('"includeBody" not allowed for viewer-response lambda triggers.')
+    }
+  })
 })

lib/addLambdaAtEdgeToCacheBehavior.js

@@ -4,6 +4,27 @@ const validLambdaTriggers = [
   'origin-response',
   'viewer-response'
 ]
+const triggersAllowedBody = ['viewer-request', 'origin-request']
+
+const makeCacheItem = (eventType, lambdaConfig) => {
+  let arn, includeBody
+  if (typeof lambdaConfig === 'string') {
+    arn = lambdaConfig
+    includeBody = triggersAllowedBody.includes(eventType)
+  } else {
+    ;({ arn, includeBody } = lambdaConfig)
+    if (includeBody && !triggersAllowedBody.includes(eventType)) {
+      throw new Error(
+        `"includeBody" not allowed for ${eventType} lambda triggers.`
+      )
+    }
+  }
+  return {
+    EventType: eventType,
+    LambdaFunctionARN: arn,
+    IncludeBody: includeBody
+  }
+}
 
 // adds lambda@edge to cache behavior passed
 module.exports = (cacheBehavior, lambdaAtEdgeConfig = {}) => {
@@ -14,12 +35,10 @@ module.exports = (cacheBehavior, lambdaAtEdgeConfig = {}) => {
       )
     }
 
+    cacheBehavior.LambdaFunctionAssociations.Items.push(
+      makeCacheItem(eventType, lambdaAtEdgeConfig[eventType])
+    )
     cacheBehavior.LambdaFunctionAssociations.Quantity =
       cacheBehavior.LambdaFunctionAssociations.Quantity + 1
-    cacheBehavior.LambdaFunctionAssociations.Items.push({
-      EventType: eventType,
-      LambdaFunctionARN: lambdaAtEdgeConfig[eventType],
-      IncludeBody: true
-    })
   })
 }