Set up 404 handling

coder · bryphe-coder · Feb 9, 2022 · Feb 5, 2022 · Feb 5, 2022 · Feb 5, 2022
commit ca515e47a46a95278d709f4a10ad984908a76570
diff --git a/nextrouter/nextrouter.go b/nextrouter/nextrouter.go
@@ -36,6 +36,9 @@ func Handler(fileSystem fs.FS, templateFunc TemplateDataFunc) http.Handler {
 	fileHandler := http.FileServer(http.FS(fileSystem))
 	router.NotFound(fileHandler.ServeHTTP)
 
+	// Finally, if there is a 404.html available, serve that
+	serve404IfAvailable(fileSystem, router)
+
 	return router
 }
 
@@ -112,8 +115,7 @@ func serveFile(router chi.Router, fileSystem fs.FS, fileName string, templateFun
 
 		// TODO(Bryan): How to handle an error here?
 		if err != nil {
-			// TODO
-			http.Error(writer, "500", 500)
+			http.Error(writer, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
 			return
 		}
 
@@ -150,6 +152,26 @@ func serveFile(router chi.Router, fileSystem fs.FS, fileName string, templateFun
 	}
 }
 
+func serve404IfAvailable(fileSystem fs.FS, router chi.Router) {
+	// Get the file contents
+	fileBytes, err := fs.ReadFile(fileSystem, "404.html")
+	if err != nil {
+		// An error is expected if the file doesn't exist
+		return
+	}
+
+	router.NotFound(func(writer http.ResponseWriter, request *http.Request) {
+		writer.WriteHeader(http.StatusNotFound)
+		_, err = writer.Write(fileBytes)
+		if err != nil {
+			// TODO(Bryan)
+			return
+		}
+	})
+}
+
+// isDynamicRoute returns true if the file is a NextJS dynamic route, like `[orgs]`
+// Returns false if the file is not a dynamic route, or if it is a catch-all route (`[[...any]]`)
 func isDynamicRoute(fileName string) bool {
 	fileWithoutExtension := removeFileExtension(fileName)
 
@@ -162,12 +184,16 @@ func isDynamicRoute(fileName string) bool {
 	return fileWithoutExtension[0] == '[' && fileWithoutExtension[1] != '[' && fileWithoutExtension[byteLen-1] == ']'
 }
 
+// isCatchAllRoute returns true if the file is a catch-all route, like `[[...any]]`
+// Return false otherwise
 func isCatchAllRoute(fileName string) bool {
 	fileWithoutExtension := removeFileExtension(fileName)
 	ret := strings.HasPrefix(fileWithoutExtension, "[[.")
 	return ret
 }
 
+// removeFileExtension removes the extension from a file
+// For example, removeFileExtension("index.html") would return "index"
 func removeFileExtension(fileName string) string {
 	return strings.TrimSuffix(fileName, filepath.Ext(fileName))
 }
diff --git a/nextrouter/nextrouter_test.go b/nextrouter/nextrouter_test.go
@@ -16,43 +16,6 @@ import (
 func TestNextRouter(t *testing.T) {
 	t.Parallel()
 
-	t.Run("Injects template parameters", func(t *testing.T) {
-		t.Parallel()
-
-		rootFS := memfs.New()
-		err := rootFS.WriteFile("test.html", []byte("{{ .CSRF.Token }}"), 0755)
-		require.NoError(t, err)
-
-		type csrfState struct {
-			Token string
-		}
-
-		type template struct {
-			CSRF csrfState
-		}
-
-		// Add custom template function
-		templateFunc := func(request *http.Request) interface{} {
-			return template{
-				CSRF: csrfState{
-					Token: "hello-csrf",
-				},
-			}
-		}
-
-		router := nextrouter.Handler(rootFS, templateFunc)
-		server := httptest.NewServer(router)
-
-		res, err := request(server, "/test.html")
-		require.NoError(t, err)
-		defer res.Body.Close()
-
-		body, err := io.ReadAll(res.Body)
-		require.NoError(t, err)
-		require.Equal(t, string(body), "hello-csrf")
-		require.Equal(t, res.StatusCode, 200)
-	})
-
 	t.Run("Serves file at root", func(t *testing.T) {
 		t.Parallel()
 		rootFS := memfs.New()
@@ -220,6 +183,41 @@ func TestNextRouter(t *testing.T) {
 		require.Equal(t, res.StatusCode, 404)
 	})
 
+	t.Run("404 if file at root is not found", func(t *testing.T) {
+		t.Parallel()
+
+		rootFS := memfs.New()
+		err := rootFS.WriteFile("test.html", []byte("test123"), 0755)
+		require.NoError(t, err)
+
+		router := nextrouter.Handler(rootFS, noopTemplateFunc)
+		server := httptest.NewServer(router)
+
+		res, err := request(server, "/test-non-existent.html")
+		require.NoError(t, err)
+		defer res.Body.Close()
+		require.Equal(t, res.StatusCode, 404)
+	})
+
+	t.Run("Serve custom 404.html if available", func(t *testing.T) {
+		t.Parallel()
+
+		rootFS := memfs.New()
+		err := rootFS.WriteFile("404.html", []byte("404 custom content"), 0755)
+		require.NoError(t, err)
+
+		router := nextrouter.Handler(rootFS, noopTemplateFunc)
+		server := httptest.NewServer(router)
+
+		res, err := request(server, "/test-non-existent.html")
+		require.NoError(t, err)
+		defer res.Body.Close()
+		body, err := io.ReadAll(res.Body)
+		require.NoError(t, err)
+		require.Equal(t, res.StatusCode, 404)
+		require.Equal(t, string(body), "404 custom content")
+	})
+
 	t.Run("Serves dynamic-routed file", func(t *testing.T) {
 		t.Parallel()
 		rootFS := memfs.New()
@@ -305,6 +303,43 @@ func TestNextRouter(t *testing.T) {
 		require.Equal(t, string(body), "test-create")
 		require.Equal(t, res.StatusCode, 200)
 	})
+
+	t.Run("Injects template parameters", func(t *testing.T) {
+		t.Parallel()
+
+		rootFS := memfs.New()
+		err := rootFS.WriteFile("test.html", []byte("{{ .CSRF.Token }}"), 0755)
+		require.NoError(t, err)
+
+		type csrfState struct {
+			Token string
+		}
+
+		type template struct {
+			CSRF csrfState
+		}
+
+		// Add custom template function
+		templateFunc := func(request *http.Request) interface{} {
+			return template{
+				CSRF: csrfState{
+					Token: "hello-csrf",
+				},
+			}
+		}
+
+		router := nextrouter.Handler(rootFS, templateFunc)
+		server := httptest.NewServer(router)
+
+		res, err := request(server, "/test.html")
+		require.NoError(t, err)
+		defer res.Body.Close()
+
+		body, err := io.ReadAll(res.Body)
+		require.NoError(t, err)
+		require.Equal(t, string(body), "hello-csrf")
+		require.Equal(t, res.StatusCode, 200)
+	})
 }
 
 func request(server *httptest.Server, path string) (*http.Response, error) {
@@ -319,6 +354,6 @@ func request(server *httptest.Server, path string) (*http.Response, error) {
 	return res, err
 }
 
-func noopTemplateFunc(request *http.Request) interface{} {
+func noopTemplateFunc(_ *http.Request) interface{} {
 	return nil
 }