Prevent users from flooding the server

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
+}
diff --git a/websock.go b/websock.go
index 9734e65..f42580f 100644
--- a/websock.go
+++ b/websock.go
@@ -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
 				}