Bluehost Screenshot service updates

functions.js

@@ -3,6 +3,7 @@ const {get} = require('axios');
 const crypto = require('crypto');
 const {readFileSync} = require('fs');
 const {Readable} = require('stream');
+const streamToArray = require('stream-to-array');
 
 const lambda = new Lambda();
 const s3 = new S3();
@@ -42,6 +43,18 @@ function bufferToStream(buffer) {
     });
 }
 
+/**
+ * Convert a stream to a buffer.
+ *
+ * @param {Readable} stream
+ * @returns {Promise<Buffer>}
+ */
+async function streamToBuffer(stream) {
+    const parts = await streamToArray(stream);
+    const buffers = parts.map(part => Buffer.isBuffer(part) ? part : Buffer.from(part));
+    return Buffer.concat(buffers);
+}
+
 /**
  * Get an MD5 hash of a string.
  *
@@ -125,6 +138,24 @@ async function captureScreenshot(url) {
     });
 }
 
+/**
+ * Asynchronously generate a screenshot given a URL.
+ *
+ * @param {string} bucket
+ * @param {string} key
+ * @param {string} url
+ * @returns {string}
+ */
+function queueScreenshotGeneration(bucket, key, url) {
+    const params = {
+        InvocationType: 'Event',
+        FunctionName: 'bluehost-url-to-screenshot-on-s3',
+        Payload: JSON.stringify({bucket, key, url}, null, 2),
+    };
+
+    lambda.invoke(params).send();
+}
+
 /**
  * Send an image as a response.
  *
@@ -169,6 +200,17 @@ function sendError(message = 'Unable to generate screenshot') {
     };
 }
 
+function sendRedirect(location, statusCode = 301) {
+    return {
+        statusCode: 301,
+        headers: {
+            Location: location,
+        },
+        body: '',
+        isBase64Encoded: false,
+    };
+}
+
 /**
  * Send a 404 response.
  *
@@ -180,17 +222,24 @@ function send404() {
     };
 }
 
+function getS3Url(bucket, file) {
+    return `http://${bucket}.s3.amazonaws.com/${file}`;
+}
+
 module.exports = {
     base64Decode,
     base64Encode,
     bufferToStream,
     captureScreenshot,
     fetchFromS3,
     fileExistsOnS3,
+    getS3Url,
     md5,
+    queueScreenshotGeneration,
     send404,
     sendError,
     sendImage,
     sendPlaceholderImage,
+    sendRedirect,
     uploadToS3,
 };

package.json

@@ -12,7 +12,8 @@
   "license": "GPL-2.0-or-later",
   "dependencies": {
     "aws-sdk": "^2.547.0",
-    "axios": "^0.19.0"
+    "axios": "^0.19.0",
+    "stream-to-array": "^2.3.0"
   },
   "devDependencies": {
     "husky": "^4.2.5",

screenshot.js

@@ -1,15 +1,9 @@
 const {
-    base64Decode,
-    base64Encode,
-    bufferToStream,
-    captureScreenshot,
-    fetchFromS3,
     fileExistsOnS3,
+    getS3Url,
     md5,
-    send404,
-    sendImage,
-    sendPlaceholderImage,
-    uploadToS3
+    queueScreenshotGeneration,
+    sendRedirect,
 } = require('./functions');
 
 exports.Screenshot = class {
@@ -37,21 +31,12 @@ exports.Screenshot = class {
         }
     }
 
-    async persist(base64) {
-        return await uploadToS3(this.bucket, this.filepath(), bufferToStream(base64Decode(base64)));
-    }
-
     async fetch() {
-        try {
-            const response = await fetchFromS3(this.bucket, this.filepath());
-            return sendImage(base64Encode(response.Body));
-        } catch (e) {
-            try {
-                return sendImage(await this.persist(await captureScreenshot(this.url)));
-            } catch (e) {
-                return sendPlaceholderImage();
-            }
+        if (!await this.hasScreenshot()) {
+            queueScreenshotGeneration(this.bucket, this.filepath(), this.url);
+            return sendRedirect(getS3Url(this.bucket, 'placeholder.png'), 302);
         }
+        return sendRedirect(getS3Url(this.bucket, this.filepath()));
     }
 
 };