From 7270d9524004ca58854ec7b0eccbb5e30e2af5de Mon Sep 17 00:00:00 2001
From: Mike Cao <mike@mikecao.com>
Date: Sun, 20 Feb 2022 23:27:50 -0800
Subject: [PATCH] Move environment variables to middleware. Closes #972.

---
 Dockerfile                           |  3 +-
 components/forms/TrackingCodeForm.js |  4 +--
 hooks/useForceSSL.js                 | 14 ---------
 next.config.js                       | 10 +------
 pages/_app.js                        |  2 --
 pages/_middleware.js                 | 44 ++++++++++++++++++++++++++++
 6 files changed, 47 insertions(+), 30 deletions(-)
 delete mode 100644 hooks/useForceSSL.js
 create mode 100644 pages/_middleware.js

diff --git a/Dockerfile b/Dockerfile
index ee931167..16dd9872 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -2,12 +2,11 @@
 FROM node:12.22-alpine AS build
 ARG BASE_PATH
 ARG DATABASE_TYPE
-ARG TRACKER_SCRIPT_NAME
 
 ENV BASE_PATH=$BASE_PATH
 ENV DATABASE_URL "postgresql://umami:umami@db:5432/umami"
 ENV DATABASE_TYPE=$DATABASE_TYPE
-ENV TRACKER_SCRIPT_NAME=$TRACKER_SCRIPT_NAME
+
 WORKDIR /build
 
 RUN yarn config set --home enableTelemetry 0
diff --git a/components/forms/TrackingCodeForm.js b/components/forms/TrackingCodeForm.js
index bfc6940d..fd679ab2 100644
--- a/components/forms/TrackingCodeForm.js
+++ b/components/forms/TrackingCodeForm.js
@@ -5,8 +5,6 @@ import Button from 'components/common/Button';
 import FormLayout, { FormButtons, FormRow } from 'components/layout/FormLayout';
 import CopyButton from 'components/common/CopyButton';
 
-const scriptName = process.env.TRACKER_SCRIPT_NAME || 'umami';
-
 export default function TrackingCodeForm({ values, onClose }) {
   const ref = useRef();
   const { basePath } = useRouter();
@@ -26,7 +24,7 @@ export default function TrackingCodeForm({ values, onClose }) {
           rows={3}
           cols={60}
           spellCheck={false}
-          defaultValue={`<script async defer data-website-id="${values.website_uuid}" src="${document.location.origin}${basePath}/${scriptName}.js"></script>`}
+          defaultValue={`<script async defer data-website-id="${values.website_uuid}" src="${document.location.origin}${basePath}/umami.js"></script>`}
           readOnly
         />
       </FormRow>
diff --git a/hooks/useForceSSL.js b/hooks/useForceSSL.js
deleted file mode 100644
index b9d95e19..00000000
--- a/hooks/useForceSSL.js
+++ /dev/null
@@ -1,14 +0,0 @@
-import { useEffect } from 'react';
-import { useRouter } from 'next/router';
-
-export default function useForceSSL(enabled) {
-  const router = useRouter();
-
-  useEffect(() => {
-    if (enabled && typeof window !== 'undefined' && /^http:\/\//.test(location.href)) {
-      router.push(location.href.replace(/^http:\/\//, 'https://'));
-    }
-  }, [enabled]);
-
-  return null;
-}
diff --git a/next.config.js b/next.config.js
index 71cb9b20..3ee8e699 100644
--- a/next.config.js
+++ b/next.config.js
@@ -1,14 +1,9 @@
 require('dotenv').config();
 const pkg = require('./package.json');
 
-const scriptName = process.env.TRACKER_SCRIPT_NAME;
-
 module.exports = {
   env: {
     VERSION: pkg.version,
-    FORCE_SSL: Boolean(process.env.FORCE_SSL),
-    DISABLE_LOGIN: Boolean(process.env.DISABLE_LOGIN),
-    TRACKER_SCRIPT_NAME: scriptName,
   },
   basePath: process.env.BASE_PATH,
   eslint: {
@@ -23,13 +18,10 @@ module.exports = {
 
     return config;
   },
-  async rewrites() {
-    return scriptName ? [{ source: `/${scriptName}.js`, destination: '/umami.js' }] : [];
-  },
   async headers() {
     return [
       {
-        source: `/${scriptName || 'umami'}.js`,
+        source: `/umami.js`,
         headers: [
           {
             key: 'Cache-Control',
diff --git a/pages/_app.js b/pages/_app.js
index a4c4776c..72364786 100644
--- a/pages/_app.js
+++ b/pages/_app.js
@@ -5,7 +5,6 @@ import { IntlProvider } from 'react-intl';
 import { Provider } from 'react-redux';
 import { useStore } from 'redux/store';
 import useLocale from 'hooks/useLocale';
-import useForceSSL from 'hooks/useForceSSL';
 import 'styles/variables.css';
 import 'styles/bootstrap-grid.css';
 import 'styles/index.css';
@@ -25,7 +24,6 @@ const Intl = ({ children }) => {
 };
 
 export default function App({ Component, pageProps }) {
-  useForceSSL(process.env.FORCE_SSL);
   const store = useStore();
   const { basePath } = useRouter();
 
diff --git a/pages/_middleware.js b/pages/_middleware.js
new file mode 100644
index 00000000..94d88360
--- /dev/null
+++ b/pages/_middleware.js
@@ -0,0 +1,44 @@
+import { NextResponse } from 'next/server';
+
+function redirectHTTPS(req) {
+  if (
+    process.env.FORCE_SSL &&
+    !req.headers.get('host').includes('localhost') &&
+    req.nextUrl.protocol !== 'https'
+  ) {
+    return NextResponse.redirect(`https://${req.headers.get('host')}${req.nextUrl.pathname}`, 301);
+  }
+}
+
+function customScriptName(req) {
+  const scriptName = process.env.TRACKER_SCRIPT_NAME;
+
+  if (scriptName) {
+    const url = req.nextUrl.clone();
+    const { pathname } = url;
+
+    if (pathname.endsWith(`/${scriptName}.js`)) {
+      url.pathname = '/umami.js';
+      return NextResponse.rewrite(url);
+    }
+  }
+}
+
+function disableLogin(req) {
+  if (process.env.DISABLE_LOGIN && req.nextUrl.pathname.endsWith('/login')) {
+    return new Response('403 Forbidden', { status: 403 });
+  }
+}
+
+export function middleware(req) {
+  const fns = [redirectHTTPS, customScriptName, disableLogin];
+
+  for (const fn of fns) {
+    const res = fn(req);
+    if (res) {
+      return res;
+    }
+  }
+
+  return NextResponse.next();
+}