Custom NotFound and Middleware with new golang 1.22 http router

router.go

package main

import (
	"fmt"
	"log/slog"
	"net/http"
	"slices"
)

type (
	Middleware func(http.Handler) http.Handler

	Router struct {
		*http.ServeMux
		chain    []Middleware
		notFound http.HandlerFunc
	}
)

func NewRouter(mx ...Middleware) *Router {
	return &Router{ServeMux: http.NewServeMux(), chain: mx, notFound: http.NotFound}
}

func (r *Router) NotFound(fn http.HandlerFunc) {
	r.notFound = fn
}

// Root is a convenience function to prevent forgetting that a root path needs to end with {$}, or else it'll grab every non-matched URL.
func (r *Router) Root(method string, fn http.HandlerFunc) {
	r.HandleFunc(fmt.Sprintf("%s /{$}", method), fn)
}

func (r *Router) Use(mx ...Middleware) {
	r.chain = append(r.chain, mx...)
}

func (r *Router) Group(fn func(r *Router)) {
	fn(&Router{ServeMux: r.ServeMux, chain: slices.Clone(r.chain), notFound: r.notFound})
}

func (r *Router) Handle(path string, fn http.Handler, mx ...Middleware) {
	r.ServeMux.Handle(path, r.wrap(fn, mx))
}

func (r *Router) HandleFunc(path string, fn http.HandlerFunc, mx ...Middleware) {
	r.ServeMux.Handle(path, r.wrapFn(fn, mx))
}

func (r *Router) wrap(h http.Handler, mx []Middleware) (out http.Handler) {
	out, mx = h, append(slices.Clone(r.chain), mx...)

	slices.Reverse(mx)

	for _, m := range mx {
		out = m(out)
	}

	return
}
func (r *Router) wrapFn(fn http.HandlerFunc, mx []Middleware) (out http.Handler) {
	return r.wrap(http.Handler(fn), mx)
}

func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	_, pattern := r.ServeMux.Handler(req)
	slog.Info("route", "pattern", pattern)
	if pattern == "" {

		r.notFound(w, req)
		return
	}

	r.ServeMux.ServeHTTP(w, req)
}

routes.go

// Simple example of setting the notfound handler to a real page handler, and ensuring that the root / doesn't greedily capture all not found routes
package main

import (
	"fmt"
	"net/http"
)

func (app *application) routes() http.Handler {
	mux := NewRouter()
	mux.NotFound(app.notFound)

	mux.Use(app.logAccess)
	mux.Use(app.recoverPanic)
	mux.Use(app.securityHeaders)

	fileServer := http.FileServerFS(assets.EmbeddedFiles)
	mux.Handle("GET /static/{file...}", fileServer)

	mux.Group(func(r *web.Router) {
		r.Use(app.preventCSRF)
		r.Use(app.authenticate)
    
    // The {$} is necessary to match *only* /, otherwise you'll never see a notfound page as / will greedily match everything.
		r.HandleFunc("GET /{$}", app.home)
		
		// Alternatively, use the helper method
		// r.Root("GET", app.home)
  })
}