package foundry import ( "context" "fmt" "log/slog" "strconv" "strings" "sync" "time" "gitea.local.lab/Lbenedar/foundry_helper_service/internal/foundry/requests" "gitea.local.lab/Lbenedar/foundry_helper_service/internal/foundry/types" "github.com/gorilla/websocket" ) type TransportCode int const ( WriterCode = TransportCode(0) ReaderCode = TransportCode(1) ) type webSocketUtil struct { wsConn *websocket.Conn currWsId int isReadReady bool logger *slog.Logger msgMap map[int](chan []byte) chanMutex sync.Mutex channels types.Channels } func NewWebSocketUtil() *webSocketUtil { return &webSocketUtil{currWsId: 0, isReadReady: false, msgMap: make(map[int]chan []byte)} } 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) (*types.WsMessage, error) { data := &types.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 } // 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.chanMutex.Lock() defer ws.chanMutex.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) } ws.logger.Debug("WS: Channel has been closed\n", "id", id, "timeout", timeout.String()) return ok } func (ws *webSocketUtil) HandleWebsocketRequest(msg *types.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 { ws.logger.Debug("WS: Data has been send\n", "msg", CodesRespToReq[code]) return ws.wsConn.WriteMessage(websocket.TextMessage, []byte(CodesRespToReq[code])) } 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(): ws.logger.Debug("WS: Data has been received\n", "msgCode", 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) IsReadReady() bool { return ws.isReadReady } func (ws *webSocketUtil) CreateWSMessageByPage(page string) *types.WsMessage { msgToSend := &types.WsMessage{Code: CodesRespToReq[RespCreateSessionCode], Id: ws.currWsId, MsgJson: fmt.Sprintf("[\"%s\"]", requests.WsTypeData[page])} ws.currWsId++ ws.logger.Debug("WS: Data to send\n", "msgToSend", msgToSend.ToString()) return msgToSend }