From 71164de3e6467e14429f84c5e320793a006e00ff Mon Sep 17 00:00:00 2001
From: StounhandJ <rmarinicev@gmail.com>
Date: Wed, 3 May 2023 19:24:09 +0300
Subject: [PATCH 1/2] add headers

---
 websocket.go | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/websocket.go b/websocket.go
index 1829001..633e3d0 100644
--- a/websocket.go
+++ b/websocket.go
@@ -118,6 +118,11 @@ func New(handler func(*Conn), config ...Config) fiber.Handler {
         c.Context().Request.Header.VisitAllCookie(func(key, value []byte) {
             conn.cookies[string(key)] = string(value)
         })
+        
+        // headers
+		c.Context().Request.Header.VisitAll(func(key, value []byte) {
+			conn.headers[string(key)] = string(value)
+		})
 
         if err := upgrader.Upgrade(c.Context(), func(fconn *websocket.Conn) {
             conn.Conn = fconn
@@ -137,6 +142,7 @@ type Conn struct {
     locals  map[string]interface{}
     params  map[string]string
     cookies map[string]string
+    headers map[string]string
     queries map[string]string
 }
 
@@ -154,6 +160,7 @@ func acquireConn() *Conn {
     conn.params = make(map[string]string)
     conn.queries = make(map[string]string)
     conn.cookies = make(map[string]string)
+    conn.headers = make(map[string]string)
     return conn
 }
 
@@ -202,6 +209,17 @@ func (conn *Conn) Cookies(key string, defaultValue ...string) string {
     return v
 }
 
+// Headers is used for getting a header value by key
+// Defaults to empty string "" if the header doesn't exist.
+// If a default value is given, it will return that value if the header doesn't exist.
+func (conn *Conn) Headers(key string, defaultValue ...string) string {
+	v, ok := conn.headers[key]
+	if !ok && len(defaultValue) > 0 {
+		return defaultValue[0]
+	}
+	return v
+}
+
 // Constants are taken from https://github.com/fasthttp/websocket/blob/master/conn.go#L43
 
 // Close codes defined in RFC 6455, section 11.7.

From bc3217f0058c691ffad6681cc824d7ad308dadf4 Mon Sep 17 00:00:00 2001
From: Roman Marinichev <58305548+StounhandJ@users.noreply.github.com>
Date: Wed, 3 May 2023 19:35:19 +0300
Subject: [PATCH 2/2] tab correction

---
 websocket.go | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/websocket.go b/websocket.go
index 633e3d0..1d1519f 100644
--- a/websocket.go
+++ b/websocket.go
@@ -120,9 +120,9 @@ func New(handler func(*Conn), config ...Config) fiber.Handler {
         })
         
         // headers
-		c.Context().Request.Header.VisitAll(func(key, value []byte) {
-			conn.headers[string(key)] = string(value)
-		})
+	c.Context().Request.Header.VisitAll(func(key, value []byte) {
+	    conn.headers[string(key)] = string(value)
+	})
 
         if err := upgrader.Upgrade(c.Context(), func(fconn *websocket.Conn) {
             conn.Conn = fconn