tenseleyflow/shithub / d1e0592

+// SPDX-License-Identifier: AGPL-3.0-or-later

+

+// Package ratelimit owns the cross-surface counter used by S35's

+// rate-limit middleware. It generalizes S05's auth/throttle pattern

+// against the shared rate_limits table so any new surface (API,

+// search, git transports) plugs in with a single Allow() call.

+//

+// Backend: Postgres-backed fixed-window counter via sqlc UPSERT.

+// At launch scale this is well within Postgres's comfort zone and

+// avoids introducing a Redis dependency. Move only if profiling

+// demands it (S36).

+package ratelimit

+

+import (

+	"context"

+	"errors"

+	"fmt"

+	"net/netip"

+	"time"

+

+	"github.com/jackc/pgx/v5/pgtype"

+	"github.com/jackc/pgx/v5/pgxpool"

+

+	ratelimitdb "github.com/tenseleyFlow/shithub/internal/ratelimit/sqlc"

+)

+

+// Limiter is the package's primary handle. Construct with New.

+type Limiter struct {

+	q    *ratelimitdb.Queries

+	pool *pgxpool.Pool

+}

+

+// New wires a Limiter against a pool. The pool is required;

+// constructing with nil panics so callers fail at boot, not at

+// first request.

+func New(pool *pgxpool.Pool) *Limiter {

+	if pool == nil {

+		panic("ratelimit: nil pool")

+	}

+	return &Limiter{q: ratelimitdb.New(), pool: pool}

+}

+

+// Policy declares a per-(scope, key) limit: at most Max hits within

+// the rolling Window.

+type Policy struct {

+	Scope  string        // e.g. "api:anon", "search", "git:https"

+	Max    int           // hits permitted within Window

+	Window time.Duration // window length (15s, 1m, 1h, …)

+}

+

+// Decision is the verdict from Allow.

+type Decision struct {

+	Allowed    bool

+	Remaining  int           // hits left in the current window (post-increment)

+	Limit      int           // == Policy.Max, surfaced for the X-RateLimit-Limit header

+	ResetIn    time.Duration // wall-clock until the current window rolls over

+	RetryAfter time.Duration // 0 when Allowed; otherwise the wait the client should respect

+}

+

+// Allow increments the (scope, key) counter and reports whether the

+// caller is under or over the configured Max. Returns the post-

+// increment Remaining + the time until the current window rolls.

+//

+// On a Postgres error the request is allowed (fail-open). The caller

+// is expected to log the error; refusing service over a transient

+// counter glitch would be worse than the brief over-limit window.

+func (l *Limiter) Allow(ctx context.Context, p Policy, key string) (Decision, error) {

+	if p.Max <= 0 || p.Window <= 0 {

+		return Decision{}, errors.New("ratelimit: Policy.Max and Window must be positive")

+	}

+	if p.Scope == "" || key == "" {

+		return Decision{}, errors.New("ratelimit: Policy.Scope and key must be non-empty")

+	}

+	row, err := l.q.BumpRateLimit(ctx, l.pool, ratelimitdb.BumpRateLimitParams{

+		Scope: p.Scope,

+		Key:   key,

+		Ttl:   pgtype.Interval{Microseconds: int64(p.Window / time.Microsecond), Valid: true},

+	})

+	if err != nil {

+		return Decision{Allowed: true, Remaining: p.Max, Limit: p.Max, ResetIn: p.Window}, fmt.Errorf("ratelimit: bump: %w", err)

+	}

+

+	hits := int(row.Hits)

+	resetIn := time.Until(row.WindowStartedAt.Time.Add(p.Window))

+	if resetIn < 0 {

+		resetIn = 0

+	}

+	d := Decision{

+		Allowed:   hits <= p.Max,

+		Limit:     p.Max,

+		Remaining: max0(p.Max - hits),

+		ResetIn:   resetIn,

+	}

+	if !d.Allowed {

+		d.RetryAfter = resetIn

+		if d.RetryAfter <= 0 {

+			d.RetryAfter = time.Second

+		}

+	}

+	return d, nil

+}

+

+// AllowSignupIP is the inet-keyed sibling of Allow against the

+// signup_ip_throttle table. ip is masked to /24 (v4) or /48 (v6)

+// so a single residential allocation shares one counter — matches

+// GitHub's approach (per-/24 signup throttle).

+func (l *Limiter) AllowSignupIP(ctx context.Context, ip netip.Addr, max int, window time.Duration) (Decision, error) {

+	if !ip.IsValid() {

+		return Decision{}, errors.New("ratelimit: invalid ip")

+	}

+	row, err := l.q.BumpSignupIPThrottle(ctx, l.pool, ratelimitdb.BumpSignupIPThrottleParams{

+		Cidr: maskToNetwork(ip),

+		Ttl:  pgtype.Interval{Microseconds: int64(window / time.Microsecond), Valid: true},

+	})

+	if err != nil {

+		return Decision{Allowed: true, Remaining: max, Limit: max, ResetIn: window}, fmt.Errorf("ratelimit: signup bump: %w", err)

+	}

+	hits := int(row.Hits)

+	resetIn := time.Until(row.WindowStartedAt.Time.Add(window))

+	if resetIn < 0 {

+		resetIn = 0

+	}

+	d := Decision{

+		Allowed:   hits <= max,

+		Limit:     max,

+		Remaining: max0(max - hits),

+		ResetIn:   resetIn,

+	}

+	if !d.Allowed {

+		d.RetryAfter = resetIn

+		if d.RetryAfter <= 0 {

+			d.RetryAfter = time.Second

+		}

+	}

+	return d, nil

+}

