From 19ad33f2576d3a07ff28984d25aca4528b2c8910 Mon Sep 17 00:00:00 2001
From: limpo1989 <limpo1989@gmail.com>
Date: Thu, 30 Nov 2023 19:17:20 +0800
Subject: [PATCH] feat: Add additional headers to the handshake

---
 README.md              |  4 ++++
 example/client.go      |  9 +++++++--
 example/http-server.go |  2 +-
 example/server.go      |  9 +++++++--
 options.go             | 33 +++++++++++++++++++++++++++++++--
 5 files changed, 50 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index b0bf4ad..212a181 100644
--- a/README.md
+++ b/README.md
@@ -30,8 +30,10 @@ type Option
     func WithBinary() Option
     func WithBufferSize(readBufferSize, writeBufferSize int) Option
     func WithCompress(compressLevel int, compressThreshold int64) Option
+    func WithClientHeader(header http.Header) Option
     func WithMaxFrameSize(maxFrameSize int64) Option
     func WithNoDelay(noDelay bool) Option
+    func WithServerHeader(header http.Header) Option
     func WithServeMux(serveMux *http.ServeMux) Option
     func WithServeTLS(tls *tls.Config) Option
     func WithValidUTF8() Option
@@ -39,6 +41,8 @@ type Option
 
 ## Easy to use
 
+> Note: `nettyws` does not support mixed text messages and binary messages, use the `WithBinary` option to switch to binary message mode.
+
 ### server :
 ```go
 // create websocket instance
diff --git a/example/client.go b/example/client.go
index 4f268ff..8daf253 100644
--- a/example/client.go
+++ b/example/client.go
@@ -18,18 +18,23 @@ package main
 
 import (
 	"fmt"
+	"net/http"
+	"time"
 
 	nettyws "github.com/go-netty/go-netty-ws"
 )
 
 func main() {
 
+	header := http.Header{}
+	header.Add("client-time", time.Now().String())
+
 	// create websocket instance
-	var ws = nettyws.NewWebsocket()
+	var ws = nettyws.NewWebsocket(nettyws.WithClientHeader(header))
 
 	// setup OnOpen handler
 	ws.OnOpen = func(conn nettyws.Conn) {
-		fmt.Println("OnOpen: ", conn.RemoteAddr())
+		fmt.Println("OnOpen: ", conn.RemoteAddr(), ", header: ", conn.Header())
 		conn.Write([]byte("hello world"))
 	}
 
diff --git a/example/http-server.go b/example/http-server.go
index 7b83492..d8112ad 100644
--- a/example/http-server.go
+++ b/example/http-server.go
@@ -35,7 +35,7 @@ func main() {
 
 	// setup OnData handler
 	ws.OnData = func(conn nettyws.Conn, data []byte) {
-		fmt.Println("OnData: ", conn.RemoteAddr(), ", message: ", string(data))
+		fmt.Println("OnData: ", conn.RemoteAddr(), ", message: ", string(data), ", header: ", conn.Header())
 		conn.Write(data)
 	}
 
diff --git a/example/server.go b/example/server.go
index 895f3e4..214bf5e 100644
--- a/example/server.go
+++ b/example/server.go
@@ -18,17 +18,22 @@ package main
 
 import (
 	"fmt"
+	"net/http"
+	"time"
 
 	nettyws "github.com/go-netty/go-netty-ws"
 )
 
 func main() {
+	header := http.Header{}
+	header.Add("server-time", time.Now().String())
+
 	// create websocket instance
-	var ws = nettyws.NewWebsocket()
+	var ws = nettyws.NewWebsocket(nettyws.WithServerHeader(header))
 
 	// setup OnOpen handler
 	ws.OnOpen = func(conn nettyws.Conn) {
-		fmt.Println("OnOpen: ", conn.RemoteAddr())
+		fmt.Println("OnOpen: ", conn.RemoteAddr(), ", header: ", conn.Header())
 	}
 
 	// setup OnData handler
diff --git a/options.go b/options.go
index 29e0a45..1c4878d 100644
--- a/options.go
+++ b/options.go
@@ -48,6 +48,8 @@ type options struct {
 	compressEnabled   bool
 	compressLevel     int
 	compressThreshold int64
+	requestHeader     http.Header
+	responseHeader    http.Header
 }
 
 func parseOptions(opt ...Option) *options {
@@ -70,6 +72,17 @@ func (wso *options) wsOptions() *websocket.Options {
 	if MsgBinary == wso.messageType {
 		opCode = ws.OpBinary
 	}
+
+	var dialer = ws.DefaultDialer
+	if wso.requestHeader != nil {
+		dialer.Header = ws.HandshakeHeaderHTTP(wso.requestHeader)
+	}
+
+	var upgrader = ws.DefaultHTTPUpgrader
+	if wso.responseHeader != nil {
+		upgrader.Header = wso.responseHeader
+	}
+
 	return &websocket.Options{
 		TLS:               wso.tls,
 		OpCode:            opCode,
@@ -82,8 +95,8 @@ func (wso *options) wsOptions() *websocket.Options {
 		CompressEnabled:   wso.compressEnabled,
 		CompressLevel:     wso.compressLevel,
 		CompressThreshold: wso.compressThreshold,
-		Dialer:            ws.DefaultDialer,
-		Upgrader:          ws.DefaultHTTPUpgrader,
+		Dialer:            dialer,
+		Upgrader:          upgrader,
 		ServeMux:          wso.serveMux,
 	}
 }
@@ -163,3 +176,19 @@ func WithCompress(compressLevel int, compressThreshold int64) Option {
 		options.compressThreshold = compressThreshold
 	}
 }
+
+// WithClientHeader is an optional http.Header mapping that could be used to
+// write additional headers to the handshake request.
+func WithClientHeader(header http.Header) Option {
+	return func(options *options) {
+		options.requestHeader = header
+	}
+}
+
+// WithServerHeader is an optional http.Header mapping that could be used to
+// write additional headers to the handshake response.
+func WithServerHeader(header http.Header) Option {
+	return func(options *options) {
+		options.responseHeader = header
+	}
+}