From 6f1a18ab19b8e912182bf4d93af085304cfa20e2 Mon Sep 17 00:00:00 2001
From: Brett Bailey <bretto36@hotmail.com>
Date: Thu, 6 Feb 2020 15:10:52 +1000
Subject: [PATCH] Custom Headers - Can now pass custom headers when using the
 Default method. This allows you to specify max age/expires/content type etc

---
 source/image-handler/image-request.js         | 19 +++++++++++-
 source/image-handler/index.js                 | 14 +++++++--
 .../image-handler/test/test-image-request.js  | 29 ++++++++++++++++++-
 3 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/source/image-handler/image-request.js b/source/image-handler/image-request.js
index 405ba6313..816b1566e 100644
--- a/source/image-handler/image-request.js
+++ b/source/image-handler/image-request.js
@@ -26,7 +26,8 @@ class ImageRequest {
             this.bucket = this.parseImageBucket(event, this.requestType);
             this.key = this.parseImageKey(event, this.requestType);
             this.edits = this.parseImageEdits(event, this.requestType);
-            this.originalImage = await this.getOriginalImage(this.bucket, this.key)
+            this.originalImage = await this.getOriginalImage(this.bucket, this.key);
+            this.headers = this.parseImageHeaders(event, this.requestType);
             return Promise.resolve(this);
         } catch (err) {
             return Promise.reject(err);
@@ -124,6 +125,22 @@ class ImageRequest {
         }
     }
 
+    /**
+     * Parses the headers to be sent with the response
+     * @param {String} event - Lambda request body.
+     * @param {String} requestType - Image handler request type.
+     */
+    parseImageHeaders(event, requestType) {
+        if (requestType === "Default") {
+            const decoded = this.decodeRequest(event);
+            if (decoded.headers !== undefined) {
+                return decoded.headers;
+            }
+        }
+
+        return [];
+    }
+
     /**
      * Parses the name of the appropriate Amazon S3 key corresponding to the
      * original image.
diff --git a/source/image-handler/index.js b/source/image-handler/index.js
index f6a9cba8a..c52a08a95 100755
--- a/source/image-handler/index.js
+++ b/source/image-handler/index.js
@@ -22,9 +22,18 @@ exports.handler = async (event) => {
         const request = await imageRequest.setup(event);
         console.log(request);
         const processedRequest = await imageHandler.process(request);
+        const headers = getResponseHeaders();
+
+        if (request.headers) {
+            // Apply the custom headers overwriting any that may need overwriting
+            for (let i = 0; i < request.headers.length; i++) {
+                headers[request.headers[i].key] = request.headers[i].value;
+            }
+        }
+
         const response = {
             "statusCode": 200,
-            "headers" : getResponseHeaders(),
+            "headers" : headers,
             "body": processedRequest,
             "isBase64Encoded": true
         }
@@ -53,7 +62,8 @@ const getResponseHeaders = (isErr) => {
         "Access-Control-Allow-Headers": "Content-Type, Authorization",
         "Access-Control-Allow-Credentials": true,
         "Content-Type": "image"
-    }
+    };
+
     if (corsEnabled) {
         headers["Access-Control-Allow-Origin"] = process.env.CORS_ORIGIN;
     }
diff --git a/source/image-handler/test/test-image-request.js b/source/image-handler/test/test-image-request.js
index f748e8df8..185463121 100644
--- a/source/image-handler/test/test-image-request.js
+++ b/source/image-handler/test/test-image-request.js
@@ -45,6 +45,7 @@ describe('setup()', function() {
                 bucket: 'validBucket',
                 key: 'validKey',
                 edits: { grayscale: true },
+                headers: [],
                 originalImage: Buffer.from('SampleImageContent\n')
             }
             // Assert
@@ -78,6 +79,7 @@ describe('setup()', function() {
                 bucket: 'allowedBucket001',
                 key: 'test-image-001.jpg',
                 edits: { grayscale: true },
+                headers: [],
                 originalImage: Buffer.from('SampleImageContent\n')
             }
             // Assert
@@ -116,6 +118,7 @@ describe('setup()', function() {
                     grayscale: true,
                     rotate: 90
                 },
+                headers: [],
                 originalImage: Buffer.from('SampleImageContent\n')
             }
             // Assert
@@ -644,4 +647,28 @@ describe('getAllowedSourceBuckets()', function() {
             });
         });
     });
-})
\ No newline at end of file
+})
+
+// ----------------------------------------------------------------------------
+// parseImageHeaders()
+// ----------------------------------------------------------------------------
+describe('parseImageHeaders()', function() {
+    describe('001/defaultCustomHeaders', function () {
+        it(`Should pass if the proper result is returned for a sample base64-
+            encoded image request`, function () {
+            // Arrange
+            const event = {
+                path: '/eyJoZWFkZXJzIjpbeyJrZXkiOiAiQ2FjaGUtQ29udHJvbCIsICJ2YWx1ZSI6ICJtYXgtYWdlPTMxNTM2MDAwLHB1YmxpYyJ9XX0='
+            }
+            // Act
+            const imageRequest = new ImageRequest();
+            const result = imageRequest.parseImageHeaders(event, 'Default');
+            // Assert
+            const expectedResult = [{
+                key: "Cache-Control",
+                value: 'max-age=31536000,public'
+            }]
+            assert.deepEqual(result, expectedResult);
+        });
+    });
+});
\ No newline at end of file