package foundry import ( "context" "fmt" "log" "net/http" "net/url" "strconv" "strings" "sync" "time" "github.com/gorilla/websocket" ) type webSocketUtil struct { wsConn *websocket.Conn currWsId int isReadReady bool msgMap map[int](chan []byte) mutex sync.Mutex channels Channels } func parseCode(msg *string) string { for j := range RequestCodes { if !strings.HasPrefix(*msg, RequestCodes[j]) { continue } return RequestCodes[j] } return "" } func parseId(msg *string, start int) (int, int) { msgLen := len(*msg) j := start for ; j < msgLen; j++ { if (*msg)[j] < '0' || (*msg)[j] > '9' { break } } if j >= msgLen || j <= 0 { return 0, 0 } msgId, err := strconv.Atoi((*msg)[start:j]) if err != nil { return 0, 0 } return msgId, j } func parseWsRespMessage(msg string) (*wsMessage, error) { data := &wsMessage{} data.code = parseCode(&msg) if data.code != RespDataCode { return data, nil } i := len(data.code) if i == 0 { return nil, ErrorMsgNotHaveNumber } data.id, i = parseId(&msg, i) if i == 0 { return nil, ErrorMsgNotHaveNumber } data.msgJson = msg[i:] return data, nil } func ConnectToWebSocket(foundry *Foundry) error { if !foundry.isAuth { return ErrorNotAuth } query := url.Values{} query.Add("session", foundry.sessionID) query.Add("EIO", "4") query.Add("transport", "websocket") u := url.URL{Scheme: "ws", Host: foundry.config.host, Path: fmt.Sprintf("%s/", socketPath), RawQuery: query.Encode()} wsHeader := http.Header{} wsHeader.Set("Cookie", fmt.Sprintf("session=%s", foundry.sessionID)) wsConn, _, err := websocket.DefaultDialer.Dial(u.String(), wsHeader) if err != nil { return err } foundry.ws.wsConn = wsConn foundry.ws.currWsId = 0 return nil } // TODO: Lookup timeout func (ws *webSocketUtil) ReceiveMessage(id int) ([]byte, error) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() for { select { case msg, ok := <-ws.msgMap[id]: if !ok { time.Sleep(5 * time.Microsecond) continue } ws.closeMsgChannel(id, 0) return msg, nil case <-ctx.Done(): err := ctx.Err() if err != nil { return nil, err } return nil, ErrorTimeout default: time.Sleep(5 * time.Microsecond) } } } func (ws *webSocketUtil) closeMsgChannel(id int, timeout time.Duration) bool { time.Sleep(timeout) ws.mutex.Lock() defer ws.mutex.Unlock() ok := true if _, ok = ws.msgMap[id]; !ok { return false } select { case _, ok = <-ws.msgMap[id]: if ok { close(ws.msgMap[id]) delete(ws.msgMap, id) } default: close(ws.msgMap[id]) delete(ws.msgMap, id) } fmt.Printf("Channel has been closed(id=%d, timeout=%d)\n", id, timeout) return ok } func (ws *webSocketUtil) HandleWebsocketRequest(msg *wsMessage) ([]byte, error) { if !ws.IsReadReady() { return nil, ErrorIsNotReady } err := ws.wsConn.WriteMessage(websocket.TextMessage, msg.toByteSlice()) if err != nil { return nil, err } data, err := ws.ReceiveMessage(msg.id) if err != nil { return nil, err } return data, nil } func (ws *webSocketUtil) sendOnlyCodeRequest(code string) error { msgToSend := []byte(CodesRespToReq[code]) log.Printf("send: %s\n", msgToSend) return ws.wsConn.WriteMessage(websocket.TextMessage, msgToSend) } func (ws *webSocketUtil) ListenWebSocket() { for { _, message, err := ws.wsConn.ReadMessage() if err != nil { ws.channels.Err() <- &FoundryError{Type: ReaderCode, Err: err, IsFatal: true} return } data, err := parseWsRespMessage(string(message)) if err != nil { ws.channels.Err() <- &FoundryError{Type: ReaderCode, Err: err} continue } ws.channels.Msg() <- data } } func (ws *webSocketUtil) ServeWebSocket() { for { select { case message := <-ws.channels.Msg(): log.Println("recv code:", message.code) switch message.code { case RespPingCode, RespSessionData: err := ws.sendOnlyCodeRequest(message.code) if err != nil { ws.isReadReady = false ws.channels.err <- &FoundryError{Type: WriterCode, Err: err} continue } case RespCreateSessionCode: ws.isReadReady = true case RespDataCode: ws.msgMap[message.id] = make(chan []byte, 1) ws.msgMap[message.id] <- []byte(message.msgJson) go ws.closeMsgChannel(message.id, 5*time.Second) default: } case <-ws.channels.Done(): return } } } func (ws *webSocketUtil) InitWsChannels() *Channels { ws.channels.done = make(chan struct{}) ws.channels.err = make(chan error, 10) ws.channels.msg = make(chan *wsMessage, 10) return &ws.channels } func NewWebSocketUtil() *webSocketUtil { return &webSocketUtil{currWsId: 0, isReadReady: false, msgMap: make(map[int]chan []byte), channels: Channels{}} } func (ws *webSocketUtil) IsReadReady() bool { return ws.isReadReady }