Add plugin websocket hooks (#18151)

* Add plugin websocket hooks

* Improve sending message mechanism and filter out of the router plugin specific messages

* Return and manage error if the request clone fails

* Wording change to avoid repetition

Co-authored-by: Mattermod <mattermod@users.noreply.github.com>

app/web_conn.go

@@ -11,6 +11,7 @@ import (
 	"net"
 	"net/http"
 	"strconv"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -18,6 +19,7 @@ import (
 	"github.com/gorilla/websocket"
 
 	"github.com/mattermost/mattermost-server/v6/model"
+	"github.com/mattermost/mattermost-server/v6/plugin"
 	"github.com/mattermost/mattermost-server/v6/shared/i18n"
 	"github.com/mattermost/mattermost-server/v6/shared/mlog"
 )
@@ -40,6 +42,14 @@ const (
 	reconnectLossless = "lossless"
 )
 
+const websocketMessagePluginPrefix = "custom_"
+
+type pluginWSPostedHook struct {
+	connectionID string
+	userID       string
+	req          *model.WebSocketRequest
+}
+
 type WebConnConfig struct {
 	WebSocket    *websocket.Conn
 	Session      model.Session
@@ -90,6 +100,7 @@ type WebConn struct {
 	connectionID atomic.Value
 	endWritePump chan struct{}
 	pumpFinished chan struct{}
+	pluginPosted chan pluginWSPostedHook
 }
 
 // CheckConnResult indicates whether a connectionID was present in the hub or not.
@@ -179,6 +190,7 @@ func (a *App) NewWebConn(cfg *WebConnConfig) *WebConn {
 		active:             cfg.Active,
 		endWritePump:       make(chan struct{}),
 		pumpFinished:       make(chan struct{}),
+		pluginPosted:       make(chan pluginWSPostedHook, 10),
 	}
 
 	wc.SetSession(&cfg.Session)
@@ -186,9 +198,31 @@ func (a *App) NewWebConn(cfg *WebConnConfig) *WebConn {
 	wc.SetSessionExpiresAt(cfg.Session.ExpiresAt)
 	wc.SetConnectionID(cfg.ConnectionID)
 
+	if pluginsEnvironment := wc.App.GetPluginsEnvironment(); pluginsEnvironment != nil {
+		wc.App.Srv().Go(func() {
+			pluginsEnvironment.RunMultiPluginHook(func(hooks plugin.Hooks) bool {
+				hooks.OnWebSocketConnect(wc.GetConnectionID(), wc.UserId)
+				return true
+			}, plugin.OnWebSocketConnectID)
+		})
+	}
+
 	return wc
 }
 
+func (wc *WebConn) pluginPostedConsumer(wg *sync.WaitGroup) {
+	defer wg.Done()
+
+	for msg := range wc.pluginPosted {
+		if pluginsEnvironment := wc.App.GetPluginsEnvironment(); pluginsEnvironment != nil {
+			pluginsEnvironment.RunMultiPluginHook(func(hooks plugin.Hooks) bool {
+				hooks.WebSocketMessageHasBeenPosted(msg.connectionID, msg.userID, msg.req)
+				return true
+			}, plugin.WebSocketMessageHasBeenPostedID)
+		}
+	}
+}
+
 // Close closes the WebConn.
 func (wc *WebConn) Close() {
 	wc.WebSocket.Close()
@@ -261,11 +295,25 @@ func (wc *WebConn) Pump() {
 		defer wg.Done()
 		wc.writePump()
 	}()
+
+	wg.Add(1)
+	go wc.pluginPostedConsumer(&wg)
+
 	wc.readPump()
 	close(wc.endWritePump)
+	close(wc.pluginPosted)
 	wg.Wait()
 	wc.App.HubUnregister(wc)
 	close(wc.pumpFinished)
+
+	if pluginsEnvironment := wc.App.GetPluginsEnvironment(); pluginsEnvironment != nil {
+		wc.App.Srv().Go(func() {
+			pluginsEnvironment.RunMultiPluginHook(func(hooks plugin.Hooks) bool {
+				hooks.OnWebSocketDisconnect(wc.GetConnectionID(), wc.UserId)
+				return true
+			}, plugin.OnWebSocketDisconnectID)
+		})
+	}
 }
 
 func (wc *WebConn) readPump() {
@@ -290,7 +338,20 @@ func (wc *WebConn) readPump() {
 			wc.logSocketErr("websocket.read", err)
 			return
 		}
-		wc.App.Srv().WebSocketRouter.ServeWebSocket(wc, &req)
+
+		// Messages which actions are prefixed with the plugin prefix
+		// should only be dispatched to the plugins
+		if !strings.HasPrefix(req.Action, websocketMessagePluginPrefix) {
+			wc.App.Srv().WebSocketRouter.ServeWebSocket(wc, &req)
+		}
+
+		clonedReq, err := req.Clone()
+		if err != nil {
+			wc.logSocketErr("websocket.cloneRequest", err)
+			continue
+		}
+
+		wc.pluginPosted <- pluginWSPostedHook{wc.GetConnectionID(), wc.UserId, clonedReq}
 	}
 }