diff --git a/checkpoint.go b/checkpoint.go index 4f3bf4d..4714ee7 100644 --- a/checkpoint.go +++ b/checkpoint.go @@ -14,3 +14,7 @@ func LimitByIp(maxTokens, refillRate, tokensPerRefill int) func(next http.Handle func LimitIpByEndpoint(maxTokens, refillRate, tokensPerRefill int) func(next http.Handler) http.Handler { return Limit(maxTokens, refillRate, tokensPerRefill) } + +func WithConfig(config Config) func(next http.Handler) http.Handler { + return config.LimitMethod(config.MaxTokens, config.RefillRate, config.TokensPerRefill) +} diff --git a/cmd/main.go b/cmd/main.go index e32918a..f6c177b 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -3,18 +3,30 @@ package main import ( "fmt" "net/http" + _ "net/http/pprof" checkpoint "github.com/aidenfine/checkpoint" ) func main() { - mux := http.NewServeMux() - mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { fmt.Fprintln(w, "Hello World") }) - rlMiddleware := checkpoint.LimitByIp(20, 1, 1) - rlHandler := rlMiddleware(mux) + // CONFIG METHOD == + config := checkpoint.Config{ + IgnorePaths: []string{"/logs"}, + MaxTokens: 25, + RefillRate: 1, + TokensPerRefill: 1, + LimitMethod: checkpoint.LimitByIp, + } + + // Quick method + // rlMiddleware := checkpoint.LimitByIp(25, 1, 1) + + rlMiddleware := checkpoint.WithConfig(config) + rlHandler := rlMiddleware(http.DefaultServeMux) http.ListenAndServe(":8080", rlHandler) } diff --git a/config.go b/config.go new file mode 100644 index 0000000..f0bd393 --- /dev/null +++ b/config.go @@ -0,0 +1,11 @@ +package checkpoint + +import "net/http" + +type Config struct { + IgnorePaths []string `json:"ignorePaths" yaml:"ignorePaths"` + MaxTokens int `json:"maxTokens" yaml:"maxTokens"` + RefillRate int `json:"refillRate" yaml:"refillRate"` + TokensPerRefill int `json:"tokensPerRefill" yaml:"tokensPerRefill"` + LimitMethod func(maxTokens, refillRate, tokensPerRefill int) func(next http.Handler) http.Handler +} diff --git a/config_test.go b/config_test.go new file mode 100644 index 0000000..e50808c --- /dev/null +++ b/config_test.go @@ -0,0 +1 @@ +package checkpoint_test