From d6f59f95b0247c7f36432421404d9fe5c248d3f9 Mon Sep 17 00:00:00 2001
From: PotatoCloud <60210021+PotatoCloud@users.noreply.github.com>
Date: Tue, 18 Jun 2024 20:00:37 +0800
Subject: [PATCH] Add Handshake.Header field

---
 dialer.go |  3 +++
 server.go | 11 +++++++++++
 2 files changed, 14 insertions(+)

diff --git a/dialer.go b/dialer.go
index e66678e..d4d80a5 100644
--- a/dialer.go
+++ b/dialer.go
@@ -31,6 +31,9 @@ type Handshake struct {
 
 	// Extensions is the list of negotiated extensions.
 	Extensions []httphead.Option
+
+	// Header all request headers obtained during handshake
+	Header http.Header
 }
 
 // Errors used by the websocket client.
diff --git a/server.go b/server.go
index 863bb22..ef3016a 100644
--- a/server.go
+++ b/server.go
@@ -241,6 +241,10 @@ func (u HTTPUpgrader) Upgrade(r *http.Request, w http.ResponseWriter) (conn net.
 	if h := u.Header; h != nil {
 		header[0] = HandshakeHeaderHTTP(h)
 	}
+
+	// set handshake header
+	hs.Header = r.Header
+
 	if err == nil {
 		httpWriteResponseUpgrade(rw.Writer, strToBytes(nonce), hs, header.WriteTo)
 		err = rw.Writer.Flush()
@@ -498,6 +502,10 @@ func (u Upgrader) Upgrade(conn io.ReadWriter) (hs Handshake, err error) {
 
 		nonce = make([]byte, nonceSize)
 	)
+
+	// init handshake headers
+	hs.Header = make(http.Header)
+
 	for err == nil {
 		line, e := readLine(br)
 		if e != nil {
@@ -514,6 +522,9 @@ func (u Upgrader) Upgrade(conn io.ReadWriter) (hs Handshake, err error) {
 			break
 		}
 
+		// copy and add header
+		hs.Header.Add(btsToString(bytes.Clone(k)), btsToString(bytes.Clone(v)))
+
 		switch btsToString(k) {
 		case headerHostCanonical:
 			headerSeen |= headerSeenHost