From b63836c89dd0e8ccb5946d8ad93cef201407258a Mon Sep 17 00:00:00 2001
From: Bobby Wibowo <bobbywibowo98@gmail.com>
Date: Tue, 3 Nov 2020 23:53:56 +0700
Subject: [PATCH] Added Content-Type override when serving with node

Closes #274
---
 config.sample.js |  8 ++++++++
 lolisafe.js      | 31 +++++++++++++++++++++++++++++--
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/config.sample.js b/config.sample.js
index e3d2a11..152c50b 100644
--- a/config.sample.js
+++ b/config.sample.js
@@ -38,6 +38,14 @@ module.exports = {
   */
   setContentDisposition: false,
 
+  /*
+    If you serve files with node, you can optionally choose to
+    override Content-Type header for certain extension names.
+  */
+  overrideContentTypes: {
+    'text/plain': ['html', 'htm', 'shtml', 'xhtml']
+  },
+
   /*
     If you are serving your files with a different domain than your lolisafe homepage,
     then fill this option with your lolisafe homepage, otherwise any falsy value.
diff --git a/lolisafe.js b/lolisafe.js
index ea28af9..0e14661 100644
--- a/lolisafe.js
+++ b/lolisafe.js
@@ -72,6 +72,21 @@ let setHeaders = res => {
   res.set('Access-Control-Allow-Origin', '*')
 }
 
+const contentTypes = config.overrideContentTypes && Object.keys(config.overrideContentTypes)
+const overrideContentTypes = (res, path) => {
+  // Do only if accessing files from uploads' root directory (i.e. not thumbs, etc.)
+  const relpath = path.replace(paths.uploads, '')
+  if (relpath.indexOf('/', 1) === -1) {
+    const name = relpath.substring(1)
+    const extname = utils.extname(name).substring(1)
+    for (const contentType of contentTypes) {
+      if (config.overrideContentTypes[contentType].includes(extname)) {
+        res.set('Content-Type', contentType)
+      }
+    }
+  }
+}
+
 const initServeStaticUploads = (opts = {}) => {
   if (config.setContentDisposition) {
     opts.preSetHeaders = async (res, req, path, stat) => {
@@ -133,8 +148,12 @@ if (config.cacheControl) {
   // If serving uploads with node
   if (config.serveFilesWithNode) {
     initServeStaticUploads({
-      setHeaders: res => {
+      setHeaders: (res, path) => {
         res.set('Access-Control-Allow-Origin', '*')
+        // Override Content-Type if necessary
+        if (contentTypes && contentTypes.length) {
+          overrideContentTypes(res, path)
+        }
         // If using CDN, cache uploads in CDN as well
         // Use with cloudflare.purgeCache enabled in config file
         if (config.cacheControl !== 2) {
@@ -164,7 +183,15 @@ if (config.cacheControl) {
     next()
   })
 } else if (config.serveFilesWithNode) {
-  initServeStaticUploads()
+  initServeStaticUploads({
+    setHeaders: (res, path) => {
+      res.set('Access-Control-Allow-Origin', '*')
+      // Override Content-Type if necessary
+      if (contentTypes && contentTypes.length) {
+        overrideContentTypes(res, path)
+      }
+    }
+  })
 }
 
 // Static assets