mod: client & server setting for context-takeover

client.go

@@ -42,6 +42,7 @@ func NewClient(netConn net.Conn, u *url.URL, requestHeader http.Header, readBufS
 		NetDial: func(net, addr string) (net.Conn, error) {
 			return netConn, nil
 		},
+		CompressionLevel: defaultCompressionLevel,
 	}
 	return d.Dial(u.String(), requestHeader)
 }
@@ -79,15 +80,18 @@ type Dialer struct {
 	// takeover" modes are supported.
 	EnableCompression bool
 
-	// EnableContextTakeover specifies specifies if the client should attempt to negotiate
-	// per message compression with context-takeover (RFC 7692).
-	// but window bits is allowed only 15, because go's flate library support 15 bits only.
-	EnableContextTakeover bool
-
 	// Jar specifies the cookie jar.
 	// If Jar is nil, cookies are not sent in requests and ignored
 	// in responses.
 	Jar http.CookieJar
+
+	// CompressionLeval is set for contextTakeoer.
+	CompressionLevel int
+
+	// EnableContextTakeover specifies specifies if the client should attempt to negotiate
+	// per message compression with context-takeover (RFC 7692).
+	// but window bits is allowed only 15, because go's flate library support 15 bits only.
+	EnableContextTakeover bool
 }
 
 var errMalformedURL = errors.New("malformed ws or wss URL")
@@ -325,7 +329,7 @@ func (d *Dialer) Dial(urlStr string, requestHeader http.Header) (*Conn, *http.Re
 			conn.contextTakeover = true
 
 			var f contextTakeoverWriterFactory
-			f.fw, _ = flate.NewWriter(&f.tw, 2) // level is specified in Dialer, Upgrader
+			f.fw, _ = flate.NewWriter(&f.tw, d.CompressionLevel) // level is specified in Dialer, Upgrader
 			conn.newCompressionWriter = f.newCompressionWriter
 
 			conn.newDecompressionReader = decompressContextTakeover

conn.go

@@ -246,7 +246,6 @@ type Conn struct {
 	enableWriteCompression bool
 	compressionLevel       int
 	newCompressionWriter   func(io.WriteCloser, int) io.WriteCloser
-	compressionWriters     map[int]*flateWriteWrapper
 
 	// Read fields
 	reader        io.ReadCloser // the current reader returned to the application
@@ -329,8 +328,6 @@ func newConnBRW(conn net.Conn, isServer bool, readBufferSize, writeBufferSize in
 		writeBuf = make([]byte, writeBufferSize+maxFrameHeaderSize)
 	}
 
-	cw := make(map[int]*flateWriteWrapper, 2)
-
 	c := &Conn{
 		isServer:               isServer,
 		br:                     br,
@@ -341,8 +338,6 @@ func newConnBRW(conn net.Conn, isServer bool, readBufferSize, writeBufferSize in
 		enableWriteCompression: true,
 		compressionLevel:       defaultCompressionLevel,
 
-		compressionWriters: cw,
-
 		rxDict: &[]byte{},
 	}
 
@@ -515,12 +510,6 @@ func (c *Conn) NextWriter(messageType int) (io.WriteCloser, error) {
 	c.writer = mw
 	if c.newCompressionWriter != nil && c.enableWriteCompression && isData(messageType) {
 		mw.compress = true
-		// For context-takeover
-		if c.contextTakeover {
-			c.writer = c.newCompressionWriter(c.writer, c.compressionLevel)
-			return c.writer, nil
-		}
-
 		c.writer = c.newCompressionWriter(c.writer, c.compressionLevel)
 	}
 	return c.writer, nil

server.go

@@ -54,6 +54,14 @@ type Upgrader struct {
 	// guarantee that compression will be supported. Currently only "no context
 	// takeover" modes are supported.
 	EnableCompression bool
+
+	// CompressionLeval is set for contextTakeoer.
+	CompressionLevel int
+
+	// EnableContextTakeover specifies specifies if the client should attempt to negotiate
+	// per message compression with context-takeover (RFC 7692).
+	// but window bits is allowed only 15, because go's flate library support 15 bits only.
+	EnableContextTakeover bool
 }
 
 func (u *Upgrader) returnError(w http.ResponseWriter, r *http.Request, status int, reason string) (*Conn, error) {
@@ -186,11 +194,11 @@ func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeade
 
 	if compress {
 		switch {
-		case contextTakeover:
+		case contextTakeover && u.EnableContextTakeover:
 			c.contextTakeover = contextTakeover
 
 			var f contextTakeoverWriterFactory
-			f.fw, _ = flate.NewWriter(&f.tw, 2) // level is specified in Dialer, Upgrader
+			f.fw, _ = flate.NewWriter(&f.tw, u.CompressionLevel) // level is specified in Dialer, Upgrader
 			c.newCompressionWriter = f.newCompressionWriter
 
 			c.newDecompressionReader = decompressContextTakeover
@@ -211,7 +219,7 @@ func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeade
 	}
 	if compress {
 		switch {
-		case contextTakeover:
+		case contextTakeover && u.EnableContextTakeover:
 			p = append(p, "Sec-Websocket-Extensions: permessage-deflate; server_max_window_bits=15; client_max_window_bits=15\r\n"...)
 		default:
 			p = append(p, "Sec-Websocket-Extensions: permessage-deflate; server_no_context_takeover; client_no_context_takeover\r\n"...)