Repository: go-chi/chi
File: middleware/realip.go (Go)

Code before:
// RemoteAddr will see the intended value.
//
// You should only use this middleware if you can trust the headers passed to
// you (in particular, the three headers this middleware uses), for example
// because you have placed a reverse proxy like HAProxy or nginx in front of
// chi. If your reverse proxies are configured to pass along arbitrary header
// values from the client, or if you use this middleware without a reverse
// proxy, malicious clients will be able to make you very sad (or, depending on
// how you're using RemoteAddr, vulnerable to an attack of some sort).
func RealIP(h http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		if rip := getRealIP(r, DefaultRealIPHeaders); rip != "" {
			r.RemoteAddr = rip
		}
		h.ServeHTTP(w, r)
	}

	return http.HandlerFunc(fn)
}

// RealIPCustomHeader is a middleware that sets a http.Request's RemoteAddr to the results
// of parsing the custom headers.
//
// usage:
// r.Use(RealIPCustomHeader([]string{"X-CUSTOM-IP"}))
// r.Use(RealIPCustomHeader(append(DefaultRealIPHeaders, "X-CUSTOM-IP")))
func RealIPCustomHeader(realIPHeaders []string) func(http.Handler) http.Handler {
	f := func(h http.Handler) http.Handler {
		fn := func(w http.ResponseWriter, r *http.Request) {
			if rip := getRealIP(r, realIPHeaders); rip != "" {
				r.RemoteAddr = rip
			}
			h.ServeHTTP(w, r)
		}
		return http.HandlerFunc(fn)
	}
	return f
}

func getRealIP(r *http.Request, realIPHeaders []string) string {
	for _, header := range realIPHeaders {
		if ip := r.Header.Get(header); ip != "" {
			ips := strings.Split(ip, ",")
			if ips[0] == "" || net.ParseIP(ips[0]) == nil {
				continue
			}
			return ips[0]
		}
	}
	return ""
}

Review comment (suggestion):
Feel free to use `CF-Connecting-IP` header in this example comment.

Code after:
//
// You should only use this middleware if you can trust the headers passed to
// you (in particular, the three headers this middleware uses), for example
// because you have placed a reverse proxy like HAProxy or nginx in front of
// chi. If your reverse proxies are configured to pass along arbitrary header
// values from the client, or if you use this middleware without a reverse
// proxy, malicious clients will be able to make you very sad (or, depending on
// how you're using RemoteAddr, vulnerable to an attack of some sort).
func RealIP(h http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		if rip := getRealIP(r, defaultHeaders); rip != "" {
			r.RemoteAddr = rip
		}
		h.ServeHTTP(w, r)
	}

	return http.HandlerFunc(fn)
}

// RealIPFromHeaders is a middleware that sets a http.Request's RemoteAddr to the results
// of parsing the custom headers.
//
// usage:
// r.Use(RealIPFromHeaders("CF-Connecting-IP"))
func RealIPFromHeaders(headers ...string) func(http.Handler) http.Handler {
	f := func(h http.Handler) http.Handler {
		fn := func(w http.ResponseWriter, r *http.Request) {
			if rip := getRealIP(r, headers); rip != "" {
				r.RemoteAddr = rip
			}
			h.ServeHTTP(w, r)
		}
		return http.HandlerFunc(fn)
	}
	return f
}

func getRealIP(r *http.Request, headers []string) string {
	for _, header := range headers {
		if ip := r.Header.Get(header); ip != "" {
			ips := strings.Split(ip, ",")
			if ips[0] == "" || net.ParseIP(ips[0]) == nil {
				continue
			}
			return ips[0]
		}
	}
	return ""
}