fix: add --real-ip-header flag (#9)

Author: deansheather

cmd/tunneld/main.go

@@ -104,6 +104,12 @@ func main() {
 				Value:   tunneld.DefaultWireguardNetworkPrefix.String(),
 				EnvVars: []string{"TUNNELD_WIREGUARD_NETWORK_PREFIX"},
 			},
+			&cli.StringFlag{
+				Name:    "real-ip-header",
+				Usage:   "Use the given header as the real IP address rather than the remote socket address.",
+				Value:   "",
+				EnvVars: []string{"TUNNELD_REAL_IP_HEADER"},
+			},
 			&cli.StringFlag{
 				Name:    "pprof-listen-address",
 				Usage:   "The address to listen on for pprof. If set to an empty string, pprof will not be enabled.",
@@ -137,6 +143,7 @@ func runApp(ctx *cli.Context) error {
 		wireguardMTU           = ctx.Int("wireguard-mtu")
 		wireguardServerIP      = ctx.String("wireguard-server-ip")
 		wireguardNetworkPrefix = ctx.String("wireguard-network-prefix")
+		realIPHeader           = ctx.String("real-ip-header")
 		pprofListenAddress     = ctx.String("pprof-listen-address")
 		tracingHoneycombTeam   = ctx.String("tracing-honeycomb-team")
 	)
@@ -240,6 +247,7 @@ func runApp(ctx *cli.Context) error {
 		WireguardMTU:           wireguardMTU,
 		WireguardServerIP:      wireguardServerIPParsed,
 		WireguardNetworkPrefix: wireguardNetworkPrefixParsed,
+		RealIPHeader:           realIPHeader,
 	}
 	td, err := tunneld.New(options)
 	if err != nil {

tunneld/api.go

@@ -33,7 +33,12 @@ func (api *API) Router() chi.Router {
 		// Post tunnel middleware, this middleware will never execute on
 		// tunneled connections.
 		httpmw.LimitBody(1<<20), // change back to 1MB
-		httpmw.RateLimit(10, 10*time.Second),
+		httpmw.RateLimit(httpmw.RateLimitConfig{
+			Log:          api.Log.Named("ratelimier"),
+			Count:        10,
+			Window:       10 * time.Second,
+			RealIPHeader: api.Options.RealIPHeader,
+		}),
 	)
 
 	r.Post("/tun", api.postTun)

tunneld/httpmw/ratelimit.go

@@ -2,31 +2,97 @@ package httpmw
 
 import (
 	"fmt"
+	"net"
 	"net/http"
+	"strings"
+	"sync"
 	"time"
 
 	"github.com/go-chi/httprate"
 
+	"cdr.dev/slog"
 	"github.com/coder/wgtunnel/tunneld/httpapi"
 	"github.com/coder/wgtunnel/tunnelsdk"
 )
 
+type RateLimitConfig struct {
+	Log slog.Logger
+
+	// Count of the amount of requests allowed in the Window. If the Count is
+	// zero, the rate limiter is disabled.
+	Count  int
+	Window time.Duration
+
+	// RealIPHeader is the header to use to get the real IP address of the
+	// request. If this is empty, the request's RemoteAddr is used.
+	RealIPHeader string
+}
+
 // RateLimit returns a handler that limits requests based on IP.
-func RateLimit(count int, window time.Duration) func(http.Handler) http.Handler {
-	if count <= 0 {
+func RateLimit(cfg RateLimitConfig) func(http.Handler) http.Handler {
+	if cfg.Count <= 0 {
 		return func(handler http.Handler) http.Handler {
 			return handler
 		}
 	}
 
+	var logMissingHeaderOnce sync.Once
+
 	return httprate.Limit(
-		count,
-		window,
-		httprate.WithKeyByIP(),
+		cfg.Count,
+		cfg.Window,
+		httprate.WithKeyFuncs(func(r *http.Request) (string, error) {
+			if cfg.RealIPHeader != "" {
+				val := r.Header.Get(cfg.RealIPHeader)
+				if val != "" {
+					val = strings.TrimSpace(strings.Split(val, ",")[0])
+					return canonicalizeIP(val), nil
+				}
+
+				logMissingHeaderOnce.Do(func() {
+					cfg.Log.Warn(r.Context(), "real IP header not found or invalid on request", slog.F("header", cfg.RealIPHeader), slog.F("value", val))
+				})
+			}
+
+			return httprate.KeyByIP(r)
+		}),
 		httprate.WithLimitHandler(func(rw http.ResponseWriter, r *http.Request) {
 			httpapi.Write(r.Context(), rw, http.StatusTooManyRequests, tunnelsdk.Response{
-				Message: fmt.Sprintf("You've been rate limited for sending more than %v requests in %v.", count, window),
+				Message: fmt.Sprintf("You've been rate limited for sending more than %v requests in %v.", cfg.Count, cfg.Window),
 			})
 		}),
 	)
 }
+
+// canonicalizeIP returns a form of ip suitable for comparison to other IPs.
+// For IPv4 addresses, this is simply the whole string.
+// For IPv6 addresses, this is the /64 prefix.
+//
+// This function is taken directly from go-chi/httprate:
+// https://github.com/go-chi/httprate/blob/0ea2148d09a46ae62efcad05b70d87418d8e4f43/httprate.go#L111
+func canonicalizeIP(ip string) string {
+	isIPv6 := false
+	// This is how net.ParseIP decides if an address is IPv6
+	// https://cs.opensource.google/go/go/+/refs/tags/go1.17.7:src/net/ip.go;l=704
+	for i := 0; !isIPv6 && i < len(ip); i++ {
+		switch ip[i] {
+		case '.':
+			// IPv4
+			return ip
+		case ':':
+			// IPv6
+			isIPv6 = true
+		}
+	}
+	if !isIPv6 {
+		// Not an IP address at all
+		return ip
+	}
+
+	ipv6 := net.ParseIP(ip)
+	if ipv6 == nil {
+		return ip
+	}
+
+	return ipv6.Mask(net.CIDRMask(64, 128)).String()
+}

tunneld/options.go

@@ -5,6 +5,7 @@ import (
 	"encoding/base32"
 	"encoding/hex"
 	"net"
+	"net/http"
 	"net/netip"
 	"net/url"
 	"strings"
@@ -60,6 +61,12 @@ type Options struct {
 	// at least 64 bits of space available. Defaults to fcca::/16.
 	WireguardNetworkPrefix netip.Prefix
 
+	// RealIPHeader is the header to use for getting a request's IP address. If
+	// not set, the request's RemoteAddr will be used.
+	//
+	// Used for rate limiting.
+	RealIPHeader string
+
 	// PeerDialTimeout is the timeout for dialing a peer on a request. Defaults
 	// to 10 seconds.
 	PeerDialTimeout time.Duration
@@ -113,6 +120,10 @@ func (options *Options) Validate() error {
 		return xerrors.New("WireguardServerIP must be contained within WireguardNetworkPrefix")
 	}
 
+	if options.RealIPHeader != "" {
+		options.RealIPHeader = http.CanonicalHeaderKey(options.RealIPHeader)
+	}
+
 	if options.PeerDialTimeout <= 0 {
 		options.PeerDialTimeout = DefaultPeerDialTimeout
 	}

tunneld/options_test.go

@@ -39,6 +39,7 @@ func Test_Option(t *testing.T) {
 				WireguardMTU:           tunneld.DefaultWireguardMTU + 1,
 				WireguardServerIP:      netip.MustParseAddr("feed::1"),
 				WireguardNetworkPrefix: netip.MustParsePrefix("feed::1/64"),
+				RealIPHeader:           "X-Real-Ip",
 				PeerDialTimeout:        1 * time.Second,
 			}
 
@@ -66,6 +67,7 @@ func Test_Option(t *testing.T) {
 				WireguardEndpoint: "localhost:1234",
 				WireguardPort:     1234,
 				WireguardKey:      key,
+				RealIPHeader:      "x-real-ip",
 			}
 
 			err := o.Validate()
@@ -78,6 +80,8 @@ func Test_Option(t *testing.T) {
 			require.EqualValues(t, tunneld.DefaultWireguardMTU, o.WireguardMTU)
 			require.Equal(t, tunneld.DefaultWireguardServerIP, o.WireguardServerIP)
 			require.Equal(t, tunneld.DefaultWireguardNetworkPrefix, o.WireguardNetworkPrefix)
+			// should be canonicalized.
+			require.Equal(t, "X-Real-Ip", o.RealIPHeader)
 		})
 
 		t.Run("Invalid", func(t *testing.T) {