From b375ab3c8c11949bb585d3a5aa361899d83d7095 Mon Sep 17 00:00:00 2001
From: watsonb8 <b.watson@watsonlabs.net>
Date: Mon, 25 Sep 2023 15:25:43 -0500
Subject: [PATCH] Fixing error where connection was re-established incorrectly

---
 internal/dispatcher/dispatcher.go | 9 +++++++++
 transport/websocket/websocket.go  | 4 ++++
 2 files changed, 13 insertions(+)

diff --git a/internal/dispatcher/dispatcher.go b/internal/dispatcher/dispatcher.go
index c90768d..6336bc7 100644
--- a/internal/dispatcher/dispatcher.go
+++ b/internal/dispatcher/dispatcher.go
@@ -113,12 +113,21 @@ func (d *Dispatcher) dispatchMessage(msg *message.Message) {
 				return
 			}
 			log.Println("Websocket terminated: reconnecting")
+			err := d.Disconnect()
+			if err != nil {
+				log.Println(err)
+			}
 			d.Connect()
+			if err != nil {
+				log.Println(err)
+			}
+
 			subsList := d.store.GetAll()
 			for i := range subsList {
 				sub := subsList[i]
 				d.Subscribe(sub.Name(), sub.AuthToken())
 			}
+			return
 		case message.MetaSubscribe:
 			//handle MetaSubscribe resp
 			d.pendingSubsMu.Lock()
diff --git a/transport/websocket/websocket.go b/transport/websocket/websocket.go
index 2087338..ec21ed2 100644
--- a/transport/websocket/websocket.go
+++ b/transport/websocket/websocket.go
@@ -46,6 +46,10 @@ func (w *Websocket) Init(endpoint string, options *transport.Options) error {
 
 	w.stopCh = make(chan error)
 	w.conn, _, err = websocket.DefaultDialer.Dial(endpoint, options.Headers)
+	if err != nil {
+		return err
+	}
+
 	err = w.conn.UnderlyingConn().(*tls.Conn).NetConn().(*net.TCPConn).SetKeepAlive(true)
 	if err != nil {
 		return err