Skip to content

Commit d745f50

Browse files
authored
(NOBIDS) better ratelimit-logic (#2851)
1 parent 442048c commit d745f50

File tree

1 file changed

+9
-13
lines changed

1 file changed

+9
-13
lines changed

ratelimit/ratelimit.go

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,7 @@ type RateLimit struct {
111111
}
112112

113113
type RateLimitResult struct {
114+
BlockRequest bool
114115
Time time.Time
115116
Weight int64
116117
Route string
@@ -323,7 +324,7 @@ func HttpMiddleware(next http.Handler) http.Handler {
323324
w.Header().Set(HeaderRateLimitRemainingMinute, strconv.FormatInt(rl.RemainingMinute, 10))
324325
w.Header().Set(HeaderRateLimitRemainingSecond, strconv.FormatInt(rl.RemainingSecond, 10))
325326

326-
if rl.Weight > rl.Remaining {
327+
if rl.BlockRequest {
327328
w.Header().Set(HeaderRetryAfter, strconv.FormatInt(rl.Reset, 10))
328329
http.Error(w, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests)
329330
err = postRateLimit(rl, http.StatusTooManyRequests)
@@ -805,16 +806,19 @@ func rateLimitRequest(r *http.Request) (*RateLimitResult, error) {
805806
res.Remaining = 0
806807
res.Reset = int64(timeUntilNextMonthUtc.Seconds())
807808
res.Window = MonthTimeWindow
809+
res.BlockRequest = true
808810
} else if res.RateLimit.Hour > 0 && rateLimitHour.Val() > res.RateLimit.Hour {
809811
res.Limit = res.RateLimit.Hour
810812
res.Remaining = 0
811813
res.Reset = int64(timeUntilNextHourUtc.Seconds())
812814
res.Window = HourTimeWindow
815+
res.BlockRequest = true
813816
} else if res.RateLimit.Second > 0 && rateLimitSecond.Val() > res.RateLimit.Second {
814817
res.Limit = res.RateLimit.Second
815818
res.Remaining = 0
816819
res.Reset = int64(1)
817820
res.Window = SecondTimeWindow
821+
res.BlockRequest = true
818822
} else {
819823
res.Limit = res.RateLimit.Second
820824
res.Remaining = res.RateLimit.Second - rateLimitSecond.Val()
@@ -846,34 +850,26 @@ func rateLimitRequest(r *http.Request) (*RateLimitResult, error) {
846850
res.LimitMonth = res.RateLimit.Month
847851
} else {
848852
res.LimitMonth = max(res.RateLimit.Month, res.RateLimit.Hour, res.RateLimit.Second)
853+
res.RemainingMonth = max(res.RemainingMonth, res.RemainingHour, res.RemainingSecond)
849854
}
850855
res.LimitDay = res.LimitMonth
856+
res.RemainingDay = res.RemainingMonth
851857

852858
if res.RateLimit.Hour > 0 {
853859
res.LimitHour = res.RateLimit.Hour
854860
} else {
855861
res.LimitHour = res.LimitMonth
862+
res.RemainingHour = res.RemainingMonth
856863
}
857864
res.LimitMinute = res.LimitHour
865+
res.RemainingMinute = res.RemainingHour
858866

859867
if res.RateLimit.Second > 0 {
860868
res.LimitSecond = res.RateLimit.Second
861869
} else {
862870
res.LimitSecond = res.LimitHour
863871
}
864872

865-
if res.RemainingMonth == 0 {
866-
res.RemainingMonth = max(res.RemainingMonth, res.RemainingHour, res.RemainingSecond)
867-
}
868-
res.RemainingDay = res.RemainingMonth
869-
if res.RemainingHour == 0 {
870-
res.RemainingHour = res.RemainingMonth
871-
}
872-
res.RemainingMinute = res.RemainingHour
873-
if res.RemainingSecond == 0 {
874-
res.RemainingSecond = res.RemainingMinute
875-
}
876-
877873
return res, nil
878874
}
879875

0 commit comments

Comments
 (0)