fix: adding SafeWebSocketClient parameter

v1/client/client.go

@@ -71,6 +71,7 @@ type SafeWebsocketClientBuilder struct {
 	rawQuery       *string
 	useTLS         *bool
 	channelSize    *int64
+	authenticateFn func(*SafeWebsocketClient) error
 }
 
 func NewSafeWebsocketClientBuilder() *SafeWebsocketClientBuilder {
@@ -92,6 +93,11 @@ func (b *SafeWebsocketClientBuilder) UseTLS(useTLS bool) *SafeWebsocketClientBui
 	return b
 }
 
+func (b *SafeWebsocketClientBuilder) AuthenticateFn(authenticateFn func(*SafeWebsocketClient) 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(*SafeWebsocketClient) error
 }
 
 func (wsClient *SafeWebsocketClient) connect() error {
@@ -200,15 +211,19 @@ 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 wsClient.authenticateFn != nil {
+		if err := wsClient.authenticateFn(wsClient); 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)
@@ -253,31 +268,32 @@ func (wsClient *SafeWebsocketClient) startReceiveHandler(ctx context.Context) {
 			log.Println("receive handler stopped")
 			return
 		default:
+			if err := wsClient.mu.ReadHandler(func() error {
 				conn := wsClient.conn
 
 				if conn == nil {
 					wsClient.triggerReconnect()
-				return
+					return fmt.Errorf("connection closed")
-				// return fmt.Errorf("connection closed")
 				}
 				_, message, err := conn.ReadMessage()
 				if err != nil {
 					wsClient.triggerReconnect()
-				// return fmt.Errorf("failed to read message: %v", err)
+					return fmt.Errorf("failed to read message: %v", err)
-				return
 				}
 				select {
 				case wsClient.dataChannel <- message:
 				case <-ctx.Done():
 					log.Println("Reconnect handler stopped")
-				return
 				default:
 					log.Println("Data channel full, dropping message")
 				}
+				return nil
+			}); err != nil {
 				return
 			}
 		}
 	}
+}
 
 func (wsClient *SafeWebsocketClient) triggerReconnect() {
 	select {