github.com/masterhung0112/hk_server/v5@v5.0.0-20220302090640-ec71aef15e1c/api4/websocket.go (about)

     1  // Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
     2  // See LICENSE.txt for license information.
     3  
     4  package api4
     5  
     6  import (
     7  	"net/http"
     8  
     9  	"github.com/gorilla/websocket"
    10  
    11  	"github.com/masterhung0112/hk_server/v5/app"
    12  	"github.com/masterhung0112/hk_server/v5/model"
    13  	"github.com/masterhung0112/hk_server/v5/shared/mlog"
    14  )
    15  
    16  const (
    17  	connectionIDParam   = "connection_id"
    18  	sequenceNumberParam = "sequence_number"
    19  )
    20  
    21  func (api *API) InitWebSocket() {
    22  	// Optionally supports a trailing slash
    23  	api.BaseRoutes.ApiRoot.Handle("/{websocket:websocket(?:\\/)?}", api.ApiHandlerTrustRequester(connectWebSocket)).Methods("GET")
    24  }
    25  
    26  func connectWebSocket(c *Context, w http.ResponseWriter, r *http.Request) {
    27  	upgrader := websocket.Upgrader{
    28  		ReadBufferSize:  model.SOCKET_MAX_MESSAGE_SIZE_KB,
    29  		WriteBufferSize: model.SOCKET_MAX_MESSAGE_SIZE_KB,
    30  		CheckOrigin:     c.App.OriginChecker(),
    31  	}
    32  
    33  	ws, err := upgrader.Upgrade(w, r, nil)
    34  	if err != nil {
    35  		c.Err = model.NewAppError("connect", "api.web_socket.connect.upgrade.app_error", nil, err.Error(), http.StatusInternalServerError)
    36  		return
    37  	}
    38  
    39  	// We initialize webconn with all the necessary data.
    40  	// If the queues are empty, they are initialized in the constructor.
    41  	cfg := &app.WebConnConfig{
    42  		WebSocket: ws,
    43  		Session:   *c.AppContext.Session(),
    44  		TFunc:     c.AppContext.T,
    45  		Locale:    "",
    46  		Active:    true,
    47  	}
    48  
    49  	if *c.App.Config().ServiceSettings.EnableReliableWebSockets {
    50  		cfg.ConnectionID = r.URL.Query().Get(connectionIDParam)
    51  		if cfg.ConnectionID == "" || c.AppContext.Session().UserId == "" {
    52  			// If not present, we assume client is not capable yet, or it's a fresh connection.
    53  			// We just create a new ID.
    54  			cfg.ConnectionID = model.NewId()
    55  			// In case of fresh connection id, sequence number is already zero.
    56  		} else {
    57  			cfg, err = c.App.PopulateWebConnConfig(c.AppContext.Session(), cfg, r.URL.Query().Get(sequenceNumberParam))
    58  			if err != nil {
    59  				mlog.Warn("Error while populating webconn config", mlog.String("id", r.URL.Query().Get(connectionIDParam)), mlog.Err(err))
    60  				ws.Close()
    61  				return
    62  			}
    63  		}
    64  	}
    65  
    66  	wc := c.App.NewWebConn(cfg)
    67  	if c.AppContext.Session().UserId != "" {
    68  		c.App.HubRegister(wc)
    69  	}
    70  
    71  	wc.Pump()
    72  }