server: Add 404 support

3 files changed, 88 insertions, 13 deletions

diff --git a/commands/commands_test.go b/commands/commands_test.go
index 97d81ec6e..d37ec1f0d 100644
--- a/commands/commands_test.go
+++ b/commands/commands_test.go
@@ -370,6 +370,11 @@ Single: {{ .Title }}
 
 `)
 
+	writeFile(t, filepath.Join(dir, "layouts", "404.html"), `
+404: {{ .Title }}|Not Found.
+
+`)
+
 	writeFile(t, filepath.Join(dir, "layouts", "_default", "list.html"), `
 
 List: {{ .Title }}
diff --git a/commands/server.go b/commands/server.go
index f082164ce..7689f03db 100644
--- a/commands/server.go
+++ b/commands/server.go
@@ -412,12 +412,18 @@ func (f *fileServer) createEndpoint(i int) (*http.ServeMux, net.Listener, string
 				// See https://docs.netlify.com/routing/redirects/rewrites-proxies/
 				if !redirect.Force {
 					path := filepath.Clean(strings.TrimPrefix(requestURI, u.Path))
-					fi, err := f.c.hugo().BaseFs.PublishFs.Stat(path)
+					if root != "" {
+						path = filepath.Join(root, path)
+					}
+					fs := f.c.publishDirServerFs
+
+					fi, err := fs.Stat(path)
+
 					if err == nil {
 						if fi.IsDir() {
 							// There will be overlapping directories, so we
 							// need to check for a file.
-							_, err = f.c.hugo().BaseFs.PublishFs.Stat(filepath.Join(path, "index.html"))
+							_, err = fs.Stat(filepath.Join(path, "index.html"))
 							doRedirect = err != nil
 						} else {
 							doRedirect = false
@@ -426,15 +432,28 @@ func (f *fileServer) createEndpoint(i int) (*http.ServeMux, net.Listener, string
 				}
 
 				if doRedirect {
-					if redirect.Status == 200 {
+					switch redirect.Status {
+					case 404:
+						w.WriteHeader(404)
+						file, err := fs.Open(filepath.FromSlash(strings.TrimPrefix(redirect.To, u.Path)))
+						if err == nil {
+							defer file.Close()
+							io.Copy(w, file)
+						} else {
+							fmt.Fprintln(w, "<h1>Page Not Found</h1>")
+						}
+						return
+					case 200:
 						if r2 := f.rewriteRequest(r, strings.TrimPrefix(redirect.To, u.Path)); r2 != nil {
 							requestURI = redirect.To
 							r = r2
 						}
-					} else {
+						fallthrough
+					default:
 						w.Header().Set("Content-Type", "")
 						http.Redirect(w, r, redirect.To, redirect.Status)
 						return
+
 					}
 				}
 
diff --git a/commands/server_test.go b/commands/server_test.go
index 56d3949ee..47cc8b9e9 100644
--- a/commands/server_test.go
+++ b/commands/server_test.go
@@ -41,12 +41,30 @@ func TestServerPanicOnConfigError(t *testing.T) {
 linenos='table'
 `
 
-	r := runServerTest(c, 0, config)
+	r := runServerTest(c,
+		serverTestOptions{
+			config: config,
+		},
+	)
 
 	c.Assert(r.err, qt.IsNotNil)
 	c.Assert(r.err.Error(), qt.Contains, "cannot parse 'Highlight.LineNos' as bool:")
 }
 
+func TestServer404(t *testing.T) {
+	c := qt.New(t)
+
+	r := runServerTest(c,
+		serverTestOptions{
+			test404:     true,
+			getNumHomes: 1,
+		},
+	)
+
+	c.Assert(r.err, qt.IsNil)
+	c.Assert(r.content404, qt.Contains, "404: 404 Page not found|Not Found.")
+}
+
 func TestServerFlags(t *testing.T) {
 	c := qt.New(t)
 
@@ -81,7 +99,13 @@ baseURL="https://example.org"
 				args = strings.Split(test.flag, "=")
 			}
 
-			r := runServerTest(c, 1, config, args...)
+			opts := serverTestOptions{
+				config:      config,
+				args:        args,
+				getNumHomes: 1,
+			}
+
+			r := runServerTest(c, opts)
 
 			test.assert(c, r)
 
@@ -140,7 +164,16 @@ baseURL="https://example.org"
 			if test.flag != "" {
 				args = strings.Split(test.flag, "=")
 			}
-			r := runServerTest(c, test.numservers, test.config, args...)
+
+			opts := serverTestOptions{
+				config:      test.config,
+				getNumHomes: test.numservers,
+				test404:     true,
+				args:        args,
+			}
+
+			r := runServerTest(c, opts)
+			c.Assert(r.content404, qt.Contains, "404: 404 Page not found|Not Found.")
 			test.assert(c, r)
 
 		})
@@ -152,11 +185,19 @@ baseURL="https://example.org"
 type serverTestResult struct {
 	err            error
 	homesContent   []string
+	content404     string
 	publicDirnames map[string]bool
 }
 
-func runServerTest(c *qt.C, getNumHomes int, config string, args ...string) (result serverTestResult) {
-	dir := createSimpleTestSite(c, testSiteConfig{configTOML: config})
+type serverTestOptions struct {
+	getNumHomes int
+	test404     bool
+	config      string
+	args        []string
+}
+
+func runServerTest(c *qt.C, opts serverTestOptions) (result serverTestResult) {
+	dir := createSimpleTestSite(c, testSiteConfig{configTOML: opts.config})
 
 	sp, err := helpers.FindAvailablePort()
 	c.Assert(err, qt.IsNil)
@@ -172,7 +213,7 @@ func runServerTest(c *qt.C, getNumHomes int, config string, args ...string) (res
 	scmd := b.newServerCmdSignaled(stop)
 
 	cmd := scmd.getCommand()
-	args = append([]string{"-s=" + dir, fmt.Sprintf("-p=%d", port)}, args...)
+	args := append([]string{"-s=" + dir, fmt.Sprintf("-p=%d", port)}, opts.args...)
 	cmd.SetArgs(args)
 
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
@@ -184,12 +225,12 @@ func runServerTest(c *qt.C, getNumHomes int, config string, args ...string) (res
 		return err
 	})
 
-	if getNumHomes > 0 {
+	if opts.getNumHomes > 0 {
 		// Esp. on slow CI machines, we need to wait a little before the web
 		// server is ready.
 		time.Sleep(567 * time.Millisecond)
-		result.homesContent = make([]string, getNumHomes)
-		for i := 0; i < getNumHomes; i++ {
+		result.homesContent = make([]string, opts.getNumHomes)
+		for i := 0; i < opts.getNumHomes; i++ {
 			func() {
 				resp, err := http.Get(fmt.Sprintf("http://localhost:%d/", port+i))
 				c.Check(err, qt.IsNil)
@@ -202,6 +243,16 @@ func runServerTest(c *qt.C, getNumHomes int, config string, args ...string) (res
 		}
 	}
 
+	if opts.test404 {
+		resp, err := http.Get(fmt.Sprintf("http://localhost:%d/this-page-does-not-exist", port))
+		c.Check(err, qt.IsNil)
+		c.Check(resp.StatusCode, qt.Equals, http.StatusNotFound)
+		if err == nil {
+			defer resp.Body.Close()
+			result.content404 = helpers.ReaderToString(resp.Body)
+		}
+	}
+
 	time.Sleep(1 * time.Second)
 
 	select {