withastro · matthewp · May 1, 2023 · Apr 28, 2023 · Apr 28, 2023 · Apr 28, 2023
diff --git a/packages/astro/src/vite-plugin-astro-server/route.ts b/packages/astro/src/vite-plugin-astro-server/route.ts
@@ -1,6 +1,6 @@
 import type http from 'http';
 import mime from 'mime';
-import type { AstroSettings, ComponentInstance, ManifestData, RouteData } from '../@types/astro';
+import type { ComponentInstance, ManifestData, RouteData } from '../@types/astro';
 import type {
 	ComponentPreload,
 	DevelopmentEnvironment,
@@ -12,12 +12,10 @@ import { call as callEndpoint } from '../core/endpoint/dev/index.js';
 import { throwIfRedirectNotAllowed } from '../core/endpoint/index.js';
 import { AstroErrorData } from '../core/errors/index.js';
 import { warn } from '../core/logger/core.js';
-import { appendForwardSlash } from '../core/path.js';
 import { preload, renderPage } from '../core/render/dev/index.js';
 import { getParamsAndProps, GetParamsAndPropsError } from '../core/render/index.js';
 import { createRequest } from '../core/request.js';
 import { matchAllRoutes } from '../core/routing/index.js';
-import { resolvePages } from '../core/util.js';
 import { log404 } from './common.js';
 import { handle404Response, writeSSRResult, writeWebResponse } from './response.js';
 
@@ -35,11 +33,9 @@ interface MatchedRoute {
 	mod: ComponentInstance;
 }
 
-function getCustom404Route({ config }: AstroSettings, manifest: ManifestData) {
-	// For Windows compat, use relative page paths to match the 404 route
-	const relPages = resolvePages(config).href.replace(config.root.href, '');
-	const pattern = new RegExp(`${appendForwardSlash(relPages)}404.(astro|md)`);
-	return manifest.routes.find((r) => r.component.match(pattern));
+function getCustom404Route(manifest: ManifestData): RouteData | undefined {
+	const route404 = /^\/404\/?$/;
+	return manifest.routes.find((r) => route404.test(r.route));
 }
 
 export async function matchRoute(
@@ -97,7 +93,7 @@ export async function matchRoute(
 	}
 
 	log404(logging, pathname);
-	const custom404 = getCustom404Route(settings, manifest);
+	const custom404 = getCustom404Route(manifest);
 
 	if (custom404) {
 		const filePath = new URL(`./${custom404.component}`, settings.config.root);

diff --git a/packages/astro/test/custom-404-injected.test.js b/packages/astro/test/custom-404-injected.test.js
@@ -0,0 +1,42 @@
+import { expect } from 'chai';
+import * as cheerio from 'cheerio';
+import { loadFixture } from './test-utils.js';
+
+describe('Custom 404 with injectRoute', () => {
+	let fixture;
+
+	before(async () => {
+		fixture = await loadFixture({
+			root: './fixtures/custom-404-injected/',
+			site: 'http://example.com',
+		});
+	});
+
+	describe('dev', () => {
+		let devServer;
+		let $;
+
+		before(async () => {
+			devServer = await fixture.startDevServer();
+		});
+
+		after(async () => {
+			await devServer.stop();
+		});
+
+		it('renders /', async () => {
+			const html = await fixture.fetch('/').then((res) => res.text());
+			$ = cheerio.load(html);
+
+			expect($('h1').text()).to.equal('Home');
+		});
+
+		it('renders 404 for /a', async () => {
+			const html = await fixture.fetch('/a').then((res) => res.text());
+			$ = cheerio.load(html);
+
+			expect($('h1').text()).to.equal('Page not found');
+			expect($('p').text()).to.equal('/a');
+		});
+	});
+});
diff --git a/packages/astro/test/fixtures/custom-404-injected/astro.config.mjs b/packages/astro/test/fixtures/custom-404-injected/astro.config.mjs
@@ -0,0 +1,18 @@
+import { defineConfig } from 'astro/config';
+
+// https://astro.build/config
+export default defineConfig({
+  integrations: [
+    {
+      name: '404-integration',
+      hooks: {
+        'astro:config:setup': ({ injectRoute }) => {
+          injectRoute({
+            pattern: '404',
+            entryPoint: '/src/404.astro',
+          });
+        },
+      },
+    },
+  ],
+});
diff --git a/packages/astro/test/fixtures/custom-404-injected/package.json b/packages/astro/test/fixtures/custom-404-injected/package.json
@@ -0,0 +1,8 @@
+{
+  "name": "@test/custom-404-injected",
+  "version": "0.0.0",
+  "private": true,
+  "dependencies": {
+    "astro": "workspace:*"
+  }
+}
diff --git a/packages/astro/test/fixtures/custom-404-injected/src/404.astro b/packages/astro/test/fixtures/custom-404-injected/src/404.astro
@@ -0,0 +1,13 @@
+---
+const canonicalURL = new URL(Astro.url.pathname, Astro.site);
+---
+
+<html lang="en">
+<head>
+  <title>Not Found - Custom 404</title>
+</head>
+<body>
+  <h1>Page not found</h1>
+	<p>{canonicalURL.pathname}</p>
+</body>
+</html>
diff --git a/packages/astro/test/fixtures/custom-404-injected/src/pages/index.astro b/packages/astro/test/fixtures/custom-404-injected/src/pages/index.astro
@@ -0,0 +1,11 @@
+---
+---
+
+<html lang="en">
+<head>
+  <title>Custom 404</title>
+</head>
+<body>
+  <h1>Home</h1>
+</body>
+</html>
diff --git a/packages/astro/test/units/dev/dev.test.js b/packages/astro/test/units/dev/dev.test.js
@@ -158,6 +158,68 @@ describe('dev container', () => {
 		);
 	});
 
+	it('Serves injected 404 route for any 404', async () => {
+		const fs = createFs(
+			{
+				'/src/components/404.astro': `<h1>Custom 404</h1>`,
+				'/src/pages/page.astro': `<h1>Regular page</h1>`,
+			},
+			root
+		);
+
+		await runInContainer(
+			{
+				fs,
+				root,
+				userConfig: {
+					output: 'server',
+					integrations: [
+						{
+							name: '@astrojs/test-integration',
+							hooks: {
+								'astro:config:setup': ({ injectRoute }) => {
+									injectRoute({
+										pattern: '/404',
+										entryPoint: './src/components/404.astro',
+									});
+								},
+							},
+						},
+					],
+				},
+			},
+			async (container) => {
+				{
+					// Regular pages are served as expected.
+					const r = createRequestAndResponse({ method: 'GET', url: '/page' });
+					container.handle(r.req, r.res);
+					await r.done;
+					const doc = await r.text();
+					expect(doc).to.match(/<h1>Regular page<\/h1>/);
+					expect(r.res.statusCode).to.equal(200);
+				}
+				{
+					// `/404` serves the custom 404 page as expected.
+					const r = createRequestAndResponse({ method: 'GET', url: '/404' });
+					container.handle(r.req, r.res);
+					await r.done;
+					const doc = await r.text();
+					expect(doc).to.match(/<h1>Custom 404<\/h1>/);
+					expect(r.res.statusCode).to.equal(200);
+				}
+				{
+					// A non-existent page also serves the custom 404 page.
+					const r = createRequestAndResponse({ method: 'GET', url: '/other-page' });
+					container.handle(r.req, r.res);
+					await r.done;
+					const doc = await r.text();
+					expect(doc).to.match(/<h1>Custom 404<\/h1>/);
+					expect(r.res.statusCode).to.equal(200);
+				}
+			}
+		);
+	});
+
 	it('items in public/ are not available from root when using a base', async () => {
 		await runInContainer(
 			{

diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml