From fb11e8eaf17a135fd6eb0fb7f012015464f396ce Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=9Cllar=20Seerme?= <yates@protonmail.ch>
Date: Thu, 1 Feb 2024 16:54:31 +0200
Subject: [PATCH] refactor: use separate handler for 'Etag'

should've done so from the start, but I wanted to get something out
there first

got the idea and most of the implementation from:
https://sidney.kochman.org/2018/etag-middleware-go/
---
 chart.go              | 15 ---------------
 cmd/gobarchar/main.go | 43 ++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 42 insertions(+), 16 deletions(-)

diff --git a/chart.go b/chart.go
index 5271324..f42c89b 100644
--- a/chart.go
+++ b/chart.go
@@ -1,9 +1,7 @@
 package gobarchar
 
 import (
-	"crypto/sha1"
 	"fmt"
-	"io"
 	"math"
 	"math/rand"
 	"net/http"
@@ -73,19 +71,6 @@ func PresentBarChart(w http.ResponseWriter, r *http.Request) {
 
 	chart := createBarChart(r)
 
-	h := sha1.New()
-	// TODO: See if 'r.URL.RawQuery' would be faster to hash
-	io.WriteString(h, chart)
-	etag := fmt.Sprintf("%x", h.Sum(nil))
-
-	if match := r.Header.Get("If-None-Match"); match != "" {
-		if strings.Contains(match, etag) {
-			w.WriteHeader(http.StatusNotModified)
-			return
-		}
-	}
-	w.Header().Set("Etag", etag)
-
 	// Skip all templating if user is requesting through 'curl' or 'wget'
 	agent := r.UserAgent()
 	if strings.HasPrefix(agent, "curl") || strings.HasPrefix(agent, "Wget") {
diff --git a/cmd/gobarchar/main.go b/cmd/gobarchar/main.go
index e28a7ca..481ea89 100644
--- a/cmd/gobarchar/main.go
+++ b/cmd/gobarchar/main.go
@@ -1,6 +1,11 @@
 package main
 
 import (
+	"bytes"
+	"crypto/sha1"
+	"fmt"
+	"hash"
+	"io"
 	"log"
 	"net/http"
 	"os"
@@ -11,13 +16,24 @@ import (
 
 var defaultPort = "8080"
 
+type etagResponseWriter struct {
+	http.ResponseWriter
+	buf  bytes.Buffer
+	hash hash.Hash
+	w    io.Writer
+}
+
+func (e *etagResponseWriter) Write(p []byte) (int, error) {
+	return e.w.Write(p)
+}
+
 func main() {
 	port := os.Getenv("PORT")
 	if port == "" {
 		port = defaultPort
 	}
 
-	http.Handle("/", timer(http.HandlerFunc(gobarchar.PresentBarChart)))
+	http.Handle("/", timer(etag(http.HandlerFunc(gobarchar.PresentBarChart))))
 
 	log.Println("listening on:", port)
 	log.Fatal(http.ListenAndServe(":"+port, nil))
@@ -31,3 +47,28 @@ func timer(h http.Handler) http.Handler {
 		log.Println("completed in:", duration)
 	})
 }
+
+func etag(h http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ew := &etagResponseWriter{
+			ResponseWriter: w,
+			buf:            bytes.Buffer{},
+			hash:           sha1.New(),
+		}
+		ew.w = io.MultiWriter(&ew.buf, ew.hash)
+
+		h.ServeHTTP(ew, r)
+
+		etag := fmt.Sprintf("%x", ew.hash.Sum(nil))
+		w.Header().Set("Etag", etag)
+
+		if r.Header.Get("If-None-Match") == etag {
+			w.WriteHeader(http.StatusNotModified)
+		} else {
+			_, err := ew.buf.WriteTo(w)
+			if err != nil {
+				log.Println("unable to write HTTP response", err)
+			}
+		}
+	})
+}