Implement basic rate limiting for incoming WebSocket messages

mattermost · Apr 28, 2022 · a5df619 · a5df619
1 parent 7fd377e
commit a5df619
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 6 deletions.
diff --git a/go.mod b/go.mod
@@ -17,6 +17,7 @@ require (
 	github.com/mattermost/rtcd v0.3.1-0.20220427165117-035048fda047
 	github.com/pion/stun v0.3.5
 	github.com/rudderlabs/analytics-go v3.3.2+incompatible
+	golang.org/x/time v0.0.0-20201208040808-7e3f01d25324
 )
 
 require (

diff --git a/go.sum b/go.sum
@@ -1918,6 +1918,7 @@ golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxb
 golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
+golang.org/x/time v0.0.0-20201208040808-7e3f01d25324 h1:Hir2P/De0WpUhtrKGGjvSb2YxUgyZ7EFOSLIcSSpiwE=
 golang.org/x/time v0.0.0-20201208040808-7e3f01d25324/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/tools v0.0.0-20180221164845-07fd8470d635/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20180525024113-a5b4c53f6e8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=

diff --git a/server/session.go b/server/session.go
@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"time"
 
+	"golang.org/x/time/rate"
+
 	"github.com/mattermost/mattermost-server/v6/model"
 )
 
@@ -25,6 +27,8 @@ type session struct {
 
 	doneCh  chan struct{}
 	closeCh chan struct{}
+
+	limiter *rate.Limiter
 }
 
 func newUserSession(userID, channelID, connID string) *session {
@@ -38,6 +42,7 @@ func newUserSession(userID, channelID, connID string) *session {
 		wsCloseCh:   make(chan struct{}),
 		closeCh:     make(chan struct{}),
 		doneCh:      make(chan struct{}),
+		limiter:     rate.NewLimiter(2, 20),
 	}
 }
 

diff --git a/server/websocket.go b/server/websocket.go
@@ -560,6 +560,18 @@ func (p *Plugin) WebSocketMessageHasBeenPosted(connID, userID string, req *model
 	var msg clientMessage
 	msg.Type = strings.TrimPrefix(req.Action, wsActionPrefix)
 
+	p.mut.RLock()
+	us := p.sessions[connID]
+	p.mut.RUnlock()
+
+	if msg.Type != clientMessageTypeJoin && us == nil {
+		return
+	}
+
+	if us != nil && !us.limiter.Allow() {
+		p.LogError("message was dropped by rate limiter", "msgType", msg.Type, "userID", us.userID, "connID", us.connID)
+	}
+
 	switch msg.Type {
 	case clientMessageTypeJoin:
 		channelID, ok := req.Data["channelID"].(string)
@@ -600,12 +612,6 @@ func (p *Plugin) WebSocketMessageHasBeenPosted(connID, userID string, req *model
 		msg.Data = []byte(msgData)
 	}
 
-	p.mut.RLock()
-	us := p.sessions[connID]
-	p.mut.RUnlock()
-	if us == nil || us.connID != connID {
-		return
-	}
 	select {
 	case us.wsMsgCh <- msg:
 	default: