Add custom connection handler for websocket keepalive

Allows passing a ConnHandler to wsServeWithConnHandler for custom
connection management, enabling active keepalive via ping messages.
WsAnnouncementServe now uses this to send periodic pings.

Author: 0x0001

v2/websocket.go

@@ -1,6 +1,7 @@
 package binance
 
 import (
+	"context"
 	"net/http"
 	"net/url"
 	"time"
@@ -28,7 +29,14 @@ func newWsConfig(endpoint string) *WsConfig {
 	}
 }
 
-var wsServe = func(cfg *WsConfig, handler WsHandler, errHandler ErrHandler) (doneC, stopC chan struct{}, err error) {
+func wsServe(cfg *WsConfig, handler WsHandler, errHandler ErrHandler) (doneC, stopC chan struct{}, err error) {
+	return wsServeWithConnHandler(cfg, handler, errHandler, nil)
+}
+
+type ConnHandler func(context.Context, *websocket.Conn)
+
+// WsServeWithConnHandler serves websocket with custom connection handler, useful for custom keepalive
+var wsServeWithConnHandler = func(cfg *WsConfig, handler WsHandler, errHandler ErrHandler, connHandler ConnHandler) (doneC, stopC chan struct{}, err error) {
 	proxy := http.ProxyFromEnvironment
 	if cfg.Proxy != nil {
 		u, err := url.Parse(*cfg.Proxy)
@@ -62,6 +70,13 @@ var wsServe = func(cfg *WsConfig, handler WsHandler, errHandler ErrHandler) (don
 			keepAlive(c, WebsocketTimeout)
 		}
 
+		// Custom connection handling, useful in active keepalive scenarios
+		if connHandler != nil {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			go connHandler(ctx, c)
+		}
+
 		// Wait for the stopC channel to be closed.  We do that in a
 		// separate goroutine because ReadMessage is a blocking
 		// operation.
@@ -88,6 +103,21 @@ var wsServe = func(cfg *WsConfig, handler WsHandler, errHandler ErrHandler) (don
 	return
 }
 
+func keepAliveHandler(duration time.Duration) ConnHandler {
+	return func(ctx context.Context, c *websocket.Conn) {
+		ticker := time.NewTicker(duration)
+		for {
+			select {
+			case <-ctx.Done():
+				ticker.Stop()
+				return
+			case <-ticker.C:
+				c.WriteMessage(websocket.PingMessage, []byte{})
+			}
+		}
+	}
+}
+
 func keepAlive(c *websocket.Conn, timeout time.Duration) {
 	ticker := time.NewTicker(timeout)
 
@@ -137,7 +167,7 @@ var WsGetReadWriteConnection = func(cfg *WsConfig) (*websocket.Conn, error) {
 		EnableCompression: false,
 	}
 
-	c, _, err := Dialer.Dial(cfg.Endpoint, nil)
+	c, _, err := Dialer.Dial(cfg.Endpoint, *cfg.Header)
 	if err != nil {
 		return nil, err
 	}

v2/websocket_service.go

@@ -946,7 +946,7 @@ func WsAnnouncementServe(params WsAnnouncementParam, handler WsAnnouncementHandl
 		json.Unmarshal([]byte(event.Data), &e)
 		handler(e)
 	}
-	return wsServe(cfg, wsHandler, errHandler)
+	return wsServeWithConnHandler(cfg, wsHandler, errHandler, keepAliveHandler(30*time.Second))
 }
 
 // getWsApiEndpoint return the base endpoint of the API WS according the UseTestnet flag

v2/websocket_service_test.go

@@ -9,7 +9,7 @@ import (
 
 type websocketServiceTestSuite struct {
 	baseTestSuite
-	origWsServe func(*WsConfig, WsHandler, ErrHandler) (chan struct{}, chan struct{}, error)
+	origWsServe func(*WsConfig, WsHandler, ErrHandler, ConnHandler) (chan struct{}, chan struct{}, error)
 	serveCount  int
 }
 
@@ -18,16 +18,16 @@ func TestWebsocketService(t *testing.T) {
 }
 
 func (s *websocketServiceTestSuite) SetupTest() {
-	s.origWsServe = wsServe
+	s.origWsServe = wsServeWithConnHandler
 }
 
 func (s *websocketServiceTestSuite) TearDownTest() {
-	wsServe = s.origWsServe
+	wsServeWithConnHandler = s.origWsServe
 	s.serveCount = 0
 }
 
 func (s *websocketServiceTestSuite) mockWsServe(data []byte, err error) {
-	wsServe = func(cfg *WsConfig, handler WsHandler, errHandler ErrHandler) (doneC, stopC chan struct{}, innerErr error) {
+	wsServeWithConnHandler = func(cfg *WsConfig, handler WsHandler, errHandler ErrHandler, connHandler ConnHandler) (doneC, stopC chan struct{}, innerErr error) {
 		s.serveCount++
 		doneC = make(chan struct{})
 		stopC = make(chan struct{})