Merge pull request #6 from go-netty/feat-handshake-headers

feat: Add additional headers to the handshake

Author: limpo1989

README.md

@@ -30,15 +30,19 @@ 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
 ```
 
 ## 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

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"))
 	}
 

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

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

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