From 30a9c83b541f6ae2ad243a46643d6269c288c8c7 Mon Sep 17 00:00:00 2001
From: jtams <jtams322@gmail.com>
Date: Thu, 21 Dec 2023 03:54:15 -0800
Subject: [PATCH] Updates route handler when using NotFound

NotFound and MethodNotAllowed follow the middleware chain without the
need to define a handler.
---
 mux.go | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/mux.go b/mux.go
index 735ab232..72430d69 100644
--- a/mux.go
+++ b/mux.go
@@ -189,6 +189,11 @@ func (mx *Mux) Trace(pattern string, handlerFn http.HandlerFunc) {
 // NotFound sets a custom http.HandlerFunc for routing paths that could
 // not be found. The default 404 handler is `http.NotFound`.
 func (mx *Mux) NotFound(handlerFn http.HandlerFunc) {
+	// Build the computed routing handler for this routing pattern.
+	if !mx.inline && mx.handler == nil {
+		mx.updateRouteHandler()
+	}
+
 	// Build NotFound handler chain
 	m := mx
 	hFn := handlerFn
@@ -209,6 +214,11 @@ func (mx *Mux) NotFound(handlerFn http.HandlerFunc) {
 // MethodNotAllowed sets a custom http.HandlerFunc for routing paths where the
 // method is unresolved. The default handler returns a 405 with an empty body.
 func (mx *Mux) MethodNotAllowed(handlerFn http.HandlerFunc) {
+	// Build the computed routing handler for this routing pattern.
+	if !mx.inline && mx.handler == nil {
+		mx.updateRouteHandler()
+	}
+
 	// Build MethodNotAllowed handler chain
 	m := mx
 	hFn := handlerFn