forked from kataras/iris
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.go
90 lines (74 loc) · 2.44 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
package main
import (
"time"
"github.com/kataras/iris/v12"
"github.com/kataras/iris/v12/middleware/rate"
)
func main() {
app := newApp()
app.Logger().SetLevel("debug")
// * http://localhost:8080/v1
// * http://localhost:8080/v1/other
// * http://localhost:8080/v2/list (with X-API-Key request header)
// Read more at: https://en.wikipedia.org/wiki/Token_bucket
//
// Alternatives:
// * https://github.com/iris-contrib/middleware/blob/master/throttler/_example/main.go
// Read more at: https://en.wikipedia.org/wiki/Generic_cell_rate_algorithm
// * https://github.com/iris-contrib/middleware/tree/master/tollboothic/_examples/limit-handler
app.Listen(":8080")
}
func newApp() *iris.Application {
app := iris.New()
v1 := app.Party("/v1")
{
// Register the rate limiter middleware at the "/v1" subrouter.
//
// Fist and second input parameters:
// Allow 1 request per second, with a maximum burst size of 5.
//
// Third optional variadic input parameter:
// Can be a cleanup function.
// Iris provides a cleanup function that will check for old entries and remove them.
// You can customize it, e.g. check every 1 minute
// if a client's last visit was 5 minutes ago ("old" entry)
// and remove it from the memory.
limitV1 := rate.Limit(1, 5, rate.PurgeEvery(time.Minute, 5*time.Minute))
// rate.Every helper:
// rate.Limit(rate.Every(time.Minute), 5)
v1.Use(limitV1)
v1.Get("/", index)
v1.Get("/other", other)
}
v2 := app.Party("/v2")
{
v2.Use(useAPIKey)
// Initialize a new rate limit middleware to limit requests
// per API Key(see `useAPIKey` below) instead of client's Remote IP Address.
limitV2 := rate.Limit(rate.Every(time.Minute), 300, rate.PurgeEvery(5*time.Minute, 15*time.Minute))
v2.Use(limitV2)
v2.Get("/list", list)
}
return app
}
func useAPIKey(ctx iris.Context) {
apiKey := ctx.GetHeader("X-API-Key")
if apiKey == "" { // [validate your API Key here...]
ctx.StopWithStatus(iris.StatusForbidden)
return
}
// Change the method that rate limit matches the requests with a specific user
// and set our own api key as theirs identifier.
rate.SetIdentifier(ctx, apiKey)
ctx.Next()
}
func list(ctx iris.Context) {
ctx.JSON(iris.Map{"key": "value"})
}
func index(ctx iris.Context) {
ctx.HTML("<h1>Index Page</h1>")
}
func other(ctx iris.Context) {
ctx.HTML("<h1>Other Page</h1>")
}
// Note: Use `ctx.SendFileWithRate` to use a download rate limiter instead.