Foundry-Scrapping-API/internal/foundry/websocket.go

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
}