// Ip请求频率校验
func RateLimitHandler(limiter *config.Limiter, next http.Handler) http.Handler {
middle := func(w http.ResponseWriter, r *http.Request) {
tollbooth.SetResponseHeaders(limiter, w)
// 检查Ip频率
remoteIP := libstring.RemoteIP(limiter.IPLookups, r)
if remoteIP != "127.0.0.1" {
// 频率校验
httpError := tollbooth.LimitByRequest(limiter, r)
if httpError != nil {
// 同步到日志
RateLogger.Info(remoteIP)
// 返回429
initCors(w, r)
w.Header().Add("Content-Type", limiter.MessageContentType)
w.WriteHeader(httpError.StatusCode)
w.Write([]byte(httpError.Message))
return
}
}
// 正常访问
next.ServeHTTP(w, r)
}
return http.HandlerFunc(middle)
}
// BuildKeys generates a slice of keys to rate-limit by given config and request structs.
func BuildKeys(limiter *config.Limiter, r *http.Request) [][]string {
remoteIP := libstring.RemoteIP(limiter.IPLookups, r)
path := r.URL.Path
sliceKeys := make([][]string, 0)
// Don't BuildKeys if remoteIP is blank.
if remoteIP == "" {
return sliceKeys
}
if limiter.Methods != nil && limiter.Headers != nil && limiter.BasicAuthUsers != nil {
// Limit by HTTP methods and HTTP headers+values and Basic Auth credentials.
if libstring.StringInSlice(limiter.Methods, r.Method) {
for headerKey, headerValues := range limiter.Headers {
if (headerValues == nil || len(headerValues) <= 0) && r.Header.Get(headerKey) != "" {
// If header values are empty, rate-limit all request with headerKey.
username, _, ok := r.BasicAuth()
if ok && libstring.StringInSlice(limiter.BasicAuthUsers, username) {
sliceKeys = append(sliceKeys, []string{remoteIP, path, r.Method, headerKey, username})
}
} else if len(headerValues) > 0 && r.Header.Get(headerKey) != "" {
// If header values are not empty, rate-limit all request with headerKey and headerValues.
for _, headerValue := range headerValues {
username, _, ok := r.BasicAuth()
if ok && libstring.StringInSlice(limiter.BasicAuthUsers, username) {
sliceKeys = append(sliceKeys, []string{remoteIP, path, r.Method, headerKey, headerValue, username})
}
}
}
}
}
} else if limiter.Methods != nil && limiter.Headers != nil {
// Limit by HTTP methods and HTTP headers+values.
if libstring.StringInSlice(limiter.Methods, r.Method) {
for headerKey, headerValues := range limiter.Headers {
if (headerValues == nil || len(headerValues) <= 0) && r.Header.Get(headerKey) != "" {
// If header values are empty, rate-limit all request with headerKey.
sliceKeys = append(sliceKeys, []string{remoteIP, path, r.Method, headerKey})
} else if len(headerValues) > 0 && r.Header.Get(headerKey) != "" {
// If header values are not empty, rate-limit all request with headerKey and headerValues.
for _, headerValue := range headerValues {
sliceKeys = append(sliceKeys, []string{remoteIP, path, r.Method, headerKey, headerValue})
}
}
}
}
} else if limiter.Methods != nil && limiter.BasicAuthUsers != nil {
// Limit by HTTP methods and Basic Auth credentials.
if libstring.StringInSlice(limiter.Methods, r.Method) {
username, _, ok := r.BasicAuth()
if ok && libstring.StringInSlice(limiter.BasicAuthUsers, username) {
sliceKeys = append(sliceKeys, []string{remoteIP, path, r.Method, username})
}
}
} else if limiter.Methods != nil {
// Limit by HTTP methods.
if libstring.StringInSlice(limiter.Methods, r.Method) {
sliceKeys = append(sliceKeys, []string{remoteIP, path, r.Method})
}
} else if limiter.Headers != nil {
// Limit by HTTP headers+values.
for headerKey, headerValues := range limiter.Headers {
if (headerValues == nil || len(headerValues) <= 0) && r.Header.Get(headerKey) != "" {
// If header values are empty, rate-limit all request with headerKey.
sliceKeys = append(sliceKeys, []string{remoteIP, path, headerKey})
} else if len(headerValues) > 0 && r.Header.Get(headerKey) != "" {
// If header values are not empty, rate-limit all request with headerKey and headerValues.
for _, headerValue := range headerValues {
sliceKeys = append(sliceKeys, []string{remoteIP, path, headerKey, headerValue})
}
}
}
} else if limiter.BasicAuthUsers != nil {
// Limit by Basic Auth credentials.
username, _, ok := r.BasicAuth()
if ok && libstring.StringInSlice(limiter.BasicAuthUsers, username) {
sliceKeys = append(sliceKeys, []string{remoteIP, path, username})
}
} else {
// Default: Limit by remoteIP and path.
sliceKeys = append(sliceKeys, []string{remoteIP, path})
}
return sliceKeys
}
