From ce4390ddd01eb6be3ebfc143daae9b1e5ebe23f9 Mon Sep 17 00:00:00 2001
From: Tom Lebreux <me@tomlebreux.com>
Date: Fri, 22 Nov 2024 10:15:23 -0500
Subject: [PATCH] Only allow some specific headers (#11)

Specifically, this prevents clients to provide a x-api-key to get access
to non-shared albums
---
 proxy/handlers.go | 24 +++++++++++++++++++-----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/proxy/handlers.go b/proxy/handlers.go
index d361e88..98c8145 100644
--- a/proxy/handlers.go
+++ b/proxy/handlers.go
@@ -12,9 +12,6 @@ import (
 )
 
 type APIReverseProxy struct {
-	upstreamURL *url.URL
-	host        string
-
 	reverseProxy *httputil.ReverseProxy
 }
 
@@ -27,6 +24,25 @@ func NewAPIReverseProxy(upstream string, host string) (*APIReverseProxy, error)
 	reverseProxy := &httputil.ReverseProxy{
 		Director: func(req *http.Request) {
 			rewriteRequestURL(req, upstreamURL)
+			headers := http.Header{}
+
+			allowedHeaders := []string{
+				"X-Forwarded-For",
+				"X-Forwarded-Host",
+				"X-Forwarded-Proto",
+				"User-Agent",
+			}
+			for _, header := range allowedHeaders {
+				vals, ok := req.Header[header]
+				if !ok {
+					continue
+				}
+
+				for _, val := range vals {
+					headers.Add(header, val)
+				}
+			}
+			req.Header = headers
 			req.Host = host
 		},
 		ModifyResponse: func(resp *http.Response) error {
@@ -38,8 +54,6 @@ func NewAPIReverseProxy(upstream string, host string) (*APIReverseProxy, error)
 	}
 
 	return &APIReverseProxy{
-		upstreamURL:  upstreamURL,
-		host:         host,
 		reverseProxy: reverseProxy,
 	}, nil
 }