From 6f635825a9b5ed2bd1ddd35110cdd28bb63c2cbf Mon Sep 17 00:00:00 2001
From: Sasha Koshka <sashakoshka@tebibyte.media>
Date: Sun, 8 Dec 2024 01:50:15 -0500
Subject: [PATCH] http: Make methods of .Data.Res actually work

---
 http.go         | 66 -------------------------------------------------
 http/handler.go | 10 +++++---
 http/http.go    | 39 +++++++++++++++++++++++++----
 3 files changed, 41 insertions(+), 74 deletions(-)
 delete mode 100644 http.go

diff --git a/http.go b/http.go
deleted file mode 100644
index d426ed9..0000000
--- a/http.go
+++ /dev/null
@@ -1,66 +0,0 @@
-package step
-
-import "io"
-import "bytes"
-import "net/http"
-
-// HTTPData represents information about an ongoing HTTP request that is made
-// available to templates as they are being executed.
-type HTTPData struct {
-	// Res is like an http.ResponseWriter.
-	Res struct {
-		Header      http.Header
-		WriteHeader func (statusCode int)
-		Reset       func ()
-	}
-	// Req is the HTTP request.
-	Req *http.Request
-}
-
-var _ http.ResponseWriter = new(HTTPResponseRecorder)
-
-// HTTPResponseRecorder is an http.ResponseWriter that can buffer a response to
-// be played back later.
-type HTTPResponseRecorder struct {
-	Status int
-	Head   http.Header
-	buffer bytes.Buffer
-}
-
-func (this *HTTPResponseRecorder) Header () http.Header {
-	if this.Head == nil {
-		this.Head = make(http.Header)
-	}
-	return this.Head
-}
-
-func (this *HTTPResponseRecorder) Write (buffer []byte) (int, error) {
-	return this.buffer.Write(buffer)
-}
-
-func (this *HTTPResponseRecorder) WriteHeader (statusCode int) {
-	this.Status = statusCode
-}
-// Play replays the response to the given http.ResponseWriter. This resets the
-// recorder.
-func (this *HTTPResponseRecorder) Play (res http.ResponseWriter) error {
-	defer this.Reset()
-	status := this.Status
-	if status == 0 {
-		status = http.StatusOK
-	}
-	header := res.Header()
-	for name, value := range this.Head {
-		header[name] = value
-	}
-	res.WriteHeader(status)
-	_, err := io.Copy(res, &this.buffer)
-	return err
-}
-
-// Reset resets this response recorder so it can be used again.
-func (this *HTTPResponseRecorder) Reset () {
-	this.buffer.Reset()
-	this.Head = nil
-	this.Status = http.StatusOK
-}
diff --git a/http/handler.go b/http/handler.go
index 28a6308..1fd42be 100644
--- a/http/handler.go
+++ b/http/handler.go
@@ -123,9 +123,13 @@ func (this *Handler) serveDocument (
 
 	// execute document
 	data := HTTPData { }
-	data.Res.Header      = recorder.Header()
-	data.Res.WriteHeader = recorder.WriteHeader
-	data.Res.Reset       = resetRecorder
+	data.Res = WrappedResponseWriter {
+		responseWriter: res,
+		resetFunc: resetRecorder,
+		Header: WrappedHeader {
+			Header: recorder.Header(),
+		},
+	}
 	data.Req = req
 	err = document.Execute(&recorder, step.ExecutionData {
 		Data: data,
diff --git a/http/http.go b/http/http.go
index a3fd9c8..0a3d3a8 100644
--- a/http/http.go
+++ b/http/http.go
@@ -8,15 +8,44 @@ import "net/http"
 // available to templates as they are being executed.
 type HTTPData struct {
 	// Res is like an http.ResponseWriter.
-	Res struct {
-		Header      http.Header
-		WriteHeader func (statusCode int)
-		Reset       func ()
-	}
+	Res WrappedResponseWriter
 	// Req is the HTTP request.
 	Req *http.Request
 }
 
+type WrappedResponseWriter struct {
+	responseWriter http.ResponseWriter
+	resetFunc func ()
+	Header WrappedHeader
+}
+
+func (this WrappedResponseWriter) WriteHeader (statusCode int) string {
+	this.responseWriter.WriteHeader(statusCode)
+	return ""
+}
+
+func (this WrappedResponseWriter) Reset () string {
+	this.resetFunc()
+	return ""
+}
+
+type WrappedHeader struct { http.Header }
+
+func (this WrappedHeader) Add (name, value string) string {
+	this.Header.Add(name, value)
+	return ""
+}
+
+func (this WrappedHeader) Del (name string) string {
+	this.Header.Del(name)
+	return ""
+}
+
+func (this WrappedHeader) Set (name, value string) string {
+	this.Header.Set(name, value)
+	return ""
+}
+
 var _ http.ResponseWriter = new(HTTPResponseRecorder)
 
 // HTTPResponseRecorder is an http.ResponseWriter that can buffer a response to