Skip to content

Commit

Permalink
Add rate limiting
Browse files Browse the repository at this point in the history
  • Loading branch information
bbedward committed Sep 2, 2022
1 parent 0a12290 commit ab3e8a3
Show file tree
Hide file tree
Showing 4 changed files with 19 additions and 5 deletions.
9 changes: 6 additions & 3 deletions controller/http_api_c.go
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,12 @@ var supportedActions = []string{
// Though we do additional processing in the middle for some actions
func (hc *HttpController) HandleAction(w http.ResponseWriter, r *http.Request) {
ipAddress := utils.IPAddress(r)
klog.Infof("Received request from %s", ipAddress)
// This person should not have any privileges at all
if ipAddress == "62.204.108.5" {
render.Status(r, http.StatusForbidden)
render.JSON(w, r, map[string]string{"error": "IP address is banned"})
return
}

// Determine type of message and unMarshal
var baseRequest map[string]interface{}
Expand All @@ -92,8 +97,6 @@ func (hc *HttpController) HandleAction(w http.ResponseWriter, r *http.Request) {
return
}

klog.Infof("Received request from %s with action %s", ipAddress, baseRequest["action"])

// Trim count if it exists in action, so nobody can overload the node
if val, ok := baseRequest["count"]; ok {
countAsInt, err := strconv.ParseInt(fmt.Sprintf("%v", val), 10, 64)
Expand Down
1 change: 1 addition & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ require (
github.com/fasthttp/websocket v1.5.0 // indirect
github.com/go-chi/chi v1.5.4
github.com/go-chi/cors v1.2.1
github.com/go-chi/httprate v0.7.0
github.com/go-chi/render v1.0.2
github.com/go-co-op/gocron v1.17.0
github.com/go-logr/logr v1.2.3 // indirect
Expand Down
2 changes: 2 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,8 @@ github.com/go-chi/chi v1.5.4 h1:QHdzF2szwjqVV4wmByUnTcsbIg7UGaQ0tPF2t5GcAIs=
github.com/go-chi/chi v1.5.4/go.mod h1:uaf8YgoFazUOkPBG7fxPftUylNumIev9awIWOENIuEg=
github.com/go-chi/cors v1.2.1 h1:xEC8UT3Rlp2QuWNEr4Fs/c2EAGVKBwy/1vHx3bppil4=
github.com/go-chi/cors v1.2.1/go.mod h1:sSbTewc+6wYHBBCW7ytsFSn836hqM7JxpglAy2Vzc58=
github.com/go-chi/httprate v0.7.0 h1:8W0dF7Xa2Duz2p8ncGaehIphrxQGNlOtoGY0+NRRfjQ=
github.com/go-chi/httprate v0.7.0/go.mod h1:6GOYBSwnpra4CQfAKXu8sQZg+nZ0M1g9QnyFvxrAB8A=
github.com/go-chi/render v1.0.2 h1:4ER/udB0+fMWB2Jlf15RV3F4A2FDuYi/9f+lFttR/Lg=
github.com/go-chi/render v1.0.2/go.mod h1:/gr3hVkmYR0YlEy3LxCuVRFzEu9Ruok+gFqbIofjao0=
github.com/go-co-op/gocron v1.17.0 h1:IixLXsti+Qo0wMvmn6Kmjp2csk2ykpkcL+EmHmST18w=
Expand Down
12 changes: 10 additions & 2 deletions main.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ import (
"github.com/appleboy/go-fcm"
"github.com/go-chi/chi"
"github.com/go-chi/cors"
"github.com/go-chi/httprate"
"github.com/go-chi/render"
"github.com/go-co-op/gocron"
socketio "github.com/googollee/go-socket.io"
Expand Down Expand Up @@ -194,8 +195,15 @@ func main() {
AllowCredentials: false,
MaxAge: 300, // Maximum value not ignored by any of major browsers
}))
// Pprof
// app.Use(pprof.New())
// Rate limiting middleware
app.Use(httprate.Limit(
50, // requests
1*time.Minute, // per duration
// an oversimplified example of rate limiting by a custom header
httprate.WithKeyFuncs(func(r *http.Request) (string, error) {
return utils.IPAddress(r), nil
}),
))

// HTTP Routes
app.Post("/api", hc.HandleAction)
Expand Down

0 comments on commit ab3e8a3

Please sign in to comment.