Custom 404 with Gorilla Mux and std http.FileServer

I have the following code and everything works fine.

var view404 = template.Must(template.ParseFiles("views/404.html"))

func NotFound(w http.ResponseWriter, r *http.Request) {
  w.WriteHeader(404)
  err := view404.Execute(w, nil)
  check(err)
}

func main() {
  router := mux.NewRouter()
  router.StrictSlash(true)
  router.NotFoundHandler = http.HandlerFunc(NotFound)
  router.Handle("/", IndexHandler).Methods("GET")
  router.PathPrefix("/public/").Handler(http.StripPrefix("/public/", http.FileServer(http.Dir("public"))))
  http.Handle("/", router)
  http.ListenAndServe(":8000", nil)
}

A request to a route like /cannot/find shows my custom 404 template. All static files inside my /public/ directory are also properly served.

I have a problem handling non-existent static files and showing my custom NotFound handler for them. A request to /public/cannot/find calls the standard http.NotFoundHandler which replies with

> 404 page not found

How can I have the same custom NotFoundHandler for normal routes and static files?

Update

I ended up implementing my own FileHandler by wrapping http.ServeFile as @Dewy Broto suggested.

type FileHandler struct {
  Path string
}

func (f FileHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
  prefix := "public"
  http.ServeFile(w, r, path.Join(prefix, f.Path))
}

// ...

router.Handle("/css/styles.css", FileHandler{"/css/styles.css"}).Methods("GET")

Now my NotFound handler catches all missing routes and even missing files.

The FileServer is generating the 404 response. The FileServer handles all requests passed to it by the mux including requests for missing files. There are a few ways to to serve static files with a custom 404 page:

• Write your own file handler using ServeContent. This handler can generate error responses in whatever way you want. It's not a lot of code if you don't generate index pages.
• Wrap the FileServer handler with another handler that hooks the ResponseWriter passed to the FileHandler. The hook writes a different body when WriteHeader(404) is called.
• Register each static resource with the mux so that not found errors are handled by the catchall in the mux. This approach requires a simple wrapper around ServeFile.

Here's a sketch of the wrapper described in the second approach:

type hookedResponseWriter struct {
    http.ResponseWriter
    ignore bool
}

func (hrw *hookedResponseWriter) WriteHeader(status int) {
    hrw.ResponseWriter.WriteHeader(status)
    if status == 404 {
	    hrw.ignore = true
	    // Write custom error here to hrw.ResponseWriter
    }
}

func (hrw *hookedResponseWriter) Write(p []byte) (int, error) {
    if hrw.ignore {
	    return len(p), nil
    }
    return hrw.ResponseWriter.Write(p)
}

type NotFoundHook struct {
    h http.Handler
}

func (nfh NotFoundHook) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    nfh.h.ServeHTTP(&hookedResponseWriter{ResponseWriter: w}, r)
}

Use the hook by wrapping the FileServer:

 router.PathPrefix("/public/").Handler(NotFoundHook{http.StripPrefix("/public/", http.FileServer(http.Dir("public")))})

One caveat of this simple hook is that it blocks an optimization in the server for copying from a file to a socket.