toki
d0754a353a
- Rename package/module from toki_socket to proto_socket in Dart, Kotlin, Python - Update crosstest implementations to use renamed packages - Add new proto_socket skill, deprecate add-toki-socket-crosstest-language skill - Update domain rules for all languages - Update documentation (README, PORTING_GUIDE, PROTOCOL, VERSIONING)
116 lines
3.3 KiB
Go
116 lines
3.3 KiB
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/common-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(error) {
|
|
c.onDisconnected()
|
|
})
|
|
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()
|
|
return
|
|
}
|
|
if msgType != websocket.MessageBinary {
|
|
continue
|
|
}
|
|
base := &packets.PacketBase{}
|
|
if err := proto.Unmarshal(b, base); err != nil {
|
|
c.onDisconnected()
|
|
return
|
|
}
|
|
c.OnReceivedData(base.GetTypeName(), base.GetData(), base.GetNonce(), base.GetResponseNonce())
|
|
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)
|
|
}
|