Lack of MFA enforcement in WebSocket connections

@@ -12,6 +12,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/dgryski/dgoogauth"
 	"github.com/gorilla/websocket"
 	"github.com/stretchr/testify/require"
 
@@ -536,3 +537,147 @@ func TestValidateDisconnectErrCode(t *testing.T) {
 		})
 	}
 }
+
+// Helper function to enable MFA enforcement in config
+func enableMFAEnforcement(th *TestHelper) {
+	th.App.UpdateConfig(func(cfg *model.Config) {
+		*cfg.ServiceSettings.EnableMultifactorAuthentication = true
+		*cfg.ServiceSettings.EnforceMultifactorAuthentication = true
+	})
+}
+
+// Helper function to set up MFA for a user
+func setupUserWithMFA(t *testing.T, th *TestHelper, user *model.User) string {
+	// Setup MFA properly - following authentication_test.go pattern
+	secret, appErr := th.App.GenerateMfaSecret(user.Id)
+	require.Nil(t, appErr)
+	err := th.Server.Store().User().UpdateMfaActive(user.Id, true)
+	require.NoError(t, err)
+	err = th.Server.Store().User().UpdateMfaSecret(user.Id, secret.Secret)
+	require.NoError(t, err)
+	return secret.Secret
+}
+
+func TestWebSocketMFAEnforcement(t *testing.T) {
+	mainHelper.Parallel(t)
+
+	t.Run("WebSocket works when MFA enforcement is disabled", func(t *testing.T) {
+		th := Setup(t).InitBasic()
+		defer th.TearDown()
+
+		// MFA enforcement disabled - should work normally
+		webSocketClient := th.CreateConnectedWebSocketClient(t)
+		defer webSocketClient.Close()
+
+		webSocketClient.GetStatuses()
+
+		select {
+		case resp := <-webSocketClient.ResponseChannel:
+			require.Nil(t, resp.Error, "WebSocket should work when MFA enforcement is disabled")
+			require.Equal(t, resp.Status, model.StatusOk)
+		case <-time.After(3 * time.Second):
+			require.Fail(t, "Expected WebSocket response but got timeout")
+		}
+	})
+
+	t.Run("WebSocket blocked when MFA required but user has no MFA", func(t *testing.T) {
+		th := SetupEnterprise(t).InitBasic()
+		defer th.TearDown()
+
+		// Enable MFA enforcement in config
+		enableMFAEnforcement(th)
+		// Defer the teardown to reset the config after the test
+		defer func() {
+			th.App.UpdateConfig(func(cfg *model.Config) {
+				*cfg.ServiceSettings.EnforceMultifactorAuthentication = false
+			})
+		}()
+
+		// Create user without MFA using existing basic user to avoid license timing issues
+		user := th.BasicUser
+
+		// Login user (this should work for initial authentication)
+		client := th.CreateClient()
+		_, _, err := client.Login(context.Background(), user.Email, "Pa$$word11")
+		require.NoError(t, err)
+
+		// Create WebSocket client - initial connection succeeds, but subsequent API requests require completed MFA
+		webSocketClient, err := th.CreateWebSocketClientWithClient(client)
+		require.NoError(t, err)
+		require.NotNil(t, webSocketClient, "webSocketClient should not be nil")
+		webSocketClient.Listen()
+		defer webSocketClient.Close()
+
+		// First, consume the successful authentication challenge response
+		authResp := <-webSocketClient.ResponseChannel
+		require.Nil(t, authResp.Error, "Authentication challenge should succeed")
+		require.Equal(t, authResp.Status, model.StatusOk)
+
+		// Individual WebSocket requests should be blocked due to MFA requirement
+		webSocketClient.GetStatuses()
+
+		// Should get authentication error due to MFA requirement on the second request
+		select {
+		case resp := <-webSocketClient.ResponseChannel:
+			t.Logf("Received response: Error=%v, Status=%s, SeqReply=%d", resp.Error, resp.Status, resp.SeqReply)
+			require.NotNil(t, resp.Error, "Should get authentication error due to MFA requirement")
+			require.Equal(t, "api.web_socket_router.not_authenticated.app_error", resp.Error.Id,
+				"Should get specific 'not authenticated' error ID due to MFA requirement")
+		case <-time.After(3 * time.Second):
+			require.Fail(t, "Expected WebSocket error response but got timeout")
+		}
+	})
+
+	t.Run("WebSocket connection allowed when user has MFA active", func(t *testing.T) {
+		th := SetupEnterprise(t).InitBasic()
+		defer th.TearDown()
+
+		// Enable MFA enforcement in config
+		enableMFAEnforcement(th)
+		// Defer the teardown to reset the config after the test
+		defer func() {
+			th.App.UpdateConfig(func(cfg *model.Config) {
+				*cfg.ServiceSettings.EnforceMultifactorAuthentication = false
+			})
+		}()
+
+		// Create user and set up MFA
+		user := &model.User{
+			Email:    th.GenerateTestEmail(),
+			Username: model.NewUsername(),
+			Password: "password123",
+		}
+		ruser, _, err := th.Client.CreateUser(context.Background(), user)
+		require.NoError(t, err)
+
+		th.LinkUserToTeam(ruser, th.BasicTeam)
+		_, err = th.App.Srv().Store().User().VerifyEmail(ruser.Id, ruser.Email)
+		require.NoError(t, err)
+
+		// Setup MFA for the user and get the secret
+		secretString := setupUserWithMFA(t, th, ruser)
+
+		// Generate TOTP token from the user's MFA secret
+		code := dgoogauth.ComputeCode(secretString, time.Now().UTC().Unix()/30)
+		token := fmt.Sprintf("%06d", code)
+
+		client := th.CreateClient()
+		_, _, err = client.LoginWithMFA(context.Background(), user.Email, user.Password, token)
+		require.NoError(t, err)
+
+		// WebSocket connection should work
+		webSocketClient := th.CreateConnectedWebSocketClientWithClient(t, client)
+		defer webSocketClient.Close()
+
+		// Should be able to get statuses
+		webSocketClient.GetStatuses()
+
+		select {
+		case resp := <-webSocketClient.ResponseChannel:
+			require.Nil(t, resp.Error, "WebSocket should work when MFA is properly set up")
+			require.Equal(t, resp.Status, model.StatusOk)
+		case <-time.After(5 * time.Second):
+			require.Fail(t, "Expected WebSocket response but got timeout")
+		}
+	})
+}

