Merge pull request #2 from josephcopenhaver/jpcope/develop

support response buffering
josephcopenhaver · May 26, 2021 · ececbcd · ececbcd
2 parents b6e7b19 + 6a5f7c4
commit ececbcd
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 9 deletions.
diff --git a/cmd/autocertproxy/main.go b/cmd/autocertproxy/main.go
@@ -68,7 +68,7 @@ func main() {
 		)
 	}
 
-	ctx := proxy.RootContext()
+	ctx, _ := proxy.RootContext()
 
 	if err := p.ListenAndServe(ctx); err != nil {
 		os.Exit(1)

diff --git a/internal/proxy/config/config.go b/internal/proxy/config/config.go
@@ -24,6 +24,7 @@ type Config struct {
 	DstHost               string        `json:"DST_HOST" envconfig:"DST_HOST" default:"127.0.0.1"`
 	DstHostHeader         string        `json:"DST_HOST_HEADER" envconfig:"DST_HOST_HEADER" default:""`
 	DstPort               int           `json:"DST_PORT" envconfig:"DST_PORT" default:"8080"`
+	ResponseBufferEnabled bool          `json:"RESPONSE_BUFFER_ENABLED" envconfig:"RESPONSE_BUFFER_ENABLED" default:"false"`
 	AdminEmail            string        `json:"ADMIN_EMAIL" envconfig:"ADMIN_EMAIL" default:""`
 	AutocertCacheDir      string        `json:"AUTOCERT_CACHE_DIR" envconfig:"AUTOCERT_CACHE_DIR" default:"./.autocert"`
 	Authorization         string        `json:"AUTHORIZATION" default:""`

diff --git a/internal/proxy/context.go b/internal/proxy/context.go
@@ -9,18 +9,26 @@ import (
 
 // RootContext returns a context that is canceled when the
 // system process receives an interrupt, sigint, or sigterm
-func RootContext() context.Context {
+//
+// Also returns a function that can be used to cancel the context.
+func RootContext() (context.Context, func()) {
 
 	ctx, cancel := context.WithCancel(context.Background())
 
-	done := make(chan os.Signal, 1)
+	procDone := make(chan os.Signal, 1)
 
-	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
+	signal.Notify(procDone, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
 
 	go func() {
-		<-done
-		cancel()
+		defer cancel()
+
+		done := ctx.Done()
+
+		select {
+		case <-procDone:
+		case <-done:
+		}
 	}()
 
-	return ctx
+	return ctx, cancel
 }
diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go
@@ -2,8 +2,10 @@ package proxy
 
 import (
 	"context"
+	"io"
 	"net"
 	"net/http"
+	"net/http/httptest"
 	"net/http/httputil"
 	"net/url"
 	"os"
@@ -227,7 +229,40 @@ func (p *proxy) ListenAndServe(ctx context.Context) error {
 			})
 		}
 
-		if p.cfg.Authorization != "" {
+		if cfg.ResponseBufferEnabled {
+			prevHandler := handler
+
+			handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+
+				// read response body till end
+				var resp *http.Response
+				{
+					rec := httptest.NewRecorder()
+
+					prevHandler.ServeHTTP(rec, r)
+
+					resp = rec.Result()
+				}
+
+				// send response headers
+				{
+					h := w.Header()
+					for k, values := range resp.Header {
+						for _, v := range values {
+							h.Add(k, v)
+						}
+					}
+				}
+
+				w.WriteHeader(resp.StatusCode)
+
+				// send response body
+				_, err := io.Copy(w, resp.Body)
+				_ = err // intentionally ignoring error, upstream or downstream must have had an issue
+			})
+		}
+
+		if cfg.Authorization != "" {
 			prevHandler := handler
 
 			handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -236,7 +271,7 @@ func (p *proxy) ListenAndServe(ctx context.Context) error {
 				{
 					user, pass, ok := r.BasicAuth()
 
-					if !ok || user+":"+pass != p.cfg.Authorization {
+					if !ok || user+":"+pass != cfg.Authorization {
 						http.Error(w, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
 						return
 					}