From 1b24308736d0426e22e20f960baa02ee7c461885 Mon Sep 17 00:00:00 2001
From: NielsAD <niels.a.d@gmail.com>
Date: Tue, 14 Jan 2020 11:20:38 +0100
Subject: [PATCH] Use X-Forwarded-Host if available

---
 main.go | 33 +++++++++++++++++++++++++--------
 1 file changed, 25 insertions(+), 8 deletions(-)

diff --git a/main.go b/main.go
index 7ca88cc..28272a5 100644
--- a/main.go
+++ b/main.go
@@ -84,16 +84,16 @@ func main() {
 		}
 	})
 
-	limit := stdlib.NewMiddleware(
-		limiter.New(memory.NewStore(), limiter.Rate{Period: 1 * time.Second, Limit: 5}),
-		stdlib.WithForwardHeader(*forwarded),
-	)
+	limit := stdlib.NewMiddleware(limiter.New(memory.NewStore(), limiter.Rate{Period: 1 * time.Second, Limit: 5}))
 
 	srv := &http.Server{Addr: *addr}
-	http.Handle("/idx/", limit.Handler(logRequest(http.StripPrefix("/idx/", fs))))
-	http.Handle("/dl/", limit.Handler(logRequest(http.StripPrefix("/dl/", nodir(http.FileServer(http.Dir(fs.Root)))))))
-	http.Handle("/urllist.txt", http.HandlerFunc(fs.Sitemap))
-	http.Handle("/", pub)
+	handleDefault := func(p string, h http.Handler) { http.Handle(p, realIP(*forwarded, h)) }
+	handleLimited := func(p string, h http.Handler) { handleDefault(p, limit.Handler(logRequest(http.StripPrefix(p, fs)))) }
+
+	handleLimited("/idx/", fs)
+	handleLimited("/dl/", nodir(http.FileServer(http.Dir(fs.Root))))
+	handleDefault("/urllist.txt", http.HandlerFunc(fs.Sitemap))
+	handleDefault("/", pub)
 
 	go func() {
 		sig := make(chan os.Signal, 1)
@@ -119,6 +119,23 @@ func orHyphen(s string) string {
 	return "-"
 }
 
+func realIP(trustForward bool, han http.Handler) http.Handler {
+	if !trustForward {
+		return han
+	}
+
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if realHost := r.Header.Get("X-Forwarded-Host"); realHost != "" {
+			r.Host = realHost
+		}
+		if realIP := r.Header.Get("X-Real-IP"); realIP != "" {
+			r.RemoteAddr = realIP + ":0"
+		}
+
+		han.ServeHTTP(w, r)
+	})
+}
+
 func logRequest(han http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		u, _, _ := r.BasicAuth()