base/apis/middlewares_rate_limit_test.go

package apis_test

import (
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/hanzoai/base/apis"
	"github.com/hanzoai/base/core"
	"github.com/hanzoai/base/tests"
)

func TestDefaultRateLimitMiddleware(t *testing.T) {
	app, _ := tests.NewTestApp()
	defer app.Cleanup()

	app.Settings().RateLimits.Enabled = true
	app.Settings().RateLimits.Rules = []core.RateLimitRule{
		{
			Label:       "/rate/",
			MaxRequests: 2,
			Duration:    1,
		},
		{
			Label:       "/rate/b",
			MaxRequests: 3,
			Duration:    1,
		},
		{
			Label:       "POST /rate/b",
			MaxRequests: 1,
			Duration:    1,
		},
		{
			Label:       "/rate/guest",
			MaxRequests: 1,
			Duration:    1,
			Audience:    core.RateLimitRuleAudienceGuest,
		},
		{
			Label:       "/rate/auth",
			MaxRequests: 1,
			Duration:    1,
			Audience:    core.RateLimitRuleAudienceAuth,
		},
	}

	baseRouter, err := apis.NewRouter(app)
	if err != nil {
		t.Fatal(err)
	}
	baseRouter.GET("/norate", func(e *core.RequestEvent) error {
		return e.String(200, "norate")
	}).BindFunc(func(e *core.RequestEvent) error {
		return e.Next()
	})
	baseRouter.GET("/rate/a", func(e *core.RequestEvent) error {
		return e.String(200, "a")
	})
	baseRouter.GET("/rate/b", func(e *core.RequestEvent) error {
		return e.String(200, "b")
	})
	baseRouter.GET("/rate/guest", func(e *core.RequestEvent) error {
		return e.String(200, "guest")
	})
	baseRouter.GET("/rate/auth", func(e *core.RequestEvent) error {
		return e.String(200, "auth")
	})

	mux, err := baseRouter.BuildMux()
	if err != nil {
		t.Fatal(err)
	}

	scenarios := []struct {
		url            string
		wait           float64 // ms
		authenticated  bool
		expectedStatus int
	}{
		{"/norate", 0, false, 200},
		{"/norate", 0, false, 200},
		{"/norate", 0, false, 200},
		{"/norate", 0, false, 200},
		{"/norate", 0, false, 200},

		{"/rate/a", 0, false, 200},
		{"/rate/a", 1100, false, 200}, // (fixed window check) wait enough to ensure that it can't fit 2 requests in 1s
		{"/rate/a", 1100, false, 200},
		{"/rate/a", 1100, false, 200},
		{"/rate/a", 0, false, 200},
		{"/rate/a", 0, false, 429},
		{"/rate/a", 0, false, 429},
		{"/rate/a", 1200, false, 200},
		{"/rate/a", 0, false, 200},
		{"/rate/a", 0, false, 429},

		{"/rate/b", 0, false, 200},
		{"/rate/b", 0, false, 200},
		{"/rate/b", 0, false, 200},
		{"/rate/b", 0, false, 429},
		{"/rate/b", 1200, false, 200},
		{"/rate/b", 0, false, 200},
		{"/rate/b", 0, false, 200},
		{"/rate/b", 0, false, 429},

		// "auth" with guest (should fallback to the /rate/ rule)
		{"/rate/auth", 0, false, 200},
		{"/rate/auth", 0, false, 200},
		{"/rate/auth", 0, false, 429},
		{"/rate/auth", 0, false, 429},

		// "auth" rule with regular user (should match the /rate/auth rule)
		{"/rate/auth", 0, true, 200},
		{"/rate/auth", 0, true, 429},
		{"/rate/auth", 0, true, 429},

		// "guest" with guest (should match the /rate/guest rule)
		{"/rate/guest", 0, false, 200},
		{"/rate/guest", 0, false, 429},
		{"/rate/guest", 0, false, 429},

		// "guest" rule with regular user (should fallback to the /rate/ rule)
		{"/rate/guest", 1200, true, 200},
		{"/rate/guest", 0, true, 200},
		{"/rate/guest", 0, true, 429},
		{"/rate/guest", 0, true, 429},
	}

	for _, s := range scenarios {
		t.Run(s.url, func(t *testing.T) {
			rec := httptest.NewRecorder()
			req := httptest.NewRequest("GET", s.url, nil)

			if s.authenticated {
				auth, err := app.FindAuthRecordByEmail("users", "[email protected]")
				if err != nil {
					t.Fatal(err)
				}

				token, err := auth.NewAuthToken()
				if err != nil {
					t.Fatal(err)
				}

				req.Header.Add("Authorization", token)
			}

			if s.wait > 0 {
				time.Sleep(time.Duration(s.wait) * time.Millisecond)
			}

			mux.ServeHTTP(rec, req)

			result := rec.Result()

			if result.StatusCode != s.expectedStatus {
				t.Fatalf("Expected response status %d, got %d", s.expectedStatus, result.StatusCode)
			}
		})
	}
}

// A limit is per client, so what the limiter does with a request that names no
// client is the whole of this test. It refuses: a rule that admits whatever it
// cannot count is not a rule, and the request that produces no address is the
// one a rule was least able to see coming.
//
// A transport that terminates elsewhere and hands the request on carries no
// socket peer, which is where an address goes missing in practice. The
// deployment says where to read it instead — Settings.TrustedProxy — and that
// is the third case here.
func TestRateLimitNeedsAClientAddress(t *testing.T) {
	app, _ := tests.NewTestApp()
	defer app.Cleanup()

	app.Settings().RateLimits.Enabled = true
	app.Settings().RateLimits.Rules = []core.RateLimitRule{
		{Label: "/probe", MaxRequests: 100, Duration: 1},
	}

	router, err := apis.NewRouter(app)
	if err != nil {
		t.Fatal(err)
	}
	router.GET("/probe", func(e *core.RequestEvent) error {
		return e.String(http.StatusOK, "probe")
	})

	mux, err := router.BuildMux()
	if err != nil {
		t.Fatal(err)
	}

	probe := func(remote string, forwarded string) int {
		req := httptest.NewRequest(http.MethodGet, "/probe", nil)
		req.RemoteAddr = remote
		if forwarded != "" {
			req.Header.Set("X-Forwarded-For", forwarded)
		}
		rec := httptest.NewRecorder()
		mux.ServeHTTP(rec, req)
		return rec.Code
	}

	// A socket peer, and a rule the request is nowhere near.
	if code := probe("192.0.2.4:1234", ""); code != http.StatusOK {
		t.Fatalf("a request carrying a client address answered %d, want 200", code)
	}

	// No socket peer.
	if code := probe("", ""); code != http.StatusTooManyRequests {
		t.Fatalf("a request carrying no client address answered %d, want 429", code)
	}

	// Still no socket peer, but the deployment names a header to read the
	// client from, so there is a client again.
	app.Settings().TrustedProxy.Headers = []string{"X-Forwarded-For"}
	if code := probe("", "203.0.113.7"); code != http.StatusOK {
		t.Fatalf("a forwarded client address answered %d, want 200", code)
	}
}