Enable rate limiting (#214)

* Actually block when IP+UA is rate limited.

* Add tests for the rate limiter code paths

Author: robertodauria

handler/handler.go

@@ -43,6 +43,10 @@ type Signer interface {
 	Sign(cl jwt.Claims) (string, error)
 }
 
+type Limiter interface {
+	IsLimited(ip, ua string) (bool, error)
+}
+
 // Client contains state needed for xyz.
 type Client struct {
 	Signer
@@ -52,7 +56,7 @@ type Client struct {
 	PrometheusClient
 	targetTmpl  *template.Template
 	agentLimits limits.Agents
-	ipLimiter   *limits.RateLimiter
+	ipLimiter   Limiter
 }
 
 // LocatorV2 defines how the Nearest handler requests machines nearest to the
@@ -86,7 +90,7 @@ func init() {
 
 // NewClient creates a new client.
 func NewClient(project string, private Signer, locatorV2 LocatorV2, client ClientLocator,
-	prom PrometheusClient, lmts limits.Agents, limiter *limits.RateLimiter) *Client {
+	prom PrometheusClient, lmts limits.Agents, limiter Limiter) *Client {
 	return &Client{
 		Signer:           private,
 		project:          project,
@@ -175,12 +179,13 @@ func (c *Client) Nearest(rw http.ResponseWriter, req *http.Request) {
 				// TODO: Add tests for this path.
 				log.Printf("Rate limiter error: %v", err)
 			} else if limited {
+				// Log IP and UA and block the request.
+				result.Error = v2.NewError("client", tooManyRequests, http.StatusTooManyRequests)
 				metrics.RequestsTotal.WithLabelValues("nearest", "rate limit",
-					http.StatusText(http.StatusTooManyRequests)).Inc()
-				// For now, we only log the rate limit exceeded message.
-				// TODO: Actually block the request and return an appropriate HTTP error
-				// code and message.
+					http.StatusText(result.Error.Status)).Inc()
 				log.Printf("Rate limit exceeded for IP %s and UA %s", ip, ua)
+				writeResult(rw, result.Error.Status, &result)
+				return
 			}
 		} else {
 			// This should never happen if Locate is deployed on AppEngine.

handler/handler_test.go

@@ -71,6 +71,18 @@ func (l *fakeAppEngineLocator) Locate(req *http.Request) (*clientgeo.Location, e
 	return l.loc, l.err
 }
 
+type fakeRateLimiter struct {
+	limited bool
+	err     error
+}
+
+func (r *fakeRateLimiter) IsLimited(ip, ua string) (bool, error) {
+	if r.err != nil {
+		return false, r.err
+	}
+	return r.limited, nil
+}
+
 func TestClient_Nearest(t *testing.T) {
 	tests := []struct {
 		name       string
@@ -81,6 +93,7 @@ func TestClient_Nearest(t *testing.T) {
 		project    string
 		latlon     string
 		limits     limits.Agents
+		ipLimiter  Limiter
 		header     http.Header
 		wantLatLon string
 		wantKey    string
@@ -206,13 +219,96 @@ func TestClient_Nearest(t *testing.T) {
 			wantKey:    "ws://:3001/ndt_protocol",
 			wantStatus: http.StatusOK,
 		},
+		{
+			name:   "error-rate-limit-exceeded",
+			path:   "ndt/ndt5",
+			signer: &fakeSigner{},
+			locator: &fakeLocatorV2{
+				targets: []v2.Target{{Machine: "mlab1-lga0t.measurement-lab.org"}},
+			},
+			header: http.Header{
+				"X-Forwarded-For": []string{"192.0.2.1"},
+				"User-Agent":      []string{"test-client"},
+			},
+			ipLimiter: &fakeRateLimiter{
+				limited: true,
+			},
+			wantStatus: http.StatusTooManyRequests,
+		},
+		{
+			name:   "success-rate-limit-not-exceeded",
+			path:   "ndt/ndt5",
+			signer: &fakeSigner{},
+			locator: &fakeLocatorV2{
+				targets: []v2.Target{{Machine: "mlab1-lga0t.measurement-lab.org"}},
+				urls: []url.URL{
+					{Scheme: "ws", Host: ":3001", Path: "/ndt_protocol"},
+					{Scheme: "wss", Host: ":3010", Path: "ndt_protocol"},
+				},
+			},
+			header: http.Header{
+				"X-AppEngine-CityLatLong": []string{"40.3,-70.4"},
+				"X-Forwarded-For":         []string{"192.168.1.1"},
+				"User-Agent":              []string{"test-client"},
+			},
+			ipLimiter: &fakeRateLimiter{
+				limited: false,
+			},
+			wantLatLon: "40.3,-70.4",
+			wantKey:    "ws://:3001/ndt_protocol",
+			wantStatus: http.StatusOK,
+		},
+		{
+			name:   "success-rate-limiter-error",
+			path:   "ndt/ndt5",
+			signer: &fakeSigner{},
+			locator: &fakeLocatorV2{
+				targets: []v2.Target{{Machine: "mlab1-lga0t.measurement-lab.org"}},
+				urls: []url.URL{
+					{Scheme: "ws", Host: ":3001", Path: "/ndt_protocol"},
+					{Scheme: "wss", Host: ":3010", Path: "/ndt_protocol"},
+				},
+			},
+			header: http.Header{
+				"X-AppEngine-CityLatLong": []string{"40.3,-70.4"},
+				"X-Forwarded-For":         []string{"192.168.1.1"},
+				"User-Agent":              []string{"test-client"},
+			},
+			ipLimiter: &fakeRateLimiter{
+				err: errors.New("redis error"),
+			},
+			wantLatLon: "40.3,-70.4",
+			wantKey:    "ws://:3001/ndt_protocol",
+			wantStatus: http.StatusOK, // Should fail open
+		},
+		{
+			name:   "success-missing-forwarded-for",
+			path:   "ndt/ndt5",
+			signer: &fakeSigner{},
+			locator: &fakeLocatorV2{
+				targets: []v2.Target{{Machine: "mlab1-lga0t.measurement-lab.org"}},
+				urls: []url.URL{
+					{Scheme: "ws", Host: ":3001", Path: "/ndt_protocol"},
+					{Scheme: "wss", Host: ":3010", Path: "/ndt_protocol"},
+				},
+			},
+			header: http.Header{
+				"X-AppEngine-CityLatLong": []string{"40.3,-70.4"},
+				// No X-Forwarded-For
+				"User-Agent": []string{"test-client"},
+			},
+			ipLimiter:  &fakeRateLimiter{limited: false},
+			wantLatLon: "40.3,-70.4",
+			wantKey:    "ws://:3001/ndt_protocol",
+			wantStatus: http.StatusOK,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			if tt.cl == nil {
 				tt.cl = clientgeo.NewAppEngineLocator()
 			}
-			c := NewClient(tt.project, tt.signer, tt.locator, tt.cl, prom.NewAPI(nil), tt.limits, nil)
+			c := NewClient(tt.project, tt.signer, tt.locator, tt.cl, prom.NewAPI(nil), tt.limits, tt.ipLimiter)
 
 			mux := http.NewServeMux()
 			mux.HandleFunc("/v2/nearest/", c.Nearest)