diff options
| -rw-r--r-- | floodstop.go | 47 | ||||
| -rw-r--r-- | websock.go | 14 |
2 files changed, 61 insertions, 0 deletions
diff --git a/floodstop.go b/floodstop.go new file mode 100644 index 0000000..ad4af0f --- /dev/null +++ b/floodstop.go @@ -0,0 +1,47 @@ +package main + +import ( + "time" +) + +type Floodstop struct { + ask chan chan bool + stop chan struct{} +} + +func NewFloodstop(reset time.Duration, countMax int) (fs *Floodstop) { + fs = &Floodstop{ + ask: make(chan chan bool), + stop: make(chan struct{}), + } + + ticker := time.NewTicker(reset) + counter := 0 + + go func() { + defer ticker.Stop() + for { + select { + case <-fs.stop: + return + case retCh := <-fs.ask: + counter++ + retCh <- (counter < countMax) + case <-ticker.C: + counter = 0 + } + } + }() + + return +} + +func (fs *Floodstop) Stop() { + fs.stop <- struct{}{} +} + +func (fs *Floodstop) Ask() bool { + ch := make(chan bool) + fs.ask <- ch + return <-ch +} @@ -5,6 +5,7 @@ import ( "github.com/gorilla/mux" "github.com/kch42/simplechat/chat" "net/http" + "time" ) type JoinResponse struct { @@ -13,6 +14,11 @@ type JoinResponse struct { Buddies []string `json:"buddies,omitempty"` } +const ( + floodstopReset = 1 * time.Second + floodstopMax = 3 +) + func AcceptWebSock(rw http.ResponseWriter, req *http.Request) { vars := mux.Vars(req) roomname := vars["chatroom"] @@ -45,6 +51,9 @@ func AcceptWebSock(rw http.ResponseWriter, req *http.Request) { exit := make(chan struct{}) + floodstop := NewFloodstop(floodstopReset, floodstopMax) + defer floodstop.Stop() + go func() { var s string for { @@ -52,6 +61,11 @@ func AcceptWebSock(rw http.ResponseWriter, req *http.Request) { break } + if !floodstop.Ask() { + // User sent too many requests in a short amount of time. Kick them! + break + } + if s == "" { continue } |