blob: 8fd6cd4578dfcc3b47f948a426827731508d0af9 [file] [log] [blame]
// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package middleware
import (
"context"
"crypto/hmac"
"crypto/sha256"
"errors"
"io"
"net"
"net/http"
"strings"
"time"
"github.com/go-redis/redis/v8"
rrate "github.com/go-redis/redis_rate/v9"
"go.opencensus.io/stats"
"go.opencensus.io/stats/view"
"go.opencensus.io/tag"
"golang.org/x/pkgsite/internal"
"golang.org/x/pkgsite/internal/config"
"golang.org/x/pkgsite/internal/log"
)
var (
keyQuotaBlocked = tag.MustNewKey("quota.blocked")
quotaResults = stats.Int64(
"go-discovery/quota_result_count",
"The result of a quota check.",
stats.UnitDimensionless,
)
// QuotaResultCount is a counter of quota results, by whether the request was blocked or not.
QuotaResultCount = &view.View{
Name: "go-discovery/quota/result_count",
Measure: quotaResults,
Aggregation: view.Count(),
Description: "quota results, by blocked or allowed",
TagKeys: []tag.Key{keyQuotaBlocked},
}
)
func recordQuotaMetric(ctx context.Context, blocked string) {
stats.RecordWithTags(ctx, []tag.Mutator{
tag.Upsert(keyQuotaBlocked, blocked),
}, quotaResults.M(1))
}
func ipKey(s string) string {
addr, _, _ := internal.Cut(s, ",")
// First field is the originating IP address.
origin := strings.TrimSpace(addr)
ip := net.ParseIP(origin)
if ip == nil {
return ""
}
// Zero out last byte, to cover ranges of IPv4 addresses. (It's less clear what the
// effect will be on IPv6 addresses: it will certainly cover a range of them, but we
// don't know if that range is likely to be allocated to a single organization.)
ip[len(ip)-1] = 0
return ip.String()
}
// Quota implements a simple IP-based rate limiter. Each set of incoming IP
// addresses with the same low-order byte gets settings.QPS requests per second.
//
// Information is kept in a redis instance.
//
// If a request is disallowed, a 429 (TooManyRequests) will be served.
func Quota(settings config.QuotaSettings, client *redis.Client) Middleware {
return func(h http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
ctx := r.Context()
if !settings.Enable {
recordQuotaMetric(ctx, "disabled")
h.ServeHTTP(w, r)
return
}
authVal := r.Header.Get(config.BypassQuotaAuthHeader)
for _, wantVal := range settings.AuthValues {
if authVal == wantVal {
recordQuotaMetric(ctx, "bypassed")
log.Infof(ctx, "Quota: accepting %q", authVal)
h.ServeHTTP(w, r)
return
}
}
header := r.Header.Get("X-Godoc-Forwarded-For")
if header == "" {
header = r.Header.Get("X-Forwarded-For")
}
blocked, reason := enforceQuota(ctx, client, settings.QPS, header, settings.HMACKey)
recordQuotaMetric(ctx, reason)
if blocked && settings.RecordOnly != nil && !*settings.RecordOnly {
const tmr = http.StatusTooManyRequests
http.Error(w, http.StatusText(tmr), tmr)
return
}
h.ServeHTTP(w, r)
})
}
}
func enforceQuota(ctx context.Context, client *redis.Client, qps int, header string, hmacKey []byte) (blocked bool, reason string) {
// Fail open if header is missing or can't be parsed.
if header == "" {
return false, "no header"
}
key := ipKey(header)
if key == "" {
return false, "bad header"
}
mac := hmac.New(sha256.New, hmacKey)
io.WriteString(mac, key)
rrateKey := string(mac.Sum(nil))
res, err := rrate.NewLimiter(client.WithTimeout(15*time.Millisecond)).Allow(ctx, rrateKey, rrate.PerSecond(qps))
if err != nil {
var nerr *net.OpError
if errors.Is(err, context.DeadlineExceeded) || (errors.As(err, &nerr) && nerr.Timeout()) {
log.Warningf(ctx, "quota: redis limiter: %v", err)
return false, "timeout"
}
log.Errorf(ctx, "quota: redis limiter: %v", err)
return false, "error"
}
if res.Allowed > 0 {
return false, "allowed"
}
return true, "blocked"
}