57 lines
1.8 KiB
Go
57 lines
1.8 KiB
Go
package app
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"fmt"
|
|
"net"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
func (a *App) allowRequest(ctx context.Context, action, subject string, limit int, window time.Duration) bool {
|
|
if limit < 1 || window < time.Second {
|
|
return false
|
|
}
|
|
windowSeconds := int64(window / time.Second)
|
|
slot := time.Now().Unix() / windowSeconds
|
|
key := sha256.Sum256([]byte(fmt.Sprintf("%s|%s|%d", action, subject, slot)))
|
|
expiresAt := time.Unix((slot+1)*windowSeconds, 0).Add(time.Minute)
|
|
_, err := a.db.ExecContext(ctx, `INSERT INTO api_rate_limits(bucket_key,action_name,hits,expires_at) VALUES(?,?,1,?) ON DUPLICATE KEY UPDATE hits=hits+1,expires_at=VALUES(expires_at)`, key[:], action, expiresAt)
|
|
if err != nil {
|
|
return false
|
|
}
|
|
var hits int
|
|
if err = a.db.QueryRowContext(ctx, `SELECT hits FROM api_rate_limits WHERE bucket_key=?`, key[:]).Scan(&hits); err != nil {
|
|
return false
|
|
}
|
|
if key[0] == 0 {
|
|
_, _ = a.db.ExecContext(ctx, `DELETE FROM api_rate_limits WHERE expires_at<NOW(3) LIMIT 1000`)
|
|
}
|
|
return hits <= limit
|
|
}
|
|
|
|
func (a *App) rateLimit(w http.ResponseWriter, r *http.Request, action, subject string, limit int, window time.Duration) bool {
|
|
if a.allowRequest(r.Context(), action, subject, limit, window) {
|
|
return true
|
|
}
|
|
w.Header().Set("Retry-After", fmt.Sprintf("%d", int(window/time.Second)))
|
|
fail(w, http.StatusTooManyRequests, 20002, "请求过于频繁，请稍后再试")
|
|
return false
|
|
}
|
|
|
|
func clientIP(r *http.Request) string {
|
|
if forwarded := strings.TrimSpace(strings.Split(r.Header.Get("X-Forwarded-For"), ",")[0]); net.ParseIP(forwarded) != nil {
|
|
return forwarded
|
|
}
|
|
if realIP := strings.TrimSpace(r.Header.Get("X-Real-IP")); net.ParseIP(realIP) != nil {
|
|
return realIP
|
|
}
|
|
host, _, err := net.SplitHostPort(strings.TrimSpace(r.RemoteAddr))
|
|
if err == nil && net.ParseIP(host) != nil {
|
|
return host
|
|
}
|
|
return "unknown"
|
|
}
|