From d794367ad3bb54a3c2c596847ef54be04fb547b0 Mon Sep 17 00:00:00 2001
From: KaKa <kaka@kakawebsitedemo.com>
Date: Fri, 1 Apr 2022 22:41:51 +0800
Subject: [PATCH] refactor: update 404 onBadUrl behavior

---
 lib/fourOhFour.js   | 33 ++++++++++++++++-----------------
 test/404s.test.js   | 16 ++++++++++++++++
 test/logger.test.js |  3 ++-
 3 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/lib/fourOhFour.js b/lib/fourOhFour.js
index 05dd4d886d..e6e828046c 100644
--- a/lib/fourOhFour.js
+++ b/lib/fourOhFour.js
@@ -30,10 +30,11 @@ const fourOhFourContext = {
  * kFourOhFourContext: the context in the reply object where the handler will be executed
  */
 function fourOhFour (options) {
-  const { logger, genReqId, disableRequestLogging } = options
+  const { logger, genReqId } = options
 
   // 404 router, used for handling encapsulated 404 handlers
-  const router = FindMyWay({ onBadUrl, defaultRoute: fourOhFourFallBack })
+  const router = FindMyWay({ onBadUrl: createOnBadUrl(), defaultRoute: fourOhFourFallBack })
+  let _onBadUrlHandler = null
 
   return { router, setNotFoundHandler, setContext, arrange404 }
 
@@ -41,6 +42,8 @@ function fourOhFour (options) {
     // Change the pointer of the fastify instance to itself, so register + prefix can add new 404 handler
     instance[kFourOhFourLevelInstance] = instance
     instance[kCanSetNotFoundHandler] = true
+    // we need to bind instance for the context
+    router.onBadUrl = router.onBadUrl.bind(instance)
   }
 
   function basic404 (request, reply) {
@@ -54,24 +57,16 @@ function fourOhFour (options) {
     })
   }
 
-  function onBadUrl (path, req, res) {
-    const { url, method } = req
-    const message = `Route ${method}:${url} not found`
-    const body = `{"error":"Not Found","message":"${message}","statusCode":404}`
-
-    // simulate normal route logging
-    if (!disableRequestLogging) {
+  function createOnBadUrl () {
+    return function onBadUrl (path, req, res) {
       const id = genReqId(req)
       const childLogger = logger.child({ reqId: id })
-      childLogger.info({ req }, 'incoming request')
-      childLogger.info({ req }, message)
-    }
+      const fourOhFourContext = this[kFourOhFourLevelInstance][kFourOhFourContext]
+      const request = new Request(id, null, req, null, childLogger, fourOhFourContext)
+      const reply = new Reply(res, request, childLogger)
 
-    res.writeHead(404, {
-      'Content-Type': 'application/json',
-      'Content-Length': body.length
-    })
-    res.end(body)
+      _onBadUrlHandler(request, reply)
+    }
   }
 
   function setContext (instance, context) {
@@ -123,8 +118,12 @@ function fourOhFour (options) {
     if (handler) {
       this[kFourOhFourLevelInstance][kCanSetNotFoundHandler] = false
       handler = handler.bind(this)
+      // update onBadUrl handler
+      _onBadUrlHandler = handler
     } else {
       handler = basic404
+      // update onBadUrl handler
+      _onBadUrlHandler = basic404
     }
 
     this.after((notHandledErr, done) => {
diff --git a/test/404s.test.js b/test/404s.test.js
index 8cec0bcd90..a5c7ec26ed 100644
--- a/test/404s.test.js
+++ b/test/404s.test.js
@@ -1872,6 +1872,22 @@ test('400 in case of bad url (pre find-my-way v2.2.0 was a 404)', t => {
     })
   })
 
+  t.test('customized 404', t => {
+    t.plan(3)
+    const fastify = Fastify({ logger: true })
+    fastify.setNotFoundHandler(function (req, reply) {
+      reply.code(404).send('this was not found')
+    })
+    fastify.inject({
+      url: '/%c0',
+      method: 'GET'
+    }, (err, response) => {
+      t.error(err)
+      t.equal(response.statusCode, 404)
+      t.same(response.payload, 'this was not found')
+    })
+  })
+
   t.end()
 })
 
diff --git a/test/logger.test.js b/test/logger.test.js
index 06e307321e..a874ce7924 100644
--- a/test/logger.test.js
+++ b/test/logger.test.js
@@ -1517,7 +1517,8 @@ test('should not log incoming request and outgoing response for 404 onBadUrl whe
     url: '/%c0',
     method: 'GET'
   }, (e, res) => {
-    t.same(lines.length, 0)
+    // it will log 1 line only because of basic404
+    t.same(lines.length, 1)
   })
 })