ing

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
+}

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