mhttp: implement AddXForwardedFor

mhttp/mhttp.go

@@ -4,11 +4,14 @@ package mhttp
 
 import (
 	"context"
+	"net"
 	"net/http"
+	"strings"
 
 	"github.com/mediocregopher/mediocre-go-lib/m"
 	"github.com/mediocregopher/mediocre-go-lib/mcfg"
 	"github.com/mediocregopher/mediocre-go-lib/mlog"
+	"github.com/mediocregopher/mediocre-go-lib/mnet"
 )
 
 // CfgServer initializes and returns an *http.Server which will initialize on
@@ -36,3 +39,18 @@ func CfgServer(cfg *mcfg.Cfg, h http.Handler) *http.Server {
 	// TODO shutdown logic
 	return &srv
 }
+
+// AddXForwardedFor populates the X-Forwarded-For header on the Request to
+// convey that the request is being proxied for IP.
+//
+// If the IP is invalid, loopback, or otherwise part of a reserved range, this
+// does nothing.
+func AddXForwardedFor(r *http.Request, ipStr string) {
+	const xff = "X-Forwarded-For"
+	ip := net.ParseIP(ipStr)
+	if ip == nil || mnet.IsReservedIP(ip) { // IsReservedIP includes loopback
+		return
+	}
+	prev, _ := r.Header[xff]
+	r.Header.Set(xff, strings.Join(append(prev, ip.String()), ", "))
+}

mhttp/mhttp_test.go

@@ -0,0 +1,41 @@
+package mhttp
+
+import (
+	"net/http/httptest"
+	. "testing"
+
+	"github.com/mediocregopher/mediocre-go-lib/mtest/massert"
+)
+
+func TestAddXForwardedFor(t *T) {
+	assertXFF := func(prev []string, ipStr, expected string) massert.Assertion {
+		r := httptest.NewRequest("GET", "/", nil)
+		for i := range prev {
+			r.Header.Add("X-Forwarded-For", prev[i])
+		}
+		AddXForwardedFor(r, ipStr)
+		var a massert.Assertion
+		if expected == "" {
+			a = massert.Len(r.Header["X-Forwarded-For"], 0)
+		} else {
+			a = massert.All(
+				massert.Len(r.Header["X-Forwarded-For"], 1),
+				massert.Equal(expected, r.Header["X-Forwarded-For"][0]),
+			)
+		}
+		return massert.Comment(a, "prev:%#v ipStr:%q", prev, ipStr)
+	}
+
+	massert.Fatal(t, massert.All(
+		assertXFF(nil, "invalid", ""),
+		assertXFF(nil, "::1", ""),
+		assertXFF([]string{"8.0.0.0"}, "invalid", "8.0.0.0"),
+		assertXFF([]string{"8.0.0.0"}, "::1", "8.0.0.0"),
+
+		assertXFF(nil, "8.0.0.0", "8.0.0.0"),
+		assertXFF([]string{"8.0.0.0"}, "8.0.0.1", "8.0.0.0, 8.0.0.1"),
+		assertXFF([]string{"8.0.0.0, 8.0.0.1"}, "8.0.0.2", "8.0.0.0, 8.0.0.1, 8.0.0.2"),
+		assertXFF([]string{"8.0.0.0, 8.0.0.1", "8.0.0.2"}, "8.0.0.3",
+			"8.0.0.0, 8.0.0.1, 8.0.0.2, 8.0.0.3"),
+	))
+}