github.com/haalcala/mattermost-server-change-repo@v0.0.0-20210713015153-16753fbeee5f/app/websocket_router.go (about)

     1  // Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
     2  // See LICENSE.txt for license information.
     3  
     4  package app
     5  
     6  import (
     7  	"net/http"
     8  
     9  	"github.com/mattermost/mattermost-server/v5/mlog"
    10  	"github.com/mattermost/mattermost-server/v5/model"
    11  	"github.com/mattermost/mattermost-server/v5/utils"
    12  )
    13  
    14  type webSocketHandler interface {
    15  	ServeWebSocket(*WebConn, *model.WebSocketRequest)
    16  }
    17  
    18  type WebSocketRouter struct {
    19  	server   *Server
    20  	app      *App
    21  	handlers map[string]webSocketHandler
    22  }
    23  
    24  func (wr *WebSocketRouter) Handle(action string, handler webSocketHandler) {
    25  	wr.handlers[action] = handler
    26  }
    27  
    28  func (wr *WebSocketRouter) ServeWebSocket(conn *WebConn, r *model.WebSocketRequest) {
    29  	wr.app.InitServer()
    30  
    31  	if r.Action == "" {
    32  		err := model.NewAppError("ServeWebSocket", "api.web_socket_router.no_action.app_error", nil, "", http.StatusBadRequest)
    33  		returnWebSocketError(wr.app, conn, r, err)
    34  		return
    35  	}
    36  
    37  	if r.Seq <= 0 {
    38  		err := model.NewAppError("ServeWebSocket", "api.web_socket_router.bad_seq.app_error", nil, "", http.StatusBadRequest)
    39  		returnWebSocketError(wr.app, conn, r, err)
    40  		return
    41  	}
    42  
    43  	if r.Action == model.WEBSOCKET_AUTHENTICATION_CHALLENGE {
    44  		if conn.GetSessionToken() != "" {
    45  			return
    46  		}
    47  
    48  		token, ok := r.Data["token"].(string)
    49  		if !ok {
    50  			conn.WebSocket.Close()
    51  			return
    52  		}
    53  
    54  		session, err := wr.app.GetSession(token)
    55  		if err != nil {
    56  			conn.WebSocket.Close()
    57  			return
    58  		}
    59  
    60  		conn.SetSession(session)
    61  		conn.SetSessionToken(session.Token)
    62  		conn.UserId = session.UserId
    63  
    64  		wr.app.HubRegister(conn)
    65  
    66  		wr.app.Srv().Go(func() {
    67  			wr.app.SetStatusOnline(session.UserId, false)
    68  			wr.app.UpdateLastActivityAtIfNeeded(*session)
    69  		})
    70  
    71  		resp := model.NewWebSocketResponse(model.STATUS_OK, r.Seq, nil)
    72  		hub := wr.app.GetHubForUserId(conn.UserId)
    73  		if hub == nil {
    74  			return
    75  		}
    76  		hub.SendMessage(conn, resp)
    77  
    78  		return
    79  	}
    80  
    81  	if !conn.IsAuthenticated() {
    82  		err := model.NewAppError("ServeWebSocket", "api.web_socket_router.not_authenticated.app_error", nil, "", http.StatusUnauthorized)
    83  		returnWebSocketError(wr.app, conn, r, err)
    84  		return
    85  	}
    86  
    87  	handler, ok := wr.handlers[r.Action]
    88  	if !ok {
    89  		err := model.NewAppError("ServeWebSocket", "api.web_socket_router.bad_action.app_error", nil, "", http.StatusInternalServerError)
    90  		returnWebSocketError(wr.app, conn, r, err)
    91  		return
    92  	}
    93  
    94  	handler.ServeWebSocket(conn, r)
    95  }
    96  
    97  func returnWebSocketError(app *App, conn *WebConn, r *model.WebSocketRequest, err *model.AppError) {
    98  	logF := mlog.Error
    99  	if err.StatusCode >= http.StatusBadRequest && err.StatusCode < http.StatusInternalServerError {
   100  		logF = mlog.Debug
   101  	}
   102  	logF(
   103  		"websocket routing error.",
   104  		mlog.Int64("seq", r.Seq),
   105  		mlog.String("user_id", conn.UserId),
   106  		mlog.String("system_message", err.SystemMessage(utils.T)),
   107  		mlog.Err(err),
   108  	)
   109  
   110  	hub := app.GetHubForUserId(conn.UserId)
   111  	if hub == nil {
   112  		return
   113  	}
   114  
   115  	err.DetailedError = ""
   116  	errorResp := model.NewWebSocketError(r.Seq, err)
   117  	hub.SendMessage(conn, errorResp)
   118  }