proto-socket/go/ws_client.go

package proto_socket

import (
	"context"
	"crypto/tls"
	"fmt"
	"net/http"
	"sync"

	"google.golang.org/protobuf/proto"
	"nhooyr.io/websocket"

	"git.toki-labs.com/toki/proto-socket/go/packets"
)

const (
	DefaultHeartbeatIntervalSec = 30
	DefaultHeartbeatWaitSec     = 10
)

type WsClient struct {
	baseClient[*WsClient]
	conn        *websocket.Conn
	writeMu     sync.Mutex // Defensive if WritePacket is called outside the queued writer.
	readCtx     context.Context
	cancelRead  context.CancelFunc
	writeCtx    context.Context
	cancelWrite context.CancelFunc
}

func NewWsClient(conn *websocket.Conn, intervalSec, waitSec int, parserMap ParserMap) *WsClient {
	readCtx, cancelRead := context.WithCancel(context.Background())
	writeCtx, cancelWrite := context.WithCancel(context.Background())
	c := &WsClient{
		conn:        conn,
		readCtx:     readCtx,
		cancelRead:  cancelRead,
		writeCtx:    writeCtx,
		cancelWrite: cancelWrite,
	}
	c.baseClient = newBaseClient(c, intervalSec, waitSec, func() error {
		c.cancelRead()
		c.cancelWrite()
		return c.conn.Close(websocket.StatusNormalClosure, "")
	})
	c.Communicator.Initialize(c, parserMap)
	c.SetWriteErrorHandler(func(err error) {
		c.onDisconnected(DisconnectReasonWriteError, err)
	})
	c.SetFrameErrorHandler(func(err error) {
		c.onDisconnected(DisconnectReasonWSPacketParse, err)
	})
	c.AddListener(TypeNameOf(&packets.HeartBeat{}), func(m proto.Message) {
		c.onHeartBeat()
	})
	go c.readLoop()
	c.sendHeartBeat()
	return c
}

func DialWs(ctx context.Context, host string, port int, path string, parserMap ParserMap) (*WsClient, error) {
	return DialWsWithHeartbeat(ctx, host, port, path, DefaultHeartbeatIntervalSec, DefaultHeartbeatWaitSec, parserMap)
}

func DialWsWithHeartbeat(ctx context.Context, host string, port int, path string, intervalSec, waitSec int, parserMap ParserMap) (*WsClient, error) {
	conn, _, err := websocket.Dial(ctx, fmt.Sprintf("ws://%s:%d%s", host, port, path), nil)
	if err != nil {
		return nil, err
	}
	return NewWsClient(conn, intervalSec, waitSec, parserMap), nil
}

func DialWss(ctx context.Context, host string, port int, path string, tlsCfg *tls.Config, parserMap ParserMap) (*WsClient, error) {
	return DialWssWithHeartbeat(ctx, host, port, path, tlsCfg, DefaultHeartbeatIntervalSec, DefaultHeartbeatWaitSec, parserMap)
}

func DialWssWithHeartbeat(ctx context.Context, host string, port int, path string, tlsCfg *tls.Config, intervalSec, waitSec int, parserMap ParserMap) (*WsClient, error) {
	opts := &websocket.DialOptions{}
	if tlsCfg != nil {
		opts.HTTPClient = &http.Client{
			Transport: &http.Transport{TLSClientConfig: tlsCfg},
		}
	}
	conn, _, err := websocket.Dial(ctx, fmt.Sprintf("wss://%s:%d%s", host, port, path), opts)
	if err != nil {
		return nil, err
	}
	return NewWsClient(conn, intervalSec, waitSec, parserMap), nil
}

func (c *WsClient) readLoop() {
	for c.IsAlive() {
		msgType, b, err := c.conn.Read(c.readCtx)
		if err != nil {
			c.onDisconnected(DisconnectReasonWSRead, err)
			return
		}
		if msgType != websocket.MessageBinary {
			continue
		}
		// OnReceivedFrame routes the raw frame through the inbound gateway when
		// one is attached, or decodes it inline otherwise. An inline decode
		// error preserves the existing parse-error disconnect semantics.
		if err := c.OnReceivedFrame(b); err != nil {
			c.onDisconnected(DisconnectReasonWSPacketParse, err)
			return
		}
		c.sendHeartBeat()
	}
}

func (c *WsClient) WritePacket(base *packets.PacketBase) error {
	b, err := proto.Marshal(base)
	if err != nil {
		return err
	}
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	return c.conn.Write(c.writeCtx, websocket.MessageBinary, b)
}