52 lines
1.1 KiB
Go
52 lines
1.1 KiB
Go
|
package rate
|
||
|
|
|
||
|
|
import (
|
||
|
|
"net/http"
|
||
|
|
"strconv"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
type Category int
|
||
|
|
|
||
|
|
const (
|
||
|
|
CheckCategory Category = iota
|
||
|
|
ReportsCategory
|
||
|
|
BlacklistCategory
|
||
|
|
ReportCategory
|
||
|
|
CheckBlockCategory
|
||
|
|
BulkReportCategory
|
||
|
|
ClearAddressCategory
|
||
|
|
)
|
||
|
|
|
||
|
|
type Limit struct {
|
||
|
|
RetryAfter time.Time `header:"Retry-After"`
|
||
|
|
XRateLimitLimit int `header:"X-RateLimit-Limit"`
|
||
|
|
XRateLimitRemaining int `header:"X-RateLimit-Remaining"`
|
||
|
|
XRateLimitReset time.Time `header:"X-RateLimit-Reset"`
|
||
|
|
}
|
||
|
|
|
||
|
|
type Limits map[Category]*Limit
|
||
|
|
|
||
|
|
func (l *Limits) Update(category Category, header http.Header) error {
|
||
|
|
limit := &Limit{}
|
||
|
|
for key, value := range header {
|
||
|
|
iValue, err := strconv.Atoi(value[0])
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
switch key {
|
||
|
|
case "X-RateLimit-Limit":
|
||
|
|
limit.XRateLimitLimit = iValue
|
||
|
|
case "X-RateLimit-Remaining":
|
||
|
|
limit.XRateLimitRemaining = iValue
|
||
|
|
case "X-RateLimit-Reset":
|
||
|
|
limit.XRateLimitReset = time.Unix(int64(iValue), 0)
|
||
|
|
case "Retry-After":
|
||
|
|
limit.RetryAfter = time.Now().Add(time.Duration(iValue) * time.Second)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
(*l)[category] = limit
|
||
|
|
return nil
|
||
|
|
}
|