Merge pull request #7 from go-netty/feat-proxy

feat: Add WithDialer to support client proxy

Author: limpo1989

README.md

@@ -31,6 +31,7 @@ type Option
     func WithBufferSize(readBufferSize, writeBufferSize int) Option
     func WithCompress(compressLevel int, compressThreshold int64) Option
     func WithClientHeader(header http.Header) Option
+    func WithDialer(dialer Dialer) Option
     func WithMaxFrameSize(maxFrameSize int64) Option
     func WithNoDelay(noDelay bool) Option
     func WithServerHeader(header http.Header) Option
@@ -125,13 +126,10 @@ ws.OnClose = func(conn nettyws.Conn, err error) {
 fmt.Println("upgrade websocket connections ....")
 
 // upgrade websocket connection from http server
-serveMux := http.NewServeMux()
-serveMux.HandleFunc("/ws", func(writer http.ResponseWriter, request *http.Request) {
-    ws.UpgradeHTTP(writer, request)
-})
+http.Handle("/ws", ws)
 
 // listen http server
-if err := http.ListenAndServe(":9527", serveMux); nil != err {
+if err := http.ListenAndServe(":9527", nil); nil != err {
     panic(err)
 }
 ```

example/http-server.go

@@ -47,13 +47,10 @@ func main() {
 	fmt.Println("upgrade websocket connections ....")
 
 	// upgrade websocket connection from http server
-	serveMux := http.NewServeMux()
-	serveMux.HandleFunc("/ws", func(writer http.ResponseWriter, request *http.Request) {
-		ws.UpgradeHTTP(writer, request)
-	})
+	http.Handle("/ws", ws)
 
 	// listen http server
-	if err := http.ListenAndServe(":9527", serveMux); nil != err {
+	if err := http.ListenAndServe(":9527", nil); nil != err {
 		panic(err)
 	}
 }

nettyws.go

@@ -102,6 +102,15 @@ func (ws *Websocket) Close() error {
 	return nil
 }
 
+func (ws *Websocket) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
+	select {
+	case <-ws.ctx.Done():
+		http.Error(writer, "http: server shutdown", http.StatusNotAcceptable)
+	default:
+		_, _ = ws.upgrader.Upgrade(writer, request)
+	}
+}
+
 // UpgradeHTTP upgrades http connection to the websocket connection
 func (ws *Websocket) UpgradeHTTP(writer http.ResponseWriter, request *http.Request) (conn Conn, err error) {
 

options.go

@@ -17,7 +17,9 @@
 package nettyws
 
 import (
+	"context"
 	"crypto/tls"
+	"net"
 	"net/http"
 
 	"github.com/go-netty/go-netty"
@@ -35,6 +37,17 @@ const (
 	MsgBinary
 )
 
+// Dialer is a means to establish a connection.
+type Dialer interface {
+	// Dial connects to the given address via the proxy.
+	Dial(network, addr string) (c net.Conn, err error)
+}
+
+// contextDialer dials using a context
+type contextDialer interface {
+	DialContext(ctx context.Context, network, address string) (net.Conn, error)
+}
+
 type options struct {
 	engine            netty.Bootstrap
 	serveMux          *http.ServeMux
@@ -50,6 +63,7 @@ type options struct {
 	compressThreshold int64
 	requestHeader     http.Header
 	responseHeader    http.Header
+	dialer            Dialer
 }
 
 func parseOptions(opt ...Option) *options {
@@ -78,6 +92,16 @@ func (wso *options) wsOptions() *websocket.Options {
 		dialer.Header = ws.HandshakeHeaderHTTP(wso.requestHeader)
 	}
 
+	if wso.dialer != nil {
+		ctxDialer, isCtxDialer := wso.dialer.(contextDialer)
+		dialer.NetDial = func(ctx context.Context, network, addr string) (net.Conn, error) {
+			if isCtxDialer {
+				return ctxDialer.DialContext(ctx, network, addr)
+			}
+			return wso.dialer.Dial(network, addr)
+		}
+	}
+
 	var upgrader = ws.DefaultHTTPUpgrader
 	if wso.responseHeader != nil {
 		upgrader.Header = wso.responseHeader
@@ -192,3 +216,10 @@ func WithServerHeader(header http.Header) Option {
 		options.responseHeader = header
 	}
 }
+
+// WithDialer specify the client to connect to the network via a dialer.
+func WithDialer(dialer Dialer) Option {
+	return func(options *options) {
+		options.dialer = dialer
+	}
+}