From 509f9ac8caf38ae950501bfec4f4804d4d4a9082 Mon Sep 17 00:00:00 2001
From: Edwin <edwin@am.is>
Date: Thu, 10 Aug 2023 15:08:31 +0800
Subject: [PATCH] pkg/types: add BeforeConnect hook function

---
 pkg/types/stream.go | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/pkg/types/stream.go b/pkg/types/stream.go
index 6bcd5e39e..668a12856 100644
--- a/pkg/types/stream.go
+++ b/pkg/types/stream.go
@@ -44,6 +44,8 @@ type Dispatcher func(e interface{})
 // HeartBeat keeps connection alive by sending the heartbeat packet.
 type HeartBeat func(ctxConn context.Context, conn *websocket.Conn, cancelConn context.CancelFunc)
 
+type BeforeConnect func(ctx context.Context) error
+
 //go:generate callbackgen -type StandardStream -interface
 type StandardStream struct {
 	parser     Parser
@@ -111,6 +113,8 @@ type StandardStream struct {
 	FuturesPositionSnapshotCallbacks []func(futuresPositions FuturesPositionMap)
 
 	heartBeat HeartBeat
+
+	beforeConnect BeforeConnect
 }
 
 type StandardStreamEmitter interface {
@@ -306,6 +310,11 @@ func (s *StandardStream) Reconnect() {
 
 // Connect starts the stream and create the websocket connection
 func (s *StandardStream) Connect(ctx context.Context) error {
+	if s.beforeConnect != nil {
+		if err := s.beforeConnect(ctx); err != nil {
+			return err
+		}
+	}
 	err := s.DialAndConnect(ctx)
 	if err != nil {
 		return err
@@ -432,6 +441,11 @@ func (s *StandardStream) SetHeartBeat(fn HeartBeat) {
 	s.heartBeat = fn
 }
 
+// SetBeforeConnect sets the custom hook function before connect
+func (s *StandardStream) SetBeforeConnect(fn BeforeConnect) {
+	s.beforeConnect = fn
+}
+
 type Depth string
 
 const (