From 38fb28f84f5a2a9b3f993ef6b178472187655aa1 Mon Sep 17 00:00:00 2001
From: crystal <crystal@noreply.codeberg.org>
Date: Sun, 12 Jun 2022 03:50:00 +0200
Subject: [PATCH] implement custom 404 pages (#81)

solves #56.

- The expected filename is `404.html`, like GitHub Pages
- Each repo/branch can have one `404.html` file at it's root
- If a repo does not have a `pages` branch, the 404.html file from the `pages` repository is used
- You get status code 404 (unless you request /404.html which returns 200)
- The error page is cached

---
close #56

Reviewed-on: https://codeberg.org/Codeberg/pages-server/pulls/81
Reviewed-by: 6543 <6543@noreply.codeberg.org>
Co-authored-by: crystal <crystal@noreply.codeberg.org>
Co-committed-by: crystal <crystal@noreply.codeberg.org>
---
 integration/get_test.go     | 13 +++++++++++++
 server/upstream/upstream.go | 26 ++++++++++++++++++++++++--
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/integration/get_test.go b/integration/get_test.go
index 3f2048e..ef00531 100644
--- a/integration/get_test.go
+++ b/integration/get_test.go
@@ -49,6 +49,19 @@ func TestGetContent(t *testing.T) {
 	assert.True(t, getSize(resp.Body) > 1000)
 }
 
+func TestGetNotFound(t *testing.T) {
+	log.Printf("== TestGetNotFound ==\n")
+	// test custom not found pages
+	resp, err := getTestHTTPSClient().Get("https://crystal.localhost.mock.directory:4430/pages-404-demo/blah")
+	assert.NoError(t, err)
+	if !assert.EqualValues(t, http.StatusNotFound, resp.StatusCode) {
+		t.FailNow()
+	}
+	assert.EqualValues(t, "text/html; charset=utf-8", resp.Header["Content-Type"][0])
+	assert.EqualValues(t, "37", resp.Header["Content-Length"][0])
+	assert.EqualValues(t, 37, getSize(resp.Body))
+}
+
 func getTestHTTPSClient() *http.Client {
 	cookieJar, _ := cookiejar.New(nil)
 	return &http.Client{
diff --git a/server/upstream/upstream.go b/server/upstream/upstream.go
index edf4f3f..da97021 100644
--- a/server/upstream/upstream.go
+++ b/server/upstream/upstream.go
@@ -21,6 +21,11 @@ var upstreamIndexPages = []string{
 	"index.html",
 }
 
+// upstreamNotFoundPages lists pages that may be considered as custom 404 Not Found pages.
+var upstreamNotFoundPages = []string{
+	"404.html",
+}
+
 // Options provides various options for the upstream request.
 type Options struct {
 	TargetOwner,
@@ -107,6 +112,21 @@ func (o *Options) Upstream(ctx *fasthttp.RequestCtx, giteaClient *gitea.Client,
 			}
 		}
 		ctx.Response.SetStatusCode(fasthttp.StatusNotFound)
+		if o.TryIndexPages {
+			// copy the o struct & try if a not found page exists
+			optionsForNotFoundPages := *o
+			optionsForNotFoundPages.TryIndexPages = false
+			optionsForNotFoundPages.appendTrailingSlash = false
+			for _, notFoundPage := range upstreamNotFoundPages {
+				optionsForNotFoundPages.TargetPath = "/" + notFoundPage
+				if optionsForNotFoundPages.Upstream(ctx, giteaClient, branchTimestampCache, fileResponseCache) {
+					_ = fileResponseCache.Set(uri+"?timestamp="+o.timestamp(), gitea.FileResponse{
+						Exists: false,
+					}, fileCacheTimeout)
+					return true
+				}
+			}
+		}
 		if res != nil {
 			// Update cache if the request is fresh
 			_ = fileResponseCache.Set(uri+"?timestamp="+o.timestamp(), gitea.FileResponse{
@@ -141,8 +161,10 @@ func (o *Options) Upstream(ctx *fasthttp.RequestCtx, giteaClient *gitea.Client,
 	mimeType := o.getMimeTypeByExtension()
 	ctx.Response.Header.SetContentType(mimeType)
 
-	// Everything's okay so far
-	ctx.Response.SetStatusCode(fasthttp.StatusOK)
+	if ctx.Response.StatusCode() != fasthttp.StatusNotFound {
+		// Everything's okay so far
+		ctx.Response.SetStatusCode(fasthttp.StatusOK)
+	}
 	ctx.Response.Header.SetLastModified(o.BranchTimestamp)
 
 	log.Debug().Msg("response preparations")