From b2ad6ca9bcab72aae19970b836c591da105cecdf Mon Sep 17 00:00:00 2001
From: Tom Moor <tom.moor@gmail.com>
Date: Wed, 1 Nov 2023 23:52:18 -0400
Subject: [PATCH] Refactor to middleware, support old routes

---
 server/middlewares/shareDomains.ts | 27 ++++++++++++++++++++++
 server/routes/index.ts             | 36 +++++++++---------------------
 2 files changed, 37 insertions(+), 26 deletions(-)
 create mode 100644 server/middlewares/shareDomains.ts

diff --git a/server/middlewares/shareDomains.ts b/server/middlewares/shareDomains.ts
new file mode 100644
index 000000000..f846fcddd
--- /dev/null
+++ b/server/middlewares/shareDomains.ts
@@ -0,0 +1,27 @@
+import { Context, Next } from "koa";
+import { Op } from "sequelize";
+import { parseDomain } from "@shared/utils/domains";
+import env from "@server/env";
+import { Share } from "@server/models";
+
+export default function shareDomains() {
+  return async function shareDomainsMiddleware(ctx: Context, next: Next) {
+    const isCustomDomain = parseDomain(ctx.host).custom;
+    const isDevelopment = env.ENVIRONMENT === "development";
+
+    if (isDevelopment || (isCustomDomain && env.isCloudHosted)) {
+      const share = await Share.unscoped().findOne({
+        where: {
+          domain: ctx.hostname,
+          published: true,
+          revokedAt: {
+            [Op.is]: null,
+          },
+        },
+      });
+      ctx.state.rootShare = share;
+    }
+
+    return next();
+  };
+}
diff --git a/server/routes/index.ts b/server/routes/index.ts
index 5a4d87df3..4d41088a2 100644
--- a/server/routes/index.ts
+++ b/server/routes/index.ts
@@ -6,13 +6,12 @@ import compress from "koa-compress";
 import Router from "koa-router";
 import send from "koa-send";
 import userAgent, { UserAgentContext } from "koa-useragent";
-import { Op } from "sequelize";
 import { languages } from "@shared/i18n";
 import { IntegrationType } from "@shared/types";
-import { parseDomain } from "@shared/utils/domains";
 import env from "@server/env";
 import { NotFoundError } from "@server/errors";
-import { Integration, Share } from "@server/models";
+import shareDomains from "@server/middlewares/shareDomains";
+import { Integration } from "@server/models";
 import { opensearchResponse } from "@server/utils/opensearch";
 import { getTeamFromContext } from "@server/utils/passport";
 import { robotsResponse } from "@server/utils/robots";
@@ -125,12 +124,16 @@ router.get("/opensearch.xml", (ctx) => {
   ctx.body = opensearchResponse(ctx.request.URL.origin);
 });
 
-router.get("/s/:shareId", renderShare);
-router.get("/s/:shareId/doc/:documentSlug", renderShare);
-router.get("/s/:shareId/*", renderShare);
+router.get("/s/:shareId", shareDomains(), renderShare);
+router.get("/s/:shareId/doc/:documentSlug", shareDomains(), renderShare);
+router.get("/s/:shareId/*", shareDomains(), renderShare);
 
 // catch all for application
-router.get("*", async (ctx, next) => {
+router.get("*", shareDomains(), async (ctx, next) => {
+  if (ctx.state?.rootShare) {
+    return renderShare(ctx, next);
+  }
+
   const team = await getTeamFromContext(ctx);
 
   // Redirect all requests to custom domain if one is set
@@ -139,25 +142,6 @@ router.get("*", async (ctx, next) => {
     return;
   }
 
-  const isCustomDomain = parseDomain(ctx.host).custom;
-  const isDevelopment = env.ENVIRONMENT === "development";
-  if (!team && (isDevelopment || (isCustomDomain && env.isCloudHosted))) {
-    const share = await Share.unscoped().findOne({
-      where: {
-        domain: ctx.hostname,
-        published: true,
-        revokedAt: {
-          [Op.is]: null,
-        },
-      },
-    });
-
-    if (share) {
-      ctx.state.rootShare = share;
-      return renderShare(ctx, next);
-    }
-  }
-
   const analytics = team
     ? await Integration.findOne({
         where: {