@@ -65,6 +65,10 @@ func (ms *mockSuite) HasPermissionToReadChannel(rctx request.CTX, userID string,
 	return true
 }
 
+func (ms *mockSuite) MFARequired(rctx request.CTX) *model.AppError {
+	return nil
+}
+
 func setupDBStore(tb testing.TB) (store.Store, *model.SqlSettings) {
 	var dbStore store.Store
 	var dbSettings *model.SqlSettings

@@ -66,6 +66,26 @@ func (_m *SuiteIFace) HasPermissionToReadChannel(rctx request.CTX, userID string
 	return r0
 }
 
+// MFARequired provides a mock function with given fields: c
+func (_m *SuiteIFace) MFARequired(rctx request.CTX) *model.AppError {
+	ret := _m.Called(rctx)
+
+	if len(ret) == 0 {
+		panic("no return value specified for MFARequired")
+	}
+
+	var r0 *model.AppError
+	if rf, ok := ret.Get(0).(func(request.CTX) *model.AppError); ok {
+		r0 = rf(rctx)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*model.AppError)
+		}
+	}
+
+	return r0
+}
+
 // RolesGrantPermission provides a mock function with given fields: roleNames, permissionId
 func (_m *SuiteIFace) RolesGrantPermission(roleNames []string, permissionId string) bool {
 	ret := _m.Called(roleNames, permissionId)

@@ -453,7 +453,7 @@ func (wc *WebConn) readPump() {
 		if err := wc.WebSocket.SetReadDeadline(time.Now().Add(pongWaitTime)); err != nil {
 			return err
 		}
-		if wc.IsAuthenticated() {
+		if wc.IsBasicAuthenticated() {
 			userID := wc.UserId
 			wc.Platform.Go(func() {
 				wc.Platform.SetStatusAwayIfNeeded(userID, false)
@@ -770,8 +770,8 @@ func (wc *WebConn) InvalidateCache() {
 	wc.SetSessionExpiresAt(0)
 }
 
-// IsAuthenticated returns whether the given WebConn is authenticated or not.
-func (wc *WebConn) IsAuthenticated() bool {
+// IsBasicAuthenticated returns whether the given WebConn has a valid session.
+func (wc *WebConn) IsBasicAuthenticated() bool {
 	// Check the expiry to see if we need to check for a new session
 	if wc.GetSessionExpiresAt() < model.GetMillis() {
 		if wc.GetSessionToken() == "" {
@@ -799,6 +799,24 @@ func (wc *WebConn) IsAuthenticated() bool {
 	return true
 }
 
+// IsMFAAuthenticated returns whether the user has completed MFA when required.
+func (wc *WebConn) IsMFAAuthenticated() bool {
+	session := wc.GetSession()
+	c := request.EmptyContext(wc.Platform.logger).WithSession(session)
+
+	// Check if MFA is required and user has NOT completed MFA
+	if appErr := wc.Suite.MFARequired(c); appErr != nil {
+		return false
+	}
+
+	return true
+}
+
+// IsAuthenticated returns whether the given WebConn is fully authenticated (session + MFA).
+func (wc *WebConn) IsAuthenticated() bool {
+	return wc.IsBasicAuthenticated() && wc.IsMFAAuthenticated()
+}
+
 func (wc *WebConn) createHelloMessage() *model.WebSocketEvent {
 	ee := wc.Platform.LicenseManager() != nil
 
@@ -856,7 +874,7 @@ func (wc *WebConn) ShouldSendEventToGuest(msg *model.WebSocketEvent) bool {
 
 // ShouldSendEvent returns whether the message should be sent or not.
 func (wc *WebConn) ShouldSendEvent(msg *model.WebSocketEvent) bool {
-	// IMPORTANT: Do not send event if WebConn does not have a session
+	// IMPORTANT: Do not send event if WebConn does not have a session and completed MFA
 	if !wc.IsAuthenticated() {
 		return false
 	}

@@ -30,6 +30,7 @@ type SuiteIFace interface {
 	RolesGrantPermission(roleNames []string, permissionId string) bool
 	HasPermissionToReadChannel(rctx request.CTX, userID string, channel *model.Channel) bool
 	UserCanSeeOtherUser(rctx request.CTX, userID string, otherUserId string) (bool, *model.AppError)
+	MFARequired(rctx request.CTX) *model.AppError
 }
 
 type webConnActivityMessage struct {
@@ -572,7 +573,7 @@ func (h *Hub) Start() {
 				}
 				atomic.StoreInt64(&h.connectionCount, int64(connIndex.AllActive()))
 
-				if webConnReg.conn.IsAuthenticated() && webConnReg.conn.reuseCount == 0 {
+				if webConnReg.conn.IsBasicAuthenticated() && webConnReg.conn.reuseCount == 0 {
 					// The hello message should only be sent when the reuseCount is 0.
 					// i.e in server restart, or long timeout, or fresh connection case.
 					// In case of seq number not found in dead queue, it is handled by

@@ -18,6 +18,7 @@ import (
 
 	"github.com/gorilla/websocket"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 
 	"github.com/mattermost/mattermost/server/public/model"
@@ -598,6 +599,7 @@ func TestHubIsRegistered(t *testing.T) {
 
 	mockSuite := &platform_mocks.SuiteIFace{}
 	mockSuite.On("GetSession", session.Token).Return(session, nil)
+	mockSuite.On("MFARequired", mock.Anything).Return(nil)
 	th.Suite = mockSuite
 
 	s := httptest.NewServer(dummyWebsocketHandler(t))
@@ -633,6 +635,7 @@ func TestHubWebConnCount(t *testing.T) {
 
 	mockSuite := &platform_mocks.SuiteIFace{}
 	mockSuite.On("GetSession", session.Token).Return(session, nil)
+	mockSuite.On("MFARequired", mock.Anything).Return(nil)
 	th.Suite = mockSuite
 
 	s := httptest.NewServer(dummyWebsocketHandler(t))