Bugfix: Websocket ping/pong improvements (#406)

* Renames func. Creates new func to setup pinghander either via gorilla style or our own

* Cleans up all ping pong handlers.......

* Clears up issues, makes naming a bit better

* Adds tests

* Adds ping support to binance

* Cleans up ping pongs and adds a comment

* Cleans up waitgroup stuff.

* DISCREETLY cleans up woeful function

* Fixes Kraken ping message type. Removes unnecessary test property. Adds `if err == websocket.ErrCloseSent {` to ping func

* +1 for +v

exchanges/websocket/wshandler/wshandler.go

@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"compress/flate"
 	"compress/gzip"
-	"encoding/json"
 	"errors"
 	"fmt"
 	"io/ioutil"
@@ -656,32 +655,82 @@ func (w *WebsocketConnection) Dial(dialer *websocket.Dialer, headers http.Header
 	return nil
 }
 
-// SendMessage the one true message request. Sends message to WS
-func (w *WebsocketConnection) SendMessage(data interface{}) error {
+// SendJSONMessage sends a JSON encoded message over the connection
+func (w *WebsocketConnection) SendJSONMessage(data interface{}) error {
 	w.Lock()
 	defer w.Unlock()
 	if !w.IsConnected() {
 		return fmt.Errorf("%v cannot send message to a disconnected websocket", w.ExchangeName)
 	}
-	json, err := json.Marshal(data)
-	if err != nil {
-		return err
-	}
 	if w.Verbose {
 		log.Debugf(log.WebsocketMgr,
-			"%v sending message to websocket %v", w.ExchangeName, string(json))
+			"%v sending message to websocket %+v", w.ExchangeName, data)
 	}
 	if w.RateLimit > 0 {
 		time.Sleep(time.Duration(w.RateLimit) * time.Millisecond)
 	}
-	return w.Connection.WriteMessage(websocket.TextMessage, json)
+	return w.Connection.WriteJSON(data)
+}
+
+// SendRawMessage sends a message over the connection without JSON encoding it
+func (w *WebsocketConnection) SendRawMessage(messageType int, message []byte) error {
+	w.Lock()
+	defer w.Unlock()
+	if !w.IsConnected() {
+		return fmt.Errorf("%v cannot send message to a disconnected websocket", w.ExchangeName)
+	}
+	if w.Verbose {
+		log.Debugf(log.WebsocketMgr,
+			"%v sending message to websocket %s", w.ExchangeName, message)
+	}
+	if w.RateLimit > 0 {
+		time.Sleep(time.Duration(w.RateLimit) * time.Millisecond)
+	}
+	return w.Connection.WriteMessage(messageType, message)
+}
+
+// SetupPingHandler will automatically send ping or pong messages based on
+// WebsocketPingHandler configuration
+func (w *WebsocketConnection) SetupPingHandler(handler WebsocketPingHandler) {
+	if handler.UseGorillaHandler {
+		h := func(msg string) error {
+			err := w.Connection.WriteControl(handler.MessageType, []byte(msg), time.Now().Add(handler.Delay))
+			if err == websocket.ErrCloseSent {
+				return nil
+			} else if e, ok := err.(net.Error); ok && e.Temporary() {
+				return nil
+			}
+			return err
+		}
+		w.Connection.SetPingHandler(h)
+		return
+	}
+	w.Wg.Add(1)
+	defer w.Wg.Done()
+	go func() {
+		ticker := time.NewTicker(handler.Delay)
+		for {
+			select {
+			case <-w.Shutdown:
+				ticker.Stop()
+				return
+			case <-ticker.C:
+				err := w.SendRawMessage(handler.MessageType, handler.Message)
+				if err != nil {
+					log.Errorf(log.WebsocketMgr,
+						"%v failed to send message to websocket %s", w.ExchangeName, handler.Message)
+					return
+				}
+			}
+		}
+	}()
 }
 
 // SendMessageReturnResponse will send a WS message to the connection
 // It will then run a goroutine to await a JSON response
 // If there is no response it will return an error
 func (w *WebsocketConnection) SendMessageReturnResponse(id int64, request interface{}) ([]byte, error) {
-	err := w.SendMessage(request)
+	err := w.SendJSONMessage(request)
 	if err != nil {
 		return nil, err
 	}

exchanges/websocket/wshandler/wshandler_test.go

@@ -156,6 +156,16 @@ func TestWebsocket(t *testing.T) {
 		t.Error("WebsocketSetup")
 	}
 
+	ws.setEnabled(false)
+	if ws.IsEnabled() {
+		t.Error("WebsocketSetup")
+	}
+
+	ws.setEnabled(true)
+	if !ws.IsEnabled() {
+		t.Error("WebsocketSetup")
+	}
+
 	if ws.GetProxyAddress() != "testProxy" {
 		t.Error("WebsocketSetup")
 	}
@@ -569,7 +579,11 @@ func TestSendMessage(t *testing.T) {
 				}
 				t.Fatal(err)
 			}
-			err = testData.WC.SendMessage("ping")
+			err = testData.WC.SendJSONMessage(Ping)
+			if err != nil {
+				t.Error(err)
+			}
+			err = testData.WC.SendRawMessage(websocket.TextMessage, []byte(Ping))
 			if err != nil {
 				t.Error(err)
 			}
@@ -602,6 +616,42 @@ func TestSendMessageWithResponse(t *testing.T) {
 	}
 }
 
+// TestSetupPingHandler logic test
+func TestSetupPingHandler(t *testing.T) {
+	if wc.ProxyURL != "" && !useProxyTests {
+		t.Skip("Proxy testing not enabled, skipping")
+	}
+	wc.Shutdown = make(chan struct{})
+	err := wc.Dial(&dialer, http.Header{})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	wc.SetupPingHandler(WebsocketPingHandler{
+		UseGorillaHandler: true,
+		MessageType:       websocket.PingMessage,
+		Delay:             1000,
+	})
+
+	err = wc.Connection.Close()
+	if err != nil {
+		t.Error(err)
+	}
+
+	err = wc.Dial(&dialer, http.Header{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	wc.SetupPingHandler(WebsocketPingHandler{
+		MessageType: websocket.TextMessage,
+		Message:     []byte(Ping),
+		Delay:       200,
+	})
+	time.Sleep(time.Millisecond * 500)
+	close(wc.Shutdown)
+	wc.Wg.Wait()
+}
+
 // TestParseBinaryResponse logic test
 func TestParseBinaryResponse(t *testing.T) {
 	var b bytes.Buffer

exchanges/websocket/wshandler/wshandler_types.go

@@ -19,6 +19,8 @@ const (
 	// connection monitor time delays and limits
 	connectionMonitorDelay             = 2 * time.Second
 	WebsocketNotAuthenticatedUsingRest = "%v - Websocket not authenticated, using REST"
+	Ping                               = "ping"
+	Pong                               = "pong"
 )
 
 // Websocket defines a return type for websocket connections via the interface
@@ -164,3 +166,11 @@ type WebsocketConnection struct {
 	ResponseMaxLimit     time.Duration
 	TrafficTimeout       time.Duration
 }
+
+// WebsocketPingHandler container for ping handler settings
+type WebsocketPingHandler struct {
+	UseGorillaHandler bool
+	MessageType       int
+	Message           []byte
+	Delay             time.Duration
+}