- support websocket decompress

lesismal · Jul 4, 2021 · 9e4b46f · 9e4b46f
1 parent c09bea2
commit 9e4b46f
Show file tree

Hide file tree

Showing 13 changed files with 336 additions and 40 deletions.
diff --git a/examples/websocket/client/client.go b/examples/websocket/client/client.go
@@ -21,7 +21,7 @@ func main() {
 
 	u := url.URL{Scheme: "ws", Host: *addr, Path: *path}
 	log.Printf("connecting to %s", u.String())
-
+	websocket.DefaultDialer.EnableCompression = true
 	c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
 	if err != nil {
 		log.Fatal("dial:", err)

diff --git a/examples/websocket/server/server.go b/examples/websocket/server/server.go
@@ -21,8 +21,7 @@ var (
 )
 
 func onWebsocket(w http.ResponseWriter, r *http.Request) {
-	isTLS := false
-	upgrader := websocket.NewUpgrader(isTLS)
+	upgrader := &websocket.Upgrader{EnableCompression: true}
 	conn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
 		panic(err)

diff --git a/examples/websocket/server_autobahn/server.go b/examples/websocket/server_autobahn/server.go
@@ -15,8 +15,7 @@ import (
 )
 
 func onWebsocketFrame(w http.ResponseWriter, r *http.Request) {
-	isTLS := true
-	upgrader := websocket.NewUpgrader(isTLS)
+	upgrader := &websocket.Upgrader{}
 	conn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
 		panic(err)
@@ -40,8 +39,7 @@ func onWebsocketFrame(w http.ResponseWriter, r *http.Request) {
 }
 
 func onWebsocketMessage(w http.ResponseWriter, r *http.Request) {
-	isTLS := true
-	upgrader := websocket.NewUpgrader(isTLS)
+	upgrader := &websocket.Upgrader{}
 	conn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
 		panic(err)

diff --git a/examples/websocket/test_websocket/websocket_test.go b/examples/websocket/test_websocket/websocket_test.go
@@ -23,8 +23,7 @@ var (
 )
 
 func onWebsocket(cancelFunc context.CancelFunc, maxCount int, w http.ResponseWriter, r *http.Request) {
-	isTLS := false
-	upgrader := websocket.NewUpgrader(isTLS)
+	upgrader := &websocket.Upgrader{}
 	conn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
 		panic(err)

diff --git a/examples/websocket_1m/server/server.go b/examples/websocket_1m/server/server.go
@@ -20,8 +20,7 @@ var (
 )
 
 func onWebsocket(w http.ResponseWriter, r *http.Request) {
-	isTLS := false
-	upgrader := websocket.NewUpgrader(isTLS)
+	upgrader := &websocket.Upgrader{}
 	conn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
 		panic(err)

diff --git a/examples/websocket_1m_tls/server/server.go b/examples/websocket_1m_tls/server/server.go
@@ -22,8 +22,7 @@ var (
 )
 
 func onWebsocket(w http.ResponseWriter, r *http.Request) {
-	isTLS := true
-	upgrader := websocket.NewUpgrader(isTLS)
+	upgrader := &websocket.Upgrader{}
 	conn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
 		panic(err)

diff --git a/examples/websocket_tls/server/server.go b/examples/websocket_tls/server/server.go
@@ -21,8 +21,7 @@ var (
 
 func onWebsocket(w http.ResponseWriter, r *http.Request) {
 	flag.Parse()
-	isTLS := true
-	upgrader := websocket.NewUpgrader(isTLS)
+	upgrader := &websocket.Upgrader{}
 	conn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
 		panic(err)

diff --git a/examples/websocket_tls/sticky/server/server.go b/examples/websocket_tls/sticky/server/server.go
@@ -20,8 +20,7 @@ var (
 )
 
 func onWebsocket(w http.ResponseWriter, r *http.Request) {
-	isTLS := true
-	upgrader := websocket.NewUpgrader(isTLS)
+	upgrader := &websocket.Upgrader{}
 	conn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
 		panic(err)

diff --git a/examples/websocket_tls/test_tls_websocket/websocket_test.go b/examples/websocket_tls/test_tls_websocket/websocket_test.go
@@ -77,8 +77,7 @@ BpA7MNLxiqss+rCbwf3NbWxEMiDQ2zRwVoafVFys7tjmv6t2Xck=
 `)
 
 func onWebsocket(cancelFunc context.CancelFunc, maxCount int, w http.ResponseWriter, r *http.Request) {
-	isTLS := false
-	upgrader := websocket.NewUpgrader(isTLS)
+	upgrader := &websocket.Upgrader{}
 	conn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
 		panic(err)

diff --git a/nbhttp/tests/poller_test.go b/nbhttp/tests/poller_test.go
@@ -27,8 +27,7 @@ var (
 )
 
 func onWebsocket(w http.ResponseWriter, r *http.Request) {
-	isTLS := true
-	upgrader := websocket.NewUpgrader(isTLS)
+	upgrader := &websocket.Upgrader{}
 	conn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
 		panic(err)

diff --git a/nbhttp/websocket/compression.go b/nbhttp/websocket/compression.go
@@ -0,0 +1,59 @@
+package websocket
+
+import (
+	"compress/flate"
+	"io"
+	"sync"
+)
+
+const (
+	minCompressionLevel     = -2
+	maxCompressionLevel     = flate.BestCompression
+	defaultCompressionLevel = 1
+)
+
+var (
+	flateReaderPool = sync.Pool{New: func() interface{} {
+		return flate.NewReader(nil)
+	}}
+)
+
+const flateReaderTail = "\x00\x00\xff\xff" + "\x01\x00\x00\xff\xff"
+
+func decompressReader(r io.Reader) io.ReadCloser {
+	fr, _ := flateReaderPool.Get().(io.ReadCloser)
+	fr.(flate.Resetter).Reset(r, nil)
+	return &flateReadWrapper{fr}
+}
+
+func isValidCompressionLevel(level int) bool {
+	return minCompressionLevel <= level && level <= maxCompressionLevel
+}
+
+type flateReadWrapper struct {
+	fr io.ReadCloser
+}
+
+func (r *flateReadWrapper) Read(p []byte) (int, error) {
+	if r.fr == nil {
+		return 0, io.ErrClosedPipe
+	}
+	n, err := r.fr.Read(p)
+	if err == io.EOF {
+		// Preemptively place the reader back in the pool. This helps with
+		// scenarios where the application does not call NextReader() soon after
+		// this final read.
+		r.Close()
+	}
+	return n, err
+}
+
+func (r *flateReadWrapper) Close() error {
+	if r.fr == nil {
+		return io.ErrClosedPipe
+	}
+	err := r.fr.Close()
+	flateReaderPool.Put(r.fr)
+	r.fr = nil
+	return err
+}