From 54aee1326f790735785dede18e8dbfd130d0821f Mon Sep 17 00:00:00 2001
From: dadav <33197631+dadav@users.noreply.github.com>
Date: Sat, 24 Aug 2024 18:59:53 +0200
Subject: [PATCH] fix: Try to fix proxy issue

---
 internal/middleware/proxy.go | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/internal/middleware/proxy.go b/internal/middleware/proxy.go
index 40edbf6..1128922 100644
--- a/internal/middleware/proxy.go
+++ b/internal/middleware/proxy.go
@@ -1,6 +1,7 @@
 package middleware
 
 import (
+	"bytes"
 	"net/http"
 	"net/http/httputil"
 	"net/url"
@@ -11,22 +12,28 @@ import (
 // capturedResponseWriter is a custom response writer that captures the response status
 type capturedResponseWriter struct {
 	http.ResponseWriter
-	body   []byte
+	body   *bytes.Buffer
 	status int
 }
 
+func NewCapturedResponseWriter(w http.ResponseWriter) *capturedResponseWriter {
+	return &capturedResponseWriter{
+		ResponseWriter: w,
+		body:           new(bytes.Buffer),
+	}
+}
+
 func (w *capturedResponseWriter) WriteHeader(code int) {
 	w.status = code
 }
 
 func (w *capturedResponseWriter) Write(body []byte) (int, error) {
-	w.body = body
-	return len(body), nil
+	return w.body.Write(body)
 }
 
 func (w *capturedResponseWriter) sendCapturedResponse() {
 	w.ResponseWriter.WriteHeader(w.status)
-	w.ResponseWriter.Write(w.body)
+	w.ResponseWriter.Write(w.body.Bytes())
 }
 
 func NewSingleHostReverseProxy(target *url.URL) *httputil.ReverseProxy {
@@ -42,7 +49,7 @@ func NewSingleHostReverseProxy(target *url.URL) *httputil.ReverseProxy {
 func ProxyFallback(upstreamHost string, forwardToProxy func(int) bool, proxiedResponseCb func(*http.Response)) func(next http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			capturedResponseWriter := &capturedResponseWriter{ResponseWriter: w}
+			capturedResponseWriter := NewCapturedResponseWriter(w)
 			next.ServeHTTP(capturedResponseWriter, r)
 
 			if forwardToProxy(capturedResponseWriter.status) {