reconnect websocket to whatsapp when it got closed

Author: krombel

conn.go

@@ -6,16 +6,17 @@ import (
 	"crypto/sha256"
 	"encoding/json"
 	"fmt"
-	"github.com/Rhymen/go-whatsapp/binary"
-	"github.com/Rhymen/go-whatsapp/crypto/cbc"
-	"github.com/gorilla/websocket"
 	"math/rand"
 	"net/http"
 	"os"
 	"strconv"
 	"strings"
 	"sync"
 	"time"
+
+	"github.com/Rhymen/go-whatsapp/binary"
+	"github.com/Rhymen/go-whatsapp/crypto/cbc"
+	"github.com/gorilla/websocket"
 )
 
 type metric byte
@@ -79,6 +80,7 @@ It holds all necessary information to make the package work internally.
 */
 type Conn struct {
 	wsConn         *websocket.Conn
+	wsConnMutex    sync.RWMutex
 	session        *Session
 	listener       map[string]chan string
 	listenerMutex  sync.RWMutex
@@ -104,20 +106,9 @@ Creates a new connection with a given timeout. The websocket connection to the W
 The goroutine for handling incoming messages is started
 */
 func NewConn(timeout time.Duration) (*Conn, error) {
-	dialer := &websocket.Dialer{
-		ReadBufferSize:   25 * 1024 * 1024,
-		WriteBufferSize:  10 * 1024 * 1024,
-		HandshakeTimeout: timeout,
-	}
-
-	headers := http.Header{"Origin": []string{"https://web.whatsapp.com"}}
-	wsConn, _, err := dialer.Dial("wss://w3.web.whatsapp.com/ws", headers)
-	if err != nil {
-		return nil, fmt.Errorf("couldn't dial whatsapp web websocket: %v", err)
-	}
-
 	wac := &Conn{
-		wsConn:        wsConn,
+		wsConn:        nil, // will be set in connect()
+		wsConnMutex:   sync.RWMutex{},
 		listener:      make(map[string]chan string),
 		listenerMutex: sync.RWMutex{},
 		writeChan:     make(chan wsMsg),
@@ -130,13 +121,77 @@ func NewConn(timeout time.Duration) (*Conn, error) {
 		shortClientName: "go-whatsapp",
 	}
 
+	if err := wac.connect(); err != nil {
+		return nil, err
+	}
+
 	go wac.readPump()
 	go wac.writePump()
 	go wac.keepAlive(20000, 90000)
 
 	return wac, nil
 }
 
+func (wac *Conn) isConnected() bool {
+	wac.wsConnMutex.RLock()
+	defer wac.wsConnMutex.RUnlock()
+	return wac.wsConn != nil
+}
+
+// connect should be guarded with wsConnMutex
+func (wac *Conn) connect() error {
+	dialer := &websocket.Dialer{
+		ReadBufferSize:   25 * 1024 * 1024,
+		WriteBufferSize:  10 * 1024 * 1024,
+		HandshakeTimeout: wac.msgTimeout,
+	}
+
+	headers := http.Header{"Origin": []string{"https://web.whatsapp.com"}}
+	wsConn, _, err := dialer.Dial("wss://w3.web.whatsapp.com/ws", headers)
+	if err != nil {
+		return fmt.Errorf("couldn't dial whatsapp web websocket: %v", err)
+	}
+
+	wsConn.SetCloseHandler(func(code int, text string) error {
+		fmt.Fprintf(os.Stderr, "websocket connection closed(%d, %s)\n", code, text)
+
+		// from default CloseHandler
+		message := websocket.FormatCloseMessage(code, "")
+		wsConn.WriteControl(websocket.CloseMessage, message, time.Now().Add(time.Second))
+
+		// our close handling
+		if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
+			fmt.Println("Trigger reconnect")
+			go wac.reconnect()
+		}
+		return nil
+	})
+
+	wac.wsConn = wsConn
+	return nil
+}
+
+// reconnect should be run as go routine
+func (wac *Conn) reconnect() {
+	wac.wsConnMutex.Lock()
+	wac.wsConn = nil
+	wac.wsConnMutex.Unlock()
+
+	// wait up to 60 seconds and then reconnect. As writePump should send immediately, it might
+	// reconnect as well. So we check its existance before reconnecting
+	for !wac.isConnected() {
+		time.Sleep(time.Duration(rand.Intn(60)) * time.Second)
+
+		wac.wsConnMutex.Lock()
+		if wac.wsConn == nil {
+			if err := wac.connect(); err != nil {
+				fmt.Fprintf(os.Stderr, "could not reconnect to websocket: %v\n", err)
+			}
+		}
+		wac.wsConnMutex.Unlock()
+	}
+}
+
 func (wac *Conn) write(data []interface{}) (<-chan string, error) {
 	d, err := json.Marshal(data)
 	if err != nil {
@@ -199,6 +254,9 @@ func (wac *Conn) readPump() {
 	defer wac.wsConn.Close()
 
 	for {
+		for !wac.isConnected() {
+			time.Sleep(1 * time.Second)
+		}
 		msgType, msg, err := wac.wsConn.ReadMessage()
 		if err != nil {
 			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
@@ -249,8 +307,26 @@ func (wac *Conn) readPump() {
 
 func (wac *Conn) writePump() {
 	for msg := range wac.writeChan {
+		for !wac.isConnected() {
+			// reconnect to send the message ASAP
+			wac.wsConnMutex.Lock()
+			if wac.wsConn == nil {
+				if err := wac.connect(); err != nil {
+					fmt.Fprintf(os.Stderr, "could not reconnect to websocket: %v\n", err)
+				}
+			}
+			wac.wsConnMutex.Unlock()
+			if !wac.isConnected() {
+				// reconnecting failed. Sleep for a while and try again afterwards
+				time.Sleep(time.Duration(rand.Intn(5)) * time.Second)
+			}
+		}
 		if err := wac.wsConn.WriteMessage(msg.messageType, msg.data); err != nil {
-			fmt.Fprintf(os.Stderr, "error writing to socket: %v", err)
+			fmt.Fprintf(os.Stderr, "error writing to socket: %v\n", err)
+			// add message to channel again to no loose it
+			go func() {
+				wac.writeChan <- msg
+			}()
 		}
 	}
 }