From 512c0b1b7e12e138d0bc4cbf0e20c7a63074a1ed Mon Sep 17 00:00:00 2001
From: crusader <crusader@loafle.com>
Date: Fri, 25 Aug 2017 17:48:14 +0900
Subject: [PATCH] ing

---
 websocket/conn.go   | 19 +++++++++++++++++++
 websocket/server.go |  2 +-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/websocket/conn.go b/websocket/conn.go
index 97e1dba..c52f259 100644
--- a/websocket/conn.go
+++ b/websocket/conn.go
@@ -16,6 +16,8 @@ import (
 	"sync"
 	"time"
 	"unicode/utf8"
+
+	"github.com/valyala/fasthttp"
 )
 
 const (
@@ -262,6 +264,8 @@ type Conn struct {
 
 	readDecompress         bool // whether last read frame had RSV1 set
 	newDecompressionReader func(io.Reader) io.ReadCloser
+
+	headers *fasthttp.RequestHeader
 }
 
 func newConn(conn net.Conn, isServer bool, readBufferSize, writeBufferSize int) *Conn {
@@ -1147,3 +1151,18 @@ func FormatCloseMessage(closeCode int, text string) []byte {
 	copy(buf[2:], text)
 	return buf
 }
+
+// SetHeaders sets request headers
+func (c *Conn) SetHeaders(h *fasthttp.RequestHeader) {
+	c.headers = h
+}
+
+// Header returns header by key
+func (c *Conn) Header(key string) (value string) {
+	return string(c.headers.Peek(key))
+}
+
+// Headers returns the RequestHeader struct
+func (c *Conn) Headers() *fasthttp.RequestHeader {
+	return c.headers
+}
diff --git a/websocket/server.go b/websocket/server.go
index 2f29fb8..00e0fd6 100644
--- a/websocket/server.go
+++ b/websocket/server.go
@@ -191,7 +191,7 @@ func (u *Upgrader) Upgrade(ctx *fasthttp.RequestCtx, responseHeader *fasthttp.Re
 
 	ctx.Hijack(func(netConn net.Conn) {
 		c := newConn(netConn, true, u.ReadBufferSize, u.WriteBufferSize)
-		//c.SetHeaders(h)
+		c.SetHeaders(h)
 		c.subprotocol = subprotocol
 		if compress {
 			c.newCompressionWriter = compressNoContextTakeover