+

+func max0(n int) int {

+	if n < 0 {

+		return 0

+	}

+	return n

+}

+

+// maskToNetwork zeros the host bits of ip so per-/24 (v4) or /48 (v6)

+// throttle keys collapse to one network row. The choice of /24 + /48

+// matches GitHub's reported anti-abuse defaults.

+func maskToNetwork(ip netip.Addr) netip.Addr {

+	if ip.Is4() {

+		b := ip.As4()

+		b[3] = 0

+		return netip.AddrFrom4(b)

+	}

+	b := ip.As16()

+	for i := 6; i < 16; i++ { // zero everything past /48

+		b[i] = 0

+	}

+	return netip.AddrFrom16(b)

+}

+// SPDX-License-Identifier: AGPL-3.0-or-later

+

+package ratelimit

+

+import (

+	"net"

+	"net/http"

+	"net/netip"

+	"strings"

+)

+

+// IPKey is the canonical anonymous-request keyer: extracts the

+// client IP from r.RemoteAddr or the X-Forwarded-For header. The

+// trust-XFF flag should be set ONLY when the deployment runs

+// behind a CDN/proxy we control; otherwise an attacker can spoof

+// the header and dodge IP-keyed limits.

+func IPKey(trustForwarded bool) KeyFunc {

+	return func(r *http.Request) string {

+		if trustForwarded {

+			if v := r.Header.Get("X-Forwarded-For"); v != "" {

+				// First entry is the client; downstream proxies append.

+				if comma := strings.IndexByte(v, ','); comma > 0 {

+					v = v[:comma]

+				}

+				return "ip:" + strings.TrimSpace(v)

+			}

+		}

+		host, _, err := net.SplitHostPort(r.RemoteAddr)

+		if err != nil {

+			host = r.RemoteAddr

+		}

+		return "ip:" + host

+	}

+}

+

+// ClientIP returns the parsed client IP using the same rules as

+// IPKey. Used by signup-throttle's CIDR keying.

+func ClientIP(r *http.Request, trustForwarded bool) (netip.Addr, bool) {

+	raw := ""

+	if trustForwarded {

+		if v := r.Header.Get("X-Forwarded-For"); v != "" {

+			if comma := strings.IndexByte(v, ','); comma > 0 {

+				v = v[:comma]

+			}

+			raw = strings.TrimSpace(v)

+		}

+	}

+	if raw == "" {

+		host, _, err := net.SplitHostPort(r.RemoteAddr)

+		if err == nil {

+			raw = host

+		} else {

+			raw = r.RemoteAddr

+		}

+	}

+	addr, err := netip.ParseAddr(raw)

+	if err != nil {

+		return netip.Addr{}, false

+	}

+	return addr, true

+}

+// SPDX-License-Identifier: AGPL-3.0-or-later

+

+package ratelimit

+

+import (

+	"net/http"

+	"strconv"

+	"time"

+)

+

+// KeyFunc derives the per-request rate-limit key. Common choices:

+// remote IP (anon), user id (authed), token id (PAT), repo id+user

+// (per-repo), etc. Returning "" skips the limiter (caller should

+// only return "" when there's nothing meaningful to throttle on).

+type KeyFunc func(*http.Request) string

+

+// Middleware wraps next with a per-request Allow check. On allow,

+// X-RateLimit-* headers are stamped on the response; on deny, the

+// response is 429 with Retry-After.

+//

+// `name` is the per-route identifier surfaced in logs and lets the

+// admin observability pages disambiguate scopes that share the same

+// counter (e.g. "api:anon" used by both /api/repos and /api/users).

+func (l *Limiter) Middleware(p Policy, key KeyFunc) func(http.Handler) http.Handler {

+	return func(next http.Handler) http.Handler {

+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

+			k := ""

+			if key != nil {

+				k = key(r)

+			}

+			if k == "" {

+				next.ServeHTTP(w, r)

+				return

+			}

+			d, err := l.Allow(r.Context(), p, k)

+			if err != nil {

+				// Fail-open on counter errors — surface in logs via

+				// the wrapped handler's logging middleware. We still

+				// stamp the headers we know.

+				StampHeaders(w, d)

+				next.ServeHTTP(w, r)

+				return

+			}

+			StampHeaders(w, d)

+			if !d.Allowed {

+				w.Header().Set("Retry-After", strconv.Itoa(int(d.RetryAfter/time.Second)))

+				// Headers MUST be set before WriteHeader for the

+				// response Content-Type to land; matches the same

+				// pattern as the writeRetryAfter fix from S05.

+				w.Header().Set("Content-Type", "text/plain; charset=utf-8")

+				w.WriteHeader(http.StatusTooManyRequests)

+				_, _ = w.Write([]byte("Rate limit exceeded. Please retry after the X-RateLimit-Reset window.\n"))

+				return

+			}

+			next.ServeHTTP(w, r)

+		})

+	}

+}

+

+// StampHeaders writes the standard X-RateLimit-* headers from a

+// Decision. Exposed so handlers that run a manual Allow can apply

+// the same stamp without going through Middleware.

+func StampHeaders(w http.ResponseWriter, d Decision) {

+	if d.Limit > 0 {

+		w.Header().Set("X-RateLimit-Limit", strconv.Itoa(d.Limit))

+		w.Header().Set("X-RateLimit-Remaining", strconv.Itoa(d.Remaining))

+		w.Header().Set("X-RateLimit-Reset", strconv.Itoa(int(d.ResetIn/time.Second)))

+	}

+}