53 lines
1.4 KiB
Go
53 lines
1.4 KiB
Go
package handler
|
|
|
|
import (
|
|
"fmt"
|
|
"net/http"
|
|
"strconv"
|
|
"time"
|
|
|
|
"git.ilapage.cn/OPC/chorus/internal/platform/ratelimit"
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func (h *Handler) limitPortalSubmission(c *gin.Context) {
|
|
state := currentSession(c)
|
|
decision := h.rateLimiter.Take(h.now(), ratelimit.Bucket{
|
|
Key: fmt.Sprintf("submission:user:%d", state.UserID), Policy: h.userRatePolicy,
|
|
})
|
|
if !decision.Allowed {
|
|
h.writeRateLimit(c, decision.RetryAt)
|
|
return
|
|
}
|
|
c.Next()
|
|
}
|
|
|
|
func (h *Handler) limitOpenAPIRequest(c *gin.Context) {
|
|
principal := currentAPIPrincipal(c)
|
|
buckets := []ratelimit.Bucket{{
|
|
Key: fmt.Sprintf("openapi:key:%d", principal.APIKeyID), Policy: h.apiKeyRatePolicy,
|
|
}}
|
|
if c.Request.Method == http.MethodPost && (c.FullPath() == "/openapi/v1/generations/text" || c.FullPath() == "/openapi/v1/generations/image") {
|
|
buckets = append(buckets, ratelimit.Bucket{
|
|
Key: fmt.Sprintf("submission:user:%d", principal.UserID), Policy: h.userRatePolicy,
|
|
})
|
|
}
|
|
decision := h.rateLimiter.Take(h.now(), buckets...)
|
|
if !decision.Allowed {
|
|
h.writeRateLimit(c, decision.RetryAt)
|
|
return
|
|
}
|
|
c.Next()
|
|
}
|
|
|
|
func (h *Handler) writeRateLimit(c *gin.Context, retryAt time.Time) {
|
|
retryAfter := retryAt.Sub(h.now())
|
|
seconds := int64((retryAfter + time.Second - 1) / time.Second)
|
|
if seconds < 1 {
|
|
seconds = 1
|
|
}
|
|
c.Header("Retry-After", strconv.FormatInt(seconds, 10))
|
|
writeError(c, http.StatusTooManyRequests, "rate_limited", "request rate limit exceeded")
|
|
c.Abort()
|
|
}
|