Compare commits

...

2 Commits

View File

@@ -71,6 +71,7 @@ type SafeWebsocketClientBuilder struct {
rawQuery *string
useTLS *bool
channelSize *int64
authenticateFn func() error
}
func NewSafeWebsocketClientBuilder() *SafeWebsocketClientBuilder {
@@ -92,6 +93,11 @@ func (b *SafeWebsocketClientBuilder) UseTLS(useTLS bool) *SafeWebsocketClientBui
return b
}
func (b *SafeWebsocketClientBuilder) AuthenticateFn(authenticateFn func() error) *SafeWebsocketClientBuilder {
b.authenticateFn = authenticateFn
return b
}
func (b *SafeWebsocketClientBuilder) Path(path string) *SafeWebsocketClientBuilder {
b.path = &path
return b
@@ -136,6 +142,10 @@ func (b *SafeWebsocketClientBuilder) Build(ctx context.Context) (*SafeWebsocketC
doneMap: NewSafeMap[string, chan struct{}](),
}
if b.authenticateFn != nil {
wsClient.authenticateFn = b.authenticateFn
}
go wsClient.reconnectHandler()
if err := wsClient.connect(); err != nil {
@@ -160,6 +170,7 @@ type SafeWebsocketClient struct {
reconnectChans []chan struct{}
isConnected bool
doneMap *SafeMap[string, chan struct{}]
authenticateFn func() error
}
func (wsClient *SafeWebsocketClient) connect() error {
@@ -200,15 +211,18 @@ func (wsClient *SafeWebsocketClient) connect() error {
return nil
})
wsClient.mu.WriteHandler(func() error {
if wsClient.conn != nil {
wsClient.conn.Close()
}
wsClient.conn = conn
wsClient.isConnected = true
if err := wsClient.authenticateFn(); err != nil {
return err
}
wsClient.mu.WriteHandler(func() error {
pingCtx, pingCancel := context.WithCancel(context.Background())
go wsClient.startPingTicker(pingCtx)
wsClient.cancelFuncs = append(wsClient.cancelFuncs, pingCancel)
@@ -269,7 +283,6 @@ func (wsClient *SafeWebsocketClient) startReceiveHandler(ctx context.Context) {
case wsClient.dataChannel <- message:
case <-ctx.Done():
log.Println("Reconnect handler stopped")
return nil
default:
log.Println("Data channel full, dropping message")
}