init

internal/foundry/websocket.go

@@ -0,0 +1,236 @@
+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
+}