From 56bb1ddc12c38c822c30eb0ee80c1cb895449d9f Mon Sep 17 00:00:00 2001
From: Justin Harms <s1d3sw1ped@gmail.com>
Date: Sat, 19 Jul 2025 05:07:36 -0500
Subject: [PATCH] Add hop-by-hop header handling in ServeHTTP method

- Introduced a map for hop-by-hop headers to be removed from responses.
- Enhanced cache serving logic to read and filter HTTP responses, ensuring only relevant headers are forwarded.
- Updated cache writing to handle full HTTP responses, improving cache integrity and performance.
---
 steamcache/steamcache.go | 89 +++++++++++++++++++++++++++++-----------
 1 file changed, 66 insertions(+), 23 deletions(-)

diff --git a/steamcache/steamcache.go b/steamcache/steamcache.go
index 307c255..cd8f7cd 100644
--- a/steamcache/steamcache.go
+++ b/steamcache/steamcache.go
@@ -2,6 +2,7 @@
 package steamcache
 
 import (
+	"bufio"
 	"context"
 	"crypto/sha1"
 	"encoding/hex"
@@ -24,6 +25,8 @@ import (
 	"sync"
 	"time"
 
+	"bytes"
+
 	"github.com/docker/go-units"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
@@ -146,6 +149,19 @@ func verifyResponseHash(resp *http.Response, bodyData []byte, expectedHash strin
 	return strings.EqualFold(actualHash, expectedHash)
 }
 
+var hopByHopHeaders = map[string]struct{}{
+	"Connection":          {},
+	"Keep-Alive":          {},
+	"Proxy-Authenticate":  {},
+	"Proxy-Authorization": {},
+	"TE":                  {},
+	"Trailer":             {},
+	"Transfer-Encoding":   {},
+	"Upgrade":             {},
+	"Date":                {},
+	"Server":              {},
+}
+
 type SteamCache struct {
 	address  string
 	upstream string
@@ -357,25 +373,41 @@ func (sc *SteamCache) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 		w.Header().Add("X-LanCache-Processed-By", "SteamCache2") // SteamPrefill uses this header to determine if the request was processed by the cache maybe steam uses it too
 
-		reader, err := sc.vfs.Open(cacheKey)
+		cachePath := cacheKey // You may want to add a .http or .cache extension for clarity
+
+		// Try to serve from cache
+		file, err := sc.vfs.Open(cachePath)
 		if err == nil {
-			defer reader.Close()
-			w.Header().Add("X-LanCache-Status", "HIT")
-
-			io.Copy(w, reader)
-
-			logger.Logger.Info().
-				Str("key", cacheKey).
-				Str("host", r.Host).
-				Str("status", "HIT").
-				Dur("duration", time.Since(tstart)).
-				Msg("request")
-
-			requestsTotal.WithLabelValues(r.Method, "200").Inc()
-			cacheStatusTotal.WithLabelValues("HIT").Inc()
-			responseTime.WithLabelValues("HIT").Observe(time.Since(tstart).Seconds())
-
-			return
+			defer file.Close()
+			buf := bufio.NewReader(file)
+			resp, err := http.ReadResponse(buf, nil)
+			if err == nil {
+				// Remove hop-by-hop and server-specific headers
+				for k, vv := range resp.Header {
+					if _, skip := hopByHopHeaders[http.CanonicalHeaderKey(k)]; skip {
+						continue
+					}
+					for _, v := range vv {
+						w.Header().Add(k, v)
+					}
+				}
+				// Add our own headers
+				w.Header().Set("X-LanCache-Status", "HIT")
+				w.Header().Set("X-LanCache-Processed-By", "SteamCache2")
+				w.WriteHeader(resp.StatusCode)
+				io.Copy(w, resp.Body)
+				resp.Body.Close()
+				logger.Logger.Info().
+					Str("key", cacheKey).
+					Str("host", r.Host).
+					Str("status", "HIT").
+					Dur("duration", time.Since(tstart)).
+					Msg("request")
+				requestsTotal.WithLabelValues(r.Method, "200").Inc()
+				cacheStatusTotal.WithLabelValues("HIT").Inc()
+				responseTime.WithLabelValues("HIT").Observe(time.Since(tstart).Seconds())
+				return
+			}
 		}
 
 		var req *http.Request
@@ -448,8 +480,6 @@ func (sc *SteamCache) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		}
 		defer resp.Body.Close()
 
-		size := resp.ContentLength
-
 		// Read the entire response body into memory for hash verification
 		bodyData, err := io.ReadAll(resp.Body)
 		if err != nil {
@@ -500,15 +530,28 @@ func (sc *SteamCache) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		}
 
 		// Write to response (always serve the file)
-		w.Header().Add("X-LanCache-Status", "MISS")
+		// Remove hop-by-hop and server-specific headers
+		for k, vv := range resp.Header {
+			if _, skip := hopByHopHeaders[http.CanonicalHeaderKey(k)]; skip {
+				continue
+			}
+			for _, v := range vv {
+				w.Header().Add(k, v)
+			}
+		}
+		// Add our own headers
+		w.Header().Set("X-LanCache-Status", "MISS")
+		w.Header().Set("X-LanCache-Processed-By", "SteamCache2")
 		w.Write(bodyData)
 
 		// Only cache the file if hash verification passed (or no hash was present)
 		if hashVerified {
-			writer, _ := sc.vfs.Create(cacheKey, size)
+			writer, _ := sc.vfs.Create(cachePath, int64(0)) // size is not known in advance
 			if writer != nil {
 				defer writer.Close()
-				writer.Write(bodyData)
+				// Write the full HTTP response to cache
+				resp.Body = io.NopCloser(bytes.NewReader(bodyData)) // Reset body for writing
+				resp.Write(writer)
 			}
 		} else {
 			logger.Logger.Warn().