From bbc5af6a3cae3a8e509b0ffadfd3830d1399f24b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lucas=20L=C3=B6ffel?= <lucas@loeffel.io>
Date: Tue, 9 Oct 2018 11:08:15 +0200
Subject: [PATCH] added connect handler and GetURL method

---
 recws.go | 44 ++++++++++++++++++++++++++++++++------------
 1 file changed, 32 insertions(+), 12 deletions(-)

diff --git a/recws.go b/recws.go
index 63c91aa..35d08b3 100644
--- a/recws.go
+++ b/recws.go
@@ -35,6 +35,8 @@ type RecConn struct {
 	HandshakeTimeout time.Duration
 	// NonVerbose suppress connecting/reconnecting messages.
 	NonVerbose bool
+	// ConnectHandler fires after the connection successfully establish.
+	ConnectHandler func(rc *RecConn) error
 
 	mu          sync.Mutex
 	url         string
@@ -50,10 +52,7 @@ type RecConn struct {
 // CloseAndReconnect will try to reconnect.
 func (rc *RecConn) closeAndReconnect() {
 	rc.Close()
-	go func() {
-		rc.connect()
-	}()
-
+	go rc.connect()
 }
 
 // Close closes the underlying network connection without
@@ -180,14 +179,17 @@ func (rc *RecConn) Dial(urlStr string, reqHeader http.Header) {
 	rc.dialer = websocket.DefaultDialer
 	rc.dialer.HandshakeTimeout = rc.HandshakeTimeout
 
-	go func() {
-		rc.connect()
-	}()
+	go rc.connect()
 
 	// wait on first attempt
 	time.Sleep(rc.HandshakeTimeout)
 }
 
+// GetURL returns current connection url
+func (rc *RecConn) GetURL() string {
+	return rc.url
+}
+
 func (rc *RecConn) connect() {
 	b := &backoff.Backoff{
 		Min:    rc.RecIntvlMin,
@@ -214,12 +216,23 @@ func (rc *RecConn) connect() {
 			if !rc.NonVerbose {
 				log.Printf("Dial: connection was successfully established with %s\n", rc.url)
 			}
-			break
-		} else {
-			if !rc.NonVerbose {
-				log.Println(err)
-				log.Println("Dial: will try again in", nextItvl, "seconds.")
+
+			if rc.ConnectHandler == nil {
+				break
 			}
+
+			if err := rc.ConnectHandler(rc); err != nil {
+				log.Fatalf("Dial: connect handler failed with %s", err.Error())
+			}
+
+			log.Printf("Dial: connect handler was successfully established with %s\n", rc.url)
+
+			return
+		}
+
+		if !rc.NonVerbose {
+			log.Println(err)
+			log.Println("Dial: will try again in", nextItvl, "seconds.")
 		}
 
 		time.Sleep(nextItvl)
@@ -252,3 +265,10 @@ func (rc *RecConn) IsConnected() bool {
 
 	return rc.isConnected
 }
+
+func (rc *RecConn) GetConnection() *RecConn {
+	rc.mu.Lock()
+	defer rc.mu.Unlock()
+
+	return rc
+}