Skip to content

Commit

Permalink
fix: don't use ip or key in redis
Browse files Browse the repository at this point in the history
  • Loading branch information
airenas committed May 22, 2024
1 parent 873124d commit f4d3d61
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 3 deletions.
13 changes: 10 additions & 3 deletions internal/pkg/handler/rateLimit.go
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ func (h *rateLimitValidate) ServeHTTP(w http.ResponseWriter, r *http.Request) {
if limit <= 0 {
limit = h.limit
}
ok, rem, retryAfter, err := h.qv.Validate(makeRateLimitKey(ctx.Key, ctx.Manual), int64(limit), int64(quotaV))
ok, rem, retryAfter, err := h.qv.Validate(makeRateLimitKey(idOrHash(ctx), ctx.Manual), int64(limit), int64(quotaV))
if err != nil {
http.Error(w, "Service error", http.StatusInternalServerError)
goapp.Log.Error("can't validate rate limit.", err)
Expand All @@ -56,8 +56,15 @@ func (h *rateLimitValidate) ServeHTTP(w http.ResponseWriter, r *http.Request) {
h.next.ServeHTTP(w, rn)
}

func makeRateLimitKey(s string, b bool) string {
return fmt.Sprintf("%s_%t", s, b)
func idOrHash(ctx *customData) string {
if ctx.KeyID != "" {
return ctx.KeyID
}
return hashKey(ctx.Key)
}

func makeRateLimitKey(key string, manual bool) string {
return fmt.Sprintf("%s:%t", key, manual)
}

func (h *rateLimitValidate) Info(pr string) string {
Expand Down
24 changes: 24 additions & 0 deletions internal/pkg/handler/rateLimit_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
package handler

import "testing"

func Test_idOrHash(t *testing.T) {
type args struct {
ctx *customData
}
tests := []struct {
name string
args args
want string
}{
{name: "id", args: args{ctx: &customData{KeyID: "id", Key: "olia"}}, want: "id"},
{name: "key", args: args{ctx: &customData{KeyID: "", Key: "olia"}}, want: "d57329cf35"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := idOrHash(tt.args.ctx); got != tt.want {
t.Errorf("idOrHash() = %v, want %v", got, tt.want)
}
})
}
}

0 comments on commit f4d3d61

Please sign in to comment.