-
Notifications
You must be signed in to change notification settings - Fork 2
/
webserver.go
118 lines (103 loc) · 3.23 KB
/
webserver.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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
package main
import (
"fmt"
"net/http"
"strings"
"time"
"github.com/zekroTJA/ratelimit"
)
const (
limiterLimit = 10 * time.Second
limiterBurst = 3
)
// webServer contains the actual HTTP server
// instance, the ServeMux handling roots and
// a map of limiters.
type webServer struct {
s *http.Server
mux *http.ServeMux
// limiters is a map binding a limiter
// to a remote address, so a limiter only
// counts for one connection
limiters map[string]*ratelimit.Limiter
}
// newWebServer creates a new instance of
// a webServer and adds the root handlers
func newWebServer(addr string) *webServer {
// Creating a new webServer with a new
// Server, a new ServeMux and the
// initialized limiters map.
ws := &webServer{
s: &http.Server{
Addr: addr,
},
mux: http.NewServeMux(),
limiters: make(map[string]*ratelimit.Limiter),
}
// Setting handlerTest as handler for
// root /api/test
ws.mux.HandleFunc("/api/test", ws.handlerTest)
// Adding ServeMux instance as Handler
// for the HTTP Server
ws.s.Handler = ws.mux
return ws
}
// start the web server blocking the current
// thread and retuning an error if it fails.
func (ws *webServer) start() error {
return ws.s.ListenAndServe()
}
// checkLimit is a helper function checking
// the availability of a limiter for the connections
// address or creating it if not existing. Then,
// the availability of tokens will be checked
// to perform an action. This state will be returned
// as boolean.
func (ws *webServer) checkLimit(w http.ResponseWriter, r *http.Request) bool {
// Getting the address of the incomming connection.
// Because you will likely test this with a local connection,
// the local port number will be attached and differ on every
// request. So, we need to cut away everything behind the last
// ":", if existent.
addr := r.RemoteAddr
if strings.Contains(addr, ":") {
split := strings.Split(addr, ":")
addr = strings.Join(split[0:len(split)-1], ":")
}
// Getting the limiter for the current connections address
// or create one if not existent.
limiter, ok := ws.limiters[addr]
if !ok {
limiter = ratelimit.NewLimiter(limiterLimit, limiterBurst)
ws.limiters[addr] = limiter
}
// Reserve a token from the limiter.
a, res := limiter.Reserve()
// Attach the reservation result to the three headers
// "X-RateLimit-Limit"
// - containing the absolute burst rate
// of the limiter,
// "X-RateLimit-Remaining"
// - the number of remaining tickets after
// the request
// "X-RateLimit-Reset"
// - the UnixNano timestamp until a new token
// will be generated (only if remaining == 0)
w.Header().Set("X-RateLimit-Limit", fmt.Sprintf("%d", res.Burst))
w.Header().Set("X-RateLimit-Remaining", fmt.Sprintf("%d", res.Remaining))
w.Header().Set("X-RateLimit-Reset", fmt.Sprintf("%d", res.Reset.UnixNano()))
// Return the succeed status of
// the token request
return a
}
// handlerTest is the handler for /api/test root
func (ws *webServer) handlerTest(w http.ResponseWriter, r *http.Request) {
// Check and consume a token from the limiter,
// if available. If succeed, return status 200,
// else status 429.
if ws.checkLimit(w, r) {
w.WriteHeader(http.StatusOK)
} else {
w.WriteHeader(http.StatusTooManyRequests)
}
}