package main import ( "errors" "log" "net/http" "time" "github.com/gorilla/websocket" "github.com/thrasher-/gocryptotrader/common" "github.com/thrasher-/gocryptotrader/config" "github.com/thrasher-/gocryptotrader/currency" ) // Const vars for websocket const ( WebsocketResponseSuccess = "OK" ) // WebsocketClient stores information related to the websocket client type WebsocketClient struct { ID int Conn *websocket.Conn LastRecv time.Time Authenticated bool } // WebsocketEvent is the struct used for websocket events type WebsocketEvent struct { Exchange string `json:"exchange,omitempty"` AssetType string `json:"assetType,omitempty"` Event string Data interface{} } // WebsocketEventResponse is the struct used for websocket event responses type WebsocketEventResponse struct { Event string `json:"event"` Data interface{} `json:"data"` Error string `json:"error"` } // WebsocketOrderbookTickerRequest is a struct used for ticker and orderbook // requests type WebsocketOrderbookTickerRequest struct { Exchange string `json:"exchangeName"` Currency string `json:"currency"` AssetType string `json:"assetType"` } // WebsocketClientHub stores an array of websocket clients var WebsocketClientHub []WebsocketClient // WebsocketClientHandler upgrades the HTTP connection to a websocket // compatible one func WebsocketClientHandler(w http.ResponseWriter, r *http.Request) { connectionLimit := bot.config.Webserver.WebsocketConnectionLimit numClients := len(WebsocketClientHub) if numClients >= connectionLimit { log.Printf("Websocket client rejected due to websocket client limit reached. Number of clients %d. Limit %d.", numClients, connectionLimit) w.WriteHeader(http.StatusForbidden) return } upgrader := websocket.Upgrader{ WriteBufferSize: 1024, ReadBufferSize: 1024, } // Allow insecure origin if the Origin request header is present and not // equal to the Host request header. Default to false if bot.config.Webserver.WebsocketAllowInsecureOrigin { upgrader.CheckOrigin = func(r *http.Request) bool { return true } } newClient := WebsocketClient{ ID: len(WebsocketClientHub), } conn, err := upgrader.Upgrade(w, r, nil) if err != nil { log.Println(err) return } newClient.Conn = conn WebsocketClientHub = append(WebsocketClientHub, newClient) numClients++ log.Printf("New websocket client connected. Connected clients: %d. Limit %d.", numClients, connectionLimit) } // DisconnectWebsocketClient disconnects a websocket client func DisconnectWebsocketClient(id int, err error) { for i := range WebsocketClientHub { if WebsocketClientHub[i].ID == id { WebsocketClientHub[i].Conn.Close() WebsocketClientHub = append(WebsocketClientHub[:i], WebsocketClientHub[i+1:]...) log.Printf("Disconnected Websocket client, error: %s", err) return } } } // SendWebsocketMessage sends a websocket message to a specific client func SendWebsocketMessage(id int, data interface{}) error { for _, x := range WebsocketClientHub { if x.ID == id { return x.Conn.WriteJSON(data) } } return nil } // BroadcastWebsocketMessage broadcasts a websocket event message to all // websocket clients func BroadcastWebsocketMessage(evt WebsocketEvent) error { for _, x := range WebsocketClientHub { x.Conn.WriteJSON(evt) } return nil } // WebsocketAuth is a struct used for type WebsocketAuth struct { Username string `json:"username"` Password string `json:"password"` } type wsCommandHandler func(wsClient *websocket.Conn, data interface{}) error var wsHandlers = map[string]wsCommandHandler{ "getconfig": wsGetConfig, "saveconfig": wsSaveConfig, "getaccountinfo": wsGetAccountInfo, "gettickers": wsGetTickers, "getticker": wsGetTicker, "getorderbooks": wsGetOrderbooks, "getorderbook": wsGetOrderbook, "getexchangerates": wsGetExchangeRates, "getportfolio": wsGetPortfolio, } func wsGetConfig(wsClient *websocket.Conn, data interface{}) error { wsResp := WebsocketEventResponse{ Event: "GetConfig", Data: bot.config, } return wsClient.WriteJSON(wsResp) } func wsSaveConfig(wsClient *websocket.Conn, data interface{}) error { wsResp := WebsocketEventResponse{ Event: "SaveConfig", } var cfg config.Config err := common.JSONDecode(data.([]byte), &cfg) if err != nil { wsResp.Error = err.Error() err = wsClient.WriteJSON(wsResp) if err != nil { return err } } err = bot.config.UpdateConfig(bot.configFile, cfg) if err != nil { wsResp.Error = err.Error() err = wsClient.WriteJSON(wsResp) if err != nil { return err } } SetupExchanges() wsResp.Data = WebsocketResponseSuccess return wsClient.WriteJSON(wsResp) } func wsGetAccountInfo(wsClient *websocket.Conn, data interface{}) error { accountInfo := GetAllEnabledExchangeAccountInfo() wsResp := WebsocketEventResponse{ Event: "GetAccountInfo", Data: accountInfo, } return wsClient.WriteJSON(wsResp) } func wsGetTickers(wsClient *websocket.Conn, data interface{}) error { wsResp := WebsocketEventResponse{ Event: "GetTickers", } wsResp.Data = GetAllActiveTickers() return wsClient.WriteJSON(wsResp) } func wsGetTicker(wsClient *websocket.Conn, data interface{}) error { wsResp := WebsocketEventResponse{ Event: "GetTicker", } var tickerReq WebsocketOrderbookTickerRequest err := common.JSONDecode(data.([]byte), &tickerReq) if err != nil { wsResp.Error = err.Error() wsClient.WriteJSON(wsResp) return err } result, err := GetSpecificTicker(tickerReq.Currency, tickerReq.Exchange, tickerReq.AssetType) if err != nil { wsResp.Error = err.Error() wsClient.WriteJSON(wsResp) return err } wsResp.Data = result return wsClient.WriteJSON(wsResp) } func wsGetOrderbooks(wsClient *websocket.Conn, data interface{}) error { wsResp := WebsocketEventResponse{ Event: "GetOrderbooks", } wsResp.Data = GetAllActiveOrderbooks() return wsClient.WriteJSON(wsResp) } func wsGetOrderbook(wsClient *websocket.Conn, data interface{}) error { wsResp := WebsocketEventResponse{ Event: "GetOrderbook", } var orderbookReq WebsocketOrderbookTickerRequest err := common.JSONDecode(data.([]byte), &orderbookReq) if err != nil { wsResp.Error = err.Error() wsClient.WriteJSON(wsResp) return err } result, err := GetSpecificOrderbook(orderbookReq.Currency, orderbookReq.Exchange, orderbookReq.AssetType) if err != nil { wsResp.Error = err.Error() wsClient.WriteJSON(wsResp) return err } wsResp.Data = result return wsClient.WriteJSON(wsResp) } func wsGetExchangeRates(wsClient *websocket.Conn, data interface{}) error { wsResp := WebsocketEventResponse{ Event: "GetExchangeRates", } if currency.YahooEnabled { wsResp.Data = currency.CurrencyStore } else { wsResp.Data = currency.CurrencyStoreFixer } return wsClient.WriteJSON(wsResp) } func wsGetPortfolio(wsClient *websocket.Conn, data interface{}) error { wsResp := WebsocketEventResponse{ Event: "GetPortfolio", } wsResp.Data = bot.portfolio.GetPortfolioSummary() return wsClient.WriteJSON(wsResp) } // WebsocketHandler Handles websocket client requests func WebsocketHandler() { for { for x := range WebsocketClientHub { var evt WebsocketEvent err := WebsocketClientHub[x].Conn.ReadJSON(&evt) if err != nil { DisconnectWebsocketClient(x, err) continue } if evt.Event == "" { DisconnectWebsocketClient(x, errors.New("Websocket client sent data we did not understand")) continue } dataJSON, err := common.JSONEncode(evt.Data) if err != nil { log.Println(err) continue } req := common.StringToLower(evt.Event) log.Printf("Websocket req: %s", req) if !WebsocketClientHub[x].Authenticated && evt.Event != "auth" { wsResp := WebsocketEventResponse{ Event: "auth", Error: "you must authenticate first", } SendWebsocketMessage(x, wsResp) DisconnectWebsocketClient(x, errors.New("Websocket client did not auth")) continue } else if !WebsocketClientHub[x].Authenticated && evt.Event == "auth" { var auth WebsocketAuth err = common.JSONDecode(dataJSON, &auth) if err != nil { log.Println(err) continue } hashPW := common.HexEncodeToString(common.GetSHA256([]byte(bot.config.Webserver.AdminPassword))) if auth.Username == bot.config.Webserver.AdminUsername && auth.Password == hashPW { WebsocketClientHub[x].Authenticated = true wsResp := WebsocketEventResponse{ Event: "auth", Data: WebsocketResponseSuccess, } SendWebsocketMessage(x, wsResp) log.Println("Websocket client authenticated successfully") continue } else { wsResp := WebsocketEventResponse{ Event: "auth", Error: "invalid username/password", } SendWebsocketMessage(x, wsResp) DisconnectWebsocketClient(x, errors.New("Websocket client sent wrong username/password")) continue } } result, ok := wsHandlers[req] if !ok { log.Printf("Websocket unsupported event") continue } err = result(WebsocketClientHub[x].Conn, dataJSON) if err != nil { log.Printf("Websocket request %s failed. Error %s", evt.Event, err) continue } } time.Sleep(time.Millisecond) } }