From dc27611b06b7e790c8dbf315c6077feb98fd981d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dar=C3=ADo?= <dgrripoll@gmail.com>
Date: Sat, 21 Nov 2020 21:02:58 +0100
Subject: [PATCH] Added UpgradeAsClient func

---
 client.go | 21 ++++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/client.go b/client.go
index 95401db..ffbbed7 100644
--- a/client.go
+++ b/client.go
@@ -28,7 +28,10 @@ func ClientWithHeaders(c net.Conn, url string, req *fasthttp.Request) (*Conn, er
 	return client(c, url, req)
 }
 
-func client(c net.Conn, url string, r *fasthttp.Request) (conn *Conn, err error) {
+// UpgradeAsClient will upgrade the connection as a client
+//
+// r can be nil.
+func UpgradeAsClient(c net.Conn, url string, r *fasthttp.Request) error {
 	req := fasthttp.AcquireRequest()
 	res := fasthttp.AcquireResponse()
 	uri := fasthttp.AcquireURI()
@@ -64,16 +67,24 @@ func client(c net.Conn, url string, r *fasthttp.Request) (conn *Conn, err error)
 	bw := bufio.NewWriter(c)
 	req.Write(bw)
 	bw.Flush()
-	err = res.Read(br)
+	err := res.Read(br)
 	if err == nil {
 		if res.StatusCode() != 101 ||
 			!equalsFold(res.Header.PeekBytes(upgradeString), websocketString) {
 			err = ErrCannotUpgrade
-		} else {
-			conn = acquireConn(c)
-			conn.server = false
 		}
 	}
+
+	return err
+}
+
+func client(c net.Conn, url string, r *fasthttp.Request) (conn *Conn, err error) {
+	err = UpgradeAsClient(c, url, r)
+	if err == nil {
+		conn = acquireConn(c)
+		conn.server = false
+	}
+
 	return conn, err
 }