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 }