forked from FabianVowie/Lithium
Roman Zipp
3 years ago
11 changed files with 230 additions and 36 deletions
-
50README.md
-
3build.sh
-
3go.mod
-
21go.sum
-
34main.go
-
2middlewares/authorization.go
-
2middlewares/authorization_test.go
-
43middlewares/ratelimiter.go
-
64middlewares/ratelimiter_test.go
-
28settings/settings.go
-
12settings/settings_test.go
@ -0,0 +1,3 @@ |
|||
#!/bin/bash |
|||
|
|||
GIT_COMMIT=$(git rev-parse --short HEAD); go build -ldflags "-X main.GitCommit=$GIT_COMMIT" |
|||
@ -1,4 +1,4 @@ |
|||
package auth |
|||
package middlewares |
|||
|
|||
import ( |
|||
"net/http" |
|||
@ -1,4 +1,4 @@ |
|||
package auth |
|||
package middlewares |
|||
|
|||
import ( |
|||
"net/http" |
|||
@ -0,0 +1,43 @@ |
|||
package middlewares |
|||
|
|||
import ( |
|||
"net/http" |
|||
|
|||
"github.com/throttled/throttled/store/memstore" |
|||
"github.com/throttled/throttled/v2" |
|||
) |
|||
|
|||
type RateLimiterMiddleware struct { |
|||
rateLimiter throttled.HTTPRateLimiter |
|||
} |
|||
|
|||
func (middleware RateLimiterMiddleware) Middleware(next http.Handler) http.Handler { |
|||
return middleware.rateLimiter.RateLimit(next) |
|||
} |
|||
|
|||
func CreateRateLimiterMiddleware(requestsPerMinute int, allowedBurst int) (*RateLimiterMiddleware, error) { |
|||
store, err := memstore.New(65536) |
|||
|
|||
if err != nil { |
|||
return nil, err |
|||
} |
|||
|
|||
quota := throttled.RateQuota{ |
|||
MaxRate: throttled.PerMin(requestsPerMinute), |
|||
MaxBurst: allowedBurst, |
|||
} |
|||
|
|||
rateLimiter, err := throttled.NewGCRARateLimiter(store, quota) |
|||
if err != nil { |
|||
return nil, err |
|||
} |
|||
|
|||
httpRateLimiter := throttled.HTTPRateLimiter{ |
|||
RateLimiter: rateLimiter, |
|||
VaryBy: &throttled.VaryBy{Path: true}, |
|||
} |
|||
|
|||
return &RateLimiterMiddleware{ |
|||
rateLimiter: httpRateLimiter, |
|||
}, nil |
|||
} |
|||
@ -0,0 +1,64 @@ |
|||
package middlewares |
|||
|
|||
import ( |
|||
"net/http" |
|||
"net/http/httptest" |
|||
"testing" |
|||
|
|||
"github.com/stretchr/testify/assert" |
|||
) |
|||
|
|||
func ExecuteRequest(middlewareHandler http.Handler) int { |
|||
request, _ := http.NewRequest("GET", "/", nil) |
|||
responseRecorder := httptest.NewRecorder() |
|||
|
|||
middlewareHandler.ServeHTTP(responseRecorder, request) |
|||
|
|||
return responseRecorder.Code |
|||
} |
|||
|
|||
func TestRateLimiterMiddleware(t *testing.T) { |
|||
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
|||
w.WriteHeader(http.StatusOK) |
|||
}) |
|||
|
|||
t.Run("AuthorizationMiddleware returns 200 response when rate limit is not hit", func(t *testing.T) { |
|||
middleware, err := CreateRateLimiterMiddleware(1, 0) |
|||
assert.Nil(t, err) |
|||
|
|||
middlewareHandler := middleware.Middleware(handler) |
|||
|
|||
assert.Equal(t, 200, ExecuteRequest(middlewareHandler)) |
|||
}) |
|||
|
|||
t.Run("AuthorizationMiddleware returns 429 response when rate limit is hit", func(t *testing.T) { |
|||
middleware, err := CreateRateLimiterMiddleware(1, 0) |
|||
assert.Nil(t, err) |
|||
|
|||
middlewareHandler := middleware.Middleware(handler) |
|||
|
|||
assert.Equal(t, 200, ExecuteRequest(middlewareHandler)) |
|||
assert.Equal(t, 429, ExecuteRequest(middlewareHandler)) |
|||
}) |
|||
|
|||
t.Run("AuthorizationMiddleware returns 200 response when rate limit with burst is not hit", func(t *testing.T) { |
|||
middleware, err := CreateRateLimiterMiddleware(1, 1) |
|||
assert.Nil(t, err) |
|||
|
|||
middlewareHandler := middleware.Middleware(handler) |
|||
|
|||
assert.Equal(t, 200, ExecuteRequest(middlewareHandler)) |
|||
assert.Equal(t, 200, ExecuteRequest(middlewareHandler)) |
|||
}) |
|||
|
|||
t.Run("AuthorizationMiddleware returns 429 response when rate limit with burst is hit", func(t *testing.T) { |
|||
middleware, err := CreateRateLimiterMiddleware(1, 1) |
|||
assert.Nil(t, err) |
|||
|
|||
middlewareHandler := middleware.Middleware(handler) |
|||
|
|||
assert.Equal(t, 200, ExecuteRequest(middlewareHandler)) |
|||
assert.Equal(t, 200, ExecuteRequest(middlewareHandler)) |
|||
assert.Equal(t, 429, ExecuteRequest(middlewareHandler)) |
|||
}) |
|||
} |
|||
Write
Preview
Loading…
Cancel
Save
Reference in new issue