Prefer high-level RequestCtx functions over low-level ones

3 files changed, 10 insertions, 9 deletions

diff --git a/handlers/handler.go b/handlers/handler.go
index 56801a3..d6e1b5b 100644
--- a/handlers/handler.go
+++ b/handlers/handler.go
@@ -76,6 +76,11 @@ func (h *RequestHandler) HandleRequest(ctx *fasthttp.RequestCtx) {
 // information.
 func (h *RequestHandler) handleUnknownRequest(ctx *fasthttp.RequestCtx) {
 	// TODO: Show error page
-	ctx.SetStatusCode(http.StatusNotFound)
-	_, _ = ctx.WriteString("Unknown request.")
+	ctx.Error("Unknown request.", http.StatusNotFound)
+}
+
+// handleUnknownRequest handles errors which occurred during a request with a generic message.
+func (h *RequestHandler) handleError(ctx *fasthttp.RequestCtx) {
+	// TODO: Show error page
+	ctx.Error("An error occurred", http.StatusInternalServerError)
 }
diff --git a/handlers/passthrough.go b/handlers/passthrough.go
index f322d3a..b2daa88 100644
--- a/handlers/passthrough.go
+++ b/handlers/passthrough.go
@@ -2,7 +2,6 @@ package handlers
 
 import (
 	"github.com/valyala/fasthttp"
-	"net/http"
 	"portgate"
 )
 
@@ -16,14 +15,13 @@ func (h *RequestHandler) handlePassthroughRequest(ctx *fasthttp.RequestCtx, p po
 	// We reuse the request given to us by the user with minor changes to route it to the
 	// destination host.
 	ctx.Request.SetRequestURI(h.config.MakeUrl(p))
-	ctx.Request.Header.Set("Host", h.config.TargetAddress(p.DestinationIdentifier))
+	ctx.Request.Header.SetHost(h.config.TargetAddress(p.DestinationIdentifier))
 
 	// We pipe the response given to us by the destination host back to the user.
 	// Since it's possible that we get a redirect, we take this into account,
 	// but only allow upto 10 redirects.
 	err := h.client.DoRedirects(&ctx.Request, &ctx.Response, 10)
 	if err != nil {
-		ctx.SetStatusCode(http.StatusInternalServerError)
-		_, _ = ctx.WriteString("An error occurred.")
+		h.handleError(ctx)
 	}
 }
diff --git a/handlers/portgate.go b/handlers/portgate.go
index 1cd9bb8..f53f3c9 100644
--- a/handlers/portgate.go
+++ b/handlers/portgate.go
@@ -2,7 +2,6 @@ package handlers
 
 import (
 	"github.com/valyala/fasthttp"
-	"net/http"
 	"portgate"
 )
 
@@ -37,8 +36,7 @@ func (h *RequestHandler) handlePortgatePageRequest(ctx *fasthttp.RequestCtx) {
 	ctx.Response.Header.SetContentType("text/html")
 	err := h.templates.ExecuteTemplate(ctx, "authenticate.template.html", nil)
 	if err != nil {
-		ctx.SetStatusCode(http.StatusInternalServerError)
-		_, _ = ctx.WriteString("An error occurred.")
+		h.handleError(ctx)
 	}
 }