From b059ec379cece3f012c52061905e1c529182d4ab Mon Sep 17 00:00:00 2001
From: Frazer Smith <frazer.dev@icloud.com>
Date: Wed, 28 Aug 2024 10:24:43 +0100
Subject: [PATCH] fix(server): cors support for 404 responses (#1912)

---
 src/server.js      | 27 ++++++++++++++++++---------
 src/server.test.js | 14 +++++++++++---
 2 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/src/server.js b/src/server.js
index 4588b6cd..b6645f1c 100644
--- a/src/server.js
+++ b/src/server.js
@@ -8,6 +8,7 @@ const { joinSafe } = require("upath");
 const accepts = require("@fastify/accepts");
 const bearer = require("@fastify/bearer-auth");
 const compress = require("@fastify/compress");
+const cors = require("@fastify/cors");
 const helmet = require("@fastify/helmet");
 const disableCache = require("fastify-disablecache");
 const flocOff = require("fastify-floc-off");
@@ -191,15 +192,23 @@ async function plugin(server, config) {
 				});
 		})
 
-		// Rate limit 404 responses
-		.setNotFoundHandler(
-			{
-				preHandler: server.rateLimit(),
-			},
-			(req, res) => {
-				res.notFound(`Route ${req.method}:${req.url} not found`);
-			}
-		)
+		/**
+		 * Encapsulate the 404 handler into a child context, so that CORS
+		 * headers can be set explicitly for 404 responses.
+		 */
+		.register(async (notFoundContext) => {
+			await notFoundContext.register(cors, config.cors);
+
+			notFoundContext.setNotFoundHandler(
+				{
+					// Rate limit 404 responses to prevent URL enumeration
+					preHandler: server.rateLimit(),
+				},
+				(req, res) => {
+					res.notFound(`Route ${req.method}:${req.url} not found`);
+				}
+			);
+		})
 
 		// Errors thrown by routes and plugins are caught here
 		// eslint-disable-next-line promise/prefer-await-to-callbacks -- False positive
diff --git a/src/server.test.js b/src/server.test.js
index 154a6abc..010c33f6 100644
--- a/src/server.test.js
+++ b/src/server.test.js
@@ -904,7 +904,7 @@ describe("Server deployment", () => {
 							delete expResHeadersCors["content-type"];
 
 							/**
-							 * Vary header should not be set if CORS_ORIGIN is a a single domain
+							 * Vary header should not be set if CORS_ORIGIN is a single domain
 							 * or wildcard.
 							 * @see {@link https://github.com/fastify/fastify-cors/issues/287}
 							 */
@@ -971,7 +971,9 @@ describe("Server deployment", () => {
 							statusCode: 404,
 						});
 						expect(response.headers).toStrictEqual(
-							expResHeaders404Errors
+							envVariables.CORS_ORIGIN
+								? expected.response.headers.json
+								: expResHeaders404Errors
 						);
 						expect(response.statusCode).toBe(404);
 					});
@@ -982,6 +984,7 @@ describe("Server deployment", () => {
 							url: "/invalid",
 							headers: {
 								accept: "application/xml",
+								origin: request.headers.origin,
 							},
 						});
 
@@ -989,7 +992,12 @@ describe("Server deployment", () => {
 							'<?xml version="1.0" encoding="UTF-8"?><response><statusCode>404</statusCode><error>Not Found</error><message>Route GET:/invalid not found</message></response>'
 						);
 						expect(response.headers).toStrictEqual(
-							expResHeaders404ErrorsXml
+							envVariables.CORS_ORIGIN
+								? {
+										...expected.response.headers.json,
+										...expResHeaders404ErrorsXml,
+									}
+								: expResHeaders404ErrorsXml
 						);
 						expect(response.statusCode).toBe(404);
 					});