Pass remote address in WebSocketMessageHasBeenPosted plugin hook (#27332)

server/channels/api4/plugin_test.go

@@ -2033,3 +2033,49 @@ func TestPluginWebSocketSession(t *testing.T) {
 	_, sessionID := hooks.MessageWillBePosted(nil, nil)
 	require.Equal(t, sessions[0].Id, sessionID)
 }
+
+func TestPluginWebSocketRemoteAddress(t *testing.T) {
+	th := Setup(t).InitBasic()
+	defer th.TearDown()
+
+	pluginID := "com.mattermost.websocket_remote_address_test"
+
+	// Compile plugin
+	fullPath := filepath.Join(server.GetPackagePath(), "channels", "app", "plugin_api_tests", "manual.test_websocket_remote_address", "main.go")
+	pluginCode, err := os.ReadFile(fullPath)
+	require.NoError(t, err)
+	require.NotEmpty(t, pluginCode)
+	pluginDir, err := filepath.Abs(*th.App.Config().PluginSettings.Directory)
+	require.NoError(t, err)
+	backend := filepath.Join(pluginDir, pluginID, "backend.exe")
+	utils.CompileGo(t, string(pluginCode), backend)
+	os.WriteFile(filepath.Join(pluginDir, pluginID, "plugin.json"), []byte(`{"id": "`+pluginID+`", "server": {"executable": "backend.exe"}}`), 0600)
+
+	// Activate the plugin
+	manifest, activated, reterr := th.App.GetPluginsEnvironment().Activate(pluginID)
+	require.NoError(t, reterr)
+	require.NotNil(t, manifest)
+	require.True(t, activated)
+
+	// Connect through WebSocket and send a message
+	reqURL := fmt.Sprintf("ws://localhost:%d", th.Server.ListenAddr.Port)
+	wsc, err := model.NewWebSocketClient4(reqURL, th.Client.AuthToken)
+	require.NoError(t, err)
+	require.NotNil(t, wsc)
+	wsc.Listen()
+	defer wsc.Close()
+	resp := <-wsc.ResponseChannel
+	require.Equal(t, resp.Status, model.StatusOk)
+	wsc.SendMessage("custom_action", map[string]any{"value": "test"})
+
+	// Verify the remote address has been set correctly. Check plugin code in
+	// channels/app/plugin_api_tests/manual.test_websocket_remote_address
+	//
+	// Here the MessageWillBePosted hook is used purely as a way to
+	// communicate with the plugin side.
+	hooks, err := th.App.GetPluginsEnvironment().HooksForPlugin(pluginID)
+	require.NoError(t, err)
+	require.NotNil(t, hooks)
+	_, remoteAddr := hooks.MessageWillBePosted(nil, nil)
+	require.NotEmpty(t, remoteAddr)
+}

server/channels/api4/websocket.go

@@ -50,6 +50,8 @@ func connectWebSocket(c *Context, w http.ResponseWriter, r *http.Request) {
 		Locale:        "",
 		Active:        true,
 		PostedAck:     r.URL.Query().Get(postedAckParam) == "true",
+		RemoteAddress: c.AppContext.IPAddress(),
+		XForwardedFor: c.AppContext.XForwardedFor(),
 	}
 	// The WebSocket upgrade request coming from mobile is missing the
 	// user agent so we need to fallback on the session's metadata.

server/channels/app/platform/web_conn.go

@@ -70,6 +70,8 @@ type WebConnConfig struct {
 	ReuseCount    int
 	OriginClient  string
 	PostedAck     bool
+	RemoteAddress string
+	XForwardedFor string
 
 	// These aren't necessary to be exported to api layer.
 	sequence         int
@@ -121,6 +123,10 @@ type WebConn struct {
 
 	// The client type behind the connection (i.e. web, desktop or mobile)
 	originClient string
+	// The remote address from the original HTTP Upgrade request
+	remoteAddress string
+	// The X-Forwarded-For HTTP header value from the origina HTTP Upgrade request
+	xForwardedFor string
 
 	activeChannelID                 atomic.Value
 	activeTeamID                    atomic.Value
@@ -245,6 +251,8 @@ func (ps *PlatformService) NewWebConn(cfg *WebConnConfig, suite SuiteIFace, runn
 		lastLogTimeSlow:    time.Now(),
 		lastLogTimeFull:    time.Now(),
 		originClient:       cfg.OriginClient,
+		remoteAddress:      cfg.RemoteAddress,
+		xForwardedFor:      cfg.XForwardedFor,
 	}
 	wc.Active.Store(cfg.Active)
 
@@ -479,6 +487,12 @@ func (wc *WebConn) readPump() {
 			clonedReq.Session.Id = session.Id
 		}
 
+		if clonedReq.Data == nil {
+			clonedReq.Data = map[string]any{}
+		}
+		clonedReq.Data[model.WebSocketRemoteAddr] = wc.remoteAddress
+		clonedReq.Data[model.WebSocketXForwardedFor] = wc.xForwardedFor
+
 		wc.pluginPosted <- pluginWSPostedHook{wc.GetConnectionID(), wc.UserId, clonedReq}
 	}
 }

server/channels/app/plugin_api_tests/manual.test_websocket_remote_address/main.go

@@ -0,0 +1,28 @@
+// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
+// See LICENSE.txt for license information.
+
+package main
+
+import (
+	"github.com/mattermost/mattermost/server/public/model"
+	"github.com/mattermost/mattermost/server/public/plugin"
+)
+
+type Plugin struct {
+	plugin.MattermostPlugin
+	addrCh chan string
+}
+
+func (p *Plugin) MessageWillBePosted(_ *plugin.Context, _ *model.Post) (*model.Post, string) {
+	return nil, <-p.addrCh
+}
+
+func (p *Plugin) WebSocketMessageHasBeenPosted(connID, userID string, req *model.WebSocketRequest) {
+	p.addrCh <- req.Data[model.WebSocketRemoteAddr].(string)
+}
+
+func main() {
+	plugin.ClientMain(&Plugin{
+		addrCh: make(chan string, 1),
+	})
+}

server/public/model/websocket_request.go

@@ -9,6 +9,11 @@ import (
 	"github.com/vmihailenco/msgpack/v5"
 )
 
+const (
+	WebSocketRemoteAddr    = "remote_addr"
+	WebSocketXForwardedFor = "x_forwarded_for"
+)
+
 // WebSocketRequest represents a request made to the server through a websocket.
 type WebSocketRequest struct {
 	// Client-provided fields