/
middleware.go
106 lines (84 loc) · 2.2 KB
/
middleware.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
package ratelimiter
import (
"errors"
"fmt"
"log"
"net"
"strings"
"github.com/gin-gonic/gin"
)
func GetIPAddress(c *gin.Context) (string, error) {
//Get IP from the X-REAL-IP header
ip := c.Request.Header.Get("X-REAL-IP")
// fmt.Printf("Here's the raw ip ==> %v", ip)
netIP := net.ParseIP(ip)
if netIP != nil {
return ip, nil
}
//Get IP from X-FORWARDED-FOR header
ips := c.Request.Header.Get("X-FORWARDED-FOR")
// fmt.Printf("Here's the raw ip ==> %v", ip)
splitIps := strings.Split(ips, ",")
for _, ip := range splitIps {
netIP := net.ParseIP(ip)
if netIP != nil {
return ip, nil
}
}
//Get IP from RemoteAddr
ip, _, err := net.SplitHostPort(c.Request.RemoteAddr)
// fmt.Printf("Here's the raw ip ==> %v", ip)
if err != nil {
return "", err
}
netIP = net.ParseIP(ip)
if netIP != nil {
return ip, nil
}
return "", fmt.Errorf("No valid ip found")
}
// enforce rate limit based on ip address of client making the request
func IPLimiter(limiter RateLimiter) gin.HandlerFunc {
return func(c *gin.Context) {
ip_addr, err := GetIPAddress(c)
if err != nil {
log.Println(err)
c.AbortWithError(401, errors.New("IP Adress Unknown"))
return
}
if !limiter.AllowRequest(ip_addr) {
c.AbortWithError(429, errors.New("Too Many Requests Made, Retry Later"))
return
}
limiter.UpdateRequest(ip_addr)
c.Next()
}
}
// enforce rate limit based on the user making the request
func UserLimiter(limiter RateLimiter, user_id string) gin.HandlerFunc {
return func(c *gin.Context) {
if !limiter.AllowRequest(user_id) {
c.AbortWithError(429, errors.New("Too Many Requests Made, Retry Later"))
return
}
limiter.UpdateRequest(user_id)
c.Next()
}
}
// enforce rate limit based on ip address of client and user making the request
func IPUserLimiter(limiter RateLimiter, user_id string) gin.HandlerFunc {
return func(c *gin.Context) {
ip_addr, err := GetIPAddress(c)
if err != nil {
log.Println(err)
c.AbortWithError(401, errors.New("IP Adress Unknown"))
return
}
if !limiter.AllowRequest(ip_addr, user_id) {
c.AbortWithError(429, errors.New("Too Many Requests Made, Retry Later"))
return
}
limiter.UpdateRequest(ip_addr, user_id)
c.Next()
}
}