From 9b603e1c10b7bf1dbed6d6ad176559bde6b73455 Mon Sep 17 00:00:00 2001
From: Divyansh Singh <40380293+brc-dd@users.noreply.github.com>
Date: Tue, 7 Jun 2022 20:11:23 +0530
Subject: [PATCH] feat: emit 404.html to dist

---
 src/client/app/router.ts |  2 +-
 src/node/build/build.ts  |  3 +-
 src/node/build/render.ts | 60 ++++++++++++++++++++++++++--------------
 3 files changed, 42 insertions(+), 23 deletions(-)

diff --git a/src/client/app/router.ts b/src/client/app/router.ts
index 2eca5b7b..b4304585 100644
--- a/src/client/app/router.ts
+++ b/src/client/app/router.ts
@@ -109,7 +109,7 @@ export function createRouter(
         }
       }
     } catch (err: any) {
-      if (!err.message.match(/fetch/)) {
+      if (!err.message.match(/fetch/) && !href.match(/404\.html/)) {
         console.error(err)
       }
 
diff --git a/src/node/build/build.ts b/src/node/build/build.ts
index 91c08742..6ccce84f 100644
--- a/src/node/build/build.ts
+++ b/src/node/build/build.ts
@@ -49,7 +49,8 @@ export async function build(
       // as JS object literal.
       const hashMapString = JSON.stringify(JSON.stringify(pageToHashMap))
 
-      for (const page of siteConfig.pages) {
+      const pages = ['404.md', ...siteConfig.pages]
+      for (const page of pages) {
         await renderPage(
           siteConfig,
           page,
diff --git a/src/node/build/render.ts b/src/node/build/render.ts
index 3345342a..84de2e02 100644
--- a/src/node/build/render.ts
+++ b/src/node/build/render.ts
@@ -5,7 +5,7 @@ import { pathToFileURL } from 'url'
 import escape from 'escape-html'
 import { normalizePath, transformWithEsbuild } from 'vite'
 import { RollupOutput, OutputChunk, OutputAsset } from 'rollup'
-import { HeadConfig, createTitle } from '../shared'
+import { HeadConfig, createTitle, PageData } from '../shared'
 import { slash } from '../utils/slash'
 import { SiteConfig, resolveSiteDataByRoute } from '../config'
 
@@ -52,28 +52,46 @@ export async function renderPage(
   const pageHash = pageToHashMap[pageName.toLowerCase()]
   const pageClientJsFileName = `assets/${pageName}.${pageHash}.lean.js`
 
-  // resolve page data so we can render head tags
-  const { __pageData } = await import(
-    pathToFileURL(path.join(config.tempDir, pageServerJsFileName)).toString()
-  )
-  const pageData = JSON.parse(__pageData)
+  let pageData: PageData
+  let hasCustom404 = true
+
+  try {
+    // resolve page data so we can render head tags
+    const { __pageData } = await import(
+      pathToFileURL(path.join(config.tempDir, pageServerJsFileName)).toString()
+    )
+    pageData = JSON.parse(__pageData)
+  } catch (e) {
+    if (page === '404.md') {
+      hasCustom404 = false
+      pageData = {
+        relativePath: '',
+        title: '404',
+        description: 'Not Found',
+        headers: [],
+        frontmatter: {},
+        lastUpdated: 0
+      }
+    } else throw e
+  }
 
-  let preloadLinks = config.mpa
-    ? appChunk
-      ? [appChunk.fileName]
+  let preloadLinks =
+    config.mpa || (!hasCustom404 && page === '404.md')
+      ? appChunk
+        ? [appChunk.fileName]
+        : []
+      : result && appChunk
+      ? [
+          ...new Set([
+            // resolve imports for index.js + page.md.js and inject script tags for
+            // them as well so we fetch everything as early as possible without having
+            // to wait for entry chunks to parse
+            ...resolvePageImports(config, page, result, appChunk),
+            pageClientJsFileName,
+            appChunk.fileName
+          ])
+        ]
       : []
-    : result && appChunk
-    ? [
-        ...new Set([
-          // resolve imports for index.js + page.md.js and inject script tags for
-          // them as well so we fetch everything as early as possible without having
-          // to wait for entry chunks to parse
-          ...resolvePageImports(config, page, result, appChunk),
-          pageClientJsFileName,
-          appChunk.fileName
-        ])
-      ]
-    : []
 
   let prefetchLinks: string[] = []