From 6a34fcf7558508b3e736a58517e39ed0b21ddd80 Mon Sep 17 00:00:00 2001
From: Divyansh Singh <40380293+brc-dd@users.noreply.github.com>
Date: Mon, 11 Jul 2022 23:32:05 +0530
Subject: [PATCH] fix: serve 404 on vitepress serve if not found

---
 src/node/serve/serve.ts | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/src/node/serve/serve.ts b/src/node/serve/serve.ts
index f4e42f10..c099b3d8 100644
--- a/src/node/serve/serve.ts
+++ b/src/node/serve/serve.ts
@@ -1,3 +1,5 @@
+import fs from 'fs'
+import path from 'path'
 import sirv from 'sirv'
 import compression from 'compression'
 import polka from 'polka'
@@ -26,13 +28,21 @@ export async function serve(options: ServeOptions = {}) {
   const site = await resolveConfig(options.root, 'serve', 'production')
   const base = trimChar(options?.base ?? site?.site?.base ?? '', '/')
 
+  const notAnAsset = (pathname: string) => !pathname.includes('/assets/')
+  const notFound = fs.readFileSync(path.resolve(site.outDir, './404.html'))
+  const onNoMatch: polka.Options['onNoMatch'] = (req, res) => {
+    res.statusCode = 404
+    if (notAnAsset(req.path)) res.write(notFound.toString())
+    res.end()
+  }
+
   const compress = compression()
   const serve = sirv(site.outDir, {
     etag: true,
     maxAge: 31536000,
     immutable: true,
     setHeaders(res, pathname) {
-      if (!pathname.includes('/assets/')) {
+      if (notAnAsset(pathname)) {
         // force server validation for non-asset files since they
         // are not fingerprinted
         res.setHeader('cache-control', 'no-cache')
@@ -41,14 +51,14 @@ export async function serve(options: ServeOptions = {}) {
   })
 
   if (base) {
-    polka()
+    polka({ onNoMatch })
       .use(base, compress, serve)
       .listen(port, (err: any) => {
         if (err) throw err
         console.log(`Built site served at http://localhost:${port}/${base}/\n`)
       })
   } else {
-    polka()
+    polka({ onNoMatch })
       .use(compress, serve)
       .listen(port, (err: any) => {
         if (err) throw err