From 9d769aeee7bdbdce03bbb0f6fb213d4b7d45a701 Mon Sep 17 00:00:00 2001
From: crapStone <crapstone@noreply.codeberg.org>
Date: Sun, 4 Dec 2022 21:24:58 +0000
Subject: [PATCH] Fix error page generation (#145)

Co-authored-by: crapStone <crapstone01@gmail.com>
Reviewed-on: https://codeberg.org/Codeberg/pages-server/pulls/145
Reviewed-by: Gusted <gusted@noreply.codeberg.org>
Reviewed-by: 6543 <6543@obermui.de>
Co-authored-by: crapStone <crapstone@noreply.codeberg.org>
Co-committed-by: crapStone <crapstone@noreply.codeberg.org>
---
 Justfile           |  4 ++--
 html/error.go      | 31 +++++++++++++++++--------------
 html/error_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 57 insertions(+), 16 deletions(-)
 create mode 100644 html/error_test.go

diff --git a/Justfile b/Justfile
index f7ea414..7d72fe8 100644
--- a/Justfile
+++ b/Justfile
@@ -38,10 +38,10 @@ tool-gofumpt:
     fi
 
 test:
-    go test -race codeberg.org/codeberg/pages/server/...
+    go test -race codeberg.org/codeberg/pages/server/... codeberg.org/codeberg/pages/html/
 
 test-run TEST:
-    go test -race -run "^{{TEST}}$" codeberg.org/codeberg/pages/server/...
+    go test -race -run "^{{TEST}}$" codeberg.org/codeberg/pages/server/... codeberg.org/codeberg/pages/html/
 
 integration:
     go test -race -tags integration codeberg.org/codeberg/pages/integration/...
diff --git a/html/error.go b/html/error.go
index ac222c4..206b123 100644
--- a/html/error.go
+++ b/html/error.go
@@ -15,16 +15,27 @@ func ReturnErrorPage(ctx *context.Context, msg string, statusCode int) {
 	ctx.RespWriter.Header().Set("Content-Type", "text/html; charset=utf-8")
 	ctx.RespWriter.WriteHeader(statusCode)
 
-	if msg == "" {
-		msg = errorBody(statusCode)
-	} else {
-		// TODO: use template engine
-		msg = strings.ReplaceAll(strings.ReplaceAll(ErrorPage, "%message%", msg), "%status%", http.StatusText(statusCode))
-	}
+	msg = generateResponse(msg, statusCode)
 
 	_, _ = ctx.RespWriter.Write([]byte(msg))
 }
 
+// TODO: use template engine
+func generateResponse(msg string, statusCode int) string {
+	if msg == "" {
+		msg = strings.ReplaceAll(NotFoundPage,
+			"%status%",
+			strconv.Itoa(statusCode)+" "+errorMessage(statusCode))
+	} else {
+		msg = strings.ReplaceAll(
+			strings.ReplaceAll(ErrorPage, "%message%", template.HTMLEscapeString(msg)),
+			"%status%",
+			http.StatusText(statusCode))
+	}
+
+	return msg
+}
+
 func errorMessage(statusCode int) string {
 	message := http.StatusText(statusCode)
 
@@ -37,11 +48,3 @@ func errorMessage(statusCode int) string {
 
 	return message
 }
-
-// TODO: use template engine
-func errorBody(statusCode int) string {
-	return template.HTMLEscapeString(
-		strings.ReplaceAll(NotFoundPage,
-			"%status%",
-			strconv.Itoa(statusCode)+" "+errorMessage(statusCode)))
-}
diff --git a/html/error_test.go b/html/error_test.go
new file mode 100644
index 0000000..f5da08c
--- /dev/null
+++ b/html/error_test.go
@@ -0,0 +1,38 @@
+package html
+
+import (
+	"net/http"
+	"strings"
+	"testing"
+)
+
+func TestValidMessage(t *testing.T) {
+	testString := "requested blacklisted path"
+	statusCode := http.StatusForbidden
+
+	expected := strings.ReplaceAll(
+		strings.ReplaceAll(ErrorPage, "%message%", testString),
+		"%status%",
+		http.StatusText(statusCode))
+	actual := generateResponse(testString, statusCode)
+
+	if expected != actual {
+		t.Errorf("generated response did not match: expected: '%s', got: '%s'", expected, actual)
+	}
+}
+
+func TestMessageWithHtml(t *testing.T) {
+	testString := `abc<img src=1 onerror=alert("xss");`
+	escapedString := "abc&lt;img src=1 onerror=alert(&#34;xss&#34;);"
+	statusCode := http.StatusNotFound
+
+	expected := strings.ReplaceAll(
+		strings.ReplaceAll(ErrorPage, "%message%", escapedString),
+		"%status%",
+		http.StatusText(statusCode))
+	actual := generateResponse(testString, statusCode)
+
+	if expected != actual {
+		t.Errorf("generated response did not match: expected: '%s', got: '%s'", expected, actual)
+	}
+}