package middleware

import (
	"net"
	"net/http"
)

func TrustedProxyRealIP(trustedProxies []string) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			realIP := r.Header.Get("X-Real-IP")
			if realIP != "" && isTrustedProxy(r.RemoteAddr, trustedProxies) {
				r.RemoteAddr = withRewrittenHost(r.RemoteAddr, realIP)
			}
			next.ServeHTTP(w, r)
		})
	}
}

func isTrustedProxy(remoteAddr string, trustedProxies []string) bool {
	host := hostOf(remoteAddr)
	for _, trusted := range trustedProxies {
		if host == trusted {
			return true
		}
	}
	return false
}

func hostOf(remoteAddr string) string {
	host, _, err := net.SplitHostPort(remoteAddr)
	if err != nil {
		return remoteAddr
	}
	return host
}

func withRewrittenHost(remoteAddr string, newHost string) string {
	_, port, err := net.SplitHostPort(remoteAddr)
	if err != nil {
		return newHost
	}
	return net.JoinHostPort(newHost, port)
}