From 87f42949b13839121213ea0442e1b4cf2d7d3bef Mon Sep 17 00:00:00 2001
From: tyhopp <hopp.ty.c@gmail.com>
Date: Fri, 13 May 2022 16:12:08 +0800
Subject: [PATCH] feat(gatsby): Refactor partytown local proxy logic

---
 packages/gatsby/src/commands/serve.ts         | 16 ++++---------
 .../partytown/gatsby-browser.tsx              | 23 +++++++++++--------
 .../internal-plugins/partytown/gatsby-node.ts | 15 +++---------
 .../src/internal-plugins/partytown/proxy.ts   | 16 +++++++++++++
 4 files changed, 38 insertions(+), 32 deletions(-)
 create mode 100644 packages/gatsby/src/internal-plugins/partytown/proxy.ts

diff --git a/packages/gatsby/src/commands/serve.ts b/packages/gatsby/src/commands/serve.ts
index 2f65a4ca804ee..1913c07e6f01b 100644
--- a/packages/gatsby/src/commands/serve.ts
+++ b/packages/gatsby/src/commands/serve.ts
@@ -25,7 +25,10 @@ import { initTracer } from "../utils/tracer"
 import { configureTrailingSlash } from "../utils/express-middlewares"
 import { getDataStore, detectLmdbStore } from "../datastore"
 import { functionMiddlewares } from "../internal-plugins/functions/middleware"
-import proxy from "express-http-proxy"
+import {
+  partytownProxyPath,
+  partytownProxy,
+} from "../internal-plugins/partytown/proxy"
 
 process.env.GATSBY_EXPERIMENTAL_LMDB_STORE = `1`
 detectLmdbStore()
@@ -124,16 +127,7 @@ module.exports = async (program: IServeProgram): Promise<void> => {
   // Proxy gatsby-script using off-main-thread strategy
   const { partytownProxiedURLs = [] } = config || {}
 
-  app.use(
-    `/__partytown-proxy`,
-    proxy(req => new URL(req.query.url as string).host as string, {
-      filter: req => partytownProxiedURLs.some(url => req.query?.url === url),
-      proxyReqPathResolver: req => {
-        const { pathname = ``, search = `` } = new URL(req.query?.url as string)
-        return pathname + search
-      },
-    })
-  )
+  app.use(partytownProxyPath, partytownProxy(partytownProxiedURLs))
 
   // eslint-disable-next-line new-cap
   const router = express.Router()
diff --git a/packages/gatsby/src/internal-plugins/partytown/gatsby-browser.tsx b/packages/gatsby/src/internal-plugins/partytown/gatsby-browser.tsx
index 0e947c441bda0..6529ee5d76dab 100644
--- a/packages/gatsby/src/internal-plugins/partytown/gatsby-browser.tsx
+++ b/packages/gatsby/src/internal-plugins/partytown/gatsby-browser.tsx
@@ -1,16 +1,21 @@
-import React, { useState } from "react"
+import React, { ReactElement, useState } from "react"
 import type { GatsbySSR } from "gatsby"
 import { Partytown } from "@builder.io/partytown/react"
 import { PartytownContext } from "gatsby-script"
 import type { PartytownProps } from "@builder.io/partytown/react"
 
-// TODO: type return and props
-// eslint-disable-next-line
-function PartytownProvider({ children }) {
-  const [{ collectedForwards, collectedAnyScript }, setState] = useState<{
-    collectedForwards: Set<string>
-    collectedAnyScript: boolean
-  }>({ collectedForwards: new Set(), collectedAnyScript: false })
+interface ICollectedForwardsState {
+  collectedForwards: Set<string>
+  collectedAnyScript: boolean
+}
+
+function PartytownProvider({ children }): ReactElement {
+  const [{ collectedForwards, collectedAnyScript }, setState] =
+    useState<ICollectedForwardsState>({
+      collectedForwards: new Set(),
+      collectedAnyScript: false,
+    })
+
   return (
     <PartytownContext.Provider
       value={{
@@ -26,7 +31,7 @@ function PartytownProvider({ children }) {
             stateShouldChange = true
           }
 
-          if (newScript.forward) {
+          if (newScript?.forward) {
             if (Array.isArray(newScript.forward)) {
               for (const singleForward of newScript.forward) {
                 if (!potentialNewState.collectedForwards.has(singleForward)) {
diff --git a/packages/gatsby/src/internal-plugins/partytown/gatsby-node.ts b/packages/gatsby/src/internal-plugins/partytown/gatsby-node.ts
index d72d9ec7a6e45..f8d81cbf3c6b7 100644
--- a/packages/gatsby/src/internal-plugins/partytown/gatsby-node.ts
+++ b/packages/gatsby/src/internal-plugins/partytown/gatsby-node.ts
@@ -1,7 +1,7 @@
 import path from "path"
 import { copyLibFiles } from "@builder.io/partytown/utils"
-import proxy from "express-http-proxy"
 import { CreateDevServerArgs } from "gatsby"
+import { partytownProxyPath, partytownProxy } from "./proxy"
 
 /**
  * Copy Partytown library files to public.
@@ -26,7 +26,7 @@ exports.createPages = ({ actions, store }): void => {
     const encodedURL: string = encodeURI(host)
 
     createRedirect({
-      fromPath: `/__partytown-proxy?url=${encodedURL}`,
+      fromPath: `${partytownProxyPath}?url=${encodedURL}`,
       toPath: encodedURL,
       statusCode: 200,
     })
@@ -40,14 +40,5 @@ export async function onCreateDevServer({
   const { config } = store.getState()
   const { partytownProxiedURLs = [] } = config || {}
 
-  app.use(
-    `/__partytown-proxy`,
-    proxy(req => new URL(req.query.url as string).host as string, {
-      filter: req => partytownProxiedURLs.some(url => req.query?.url === url),
-      proxyReqPathResolver: req => {
-        const { pathname = ``, search = `` } = new URL(req.query?.url as string)
-        return pathname + search
-      },
-    })
-  )
+  app.use(partytownProxyPath, partytownProxy(partytownProxiedURLs))
 }
diff --git a/packages/gatsby/src/internal-plugins/partytown/proxy.ts b/packages/gatsby/src/internal-plugins/partytown/proxy.ts
new file mode 100644
index 0000000000000..c9152de1fae9a
--- /dev/null
+++ b/packages/gatsby/src/internal-plugins/partytown/proxy.ts
@@ -0,0 +1,16 @@
+import proxy from "express-http-proxy"
+import type { RequestHandler } from "express"
+
+export const partytownProxyPath = `/__partytown-proxy`
+
+export function partytownProxy(
+  partytownProxiedURLs: Array<string>
+): RequestHandler {
+  return proxy(req => new URL(req.query.url as string).host as string, {
+    filter: req => partytownProxiedURLs.some(url => req.query?.url === url),
+    proxyReqPathResolver: req => {
+      const { pathname = ``, search = `` } = new URL(req.query?.url as string)
+      return pathname + search
+    },
+  })
+}