rfc7692 context-takeover implemention

client.go

@@ -6,6 +6,7 @@ package websocket
 
 import (
 	"bytes"
+	"compress/flate"
 	"context"
 	"crypto/tls"
 	"errors"
@@ -100,16 +101,19 @@ type Dialer struct {
 	// Subprotocols specifies the client's requested subprotocols.
 	Subprotocols []string
 
-	// EnableCompression specifies if the client should attempt to negotiate
-	// per message compression (RFC 7692). Setting this value to true does not
-	// guarantee that compression will be supported. Currently only "no context
-	// takeover" modes are supported.
+	// EnableCompression specify if the server should attempt to negotiate per
+	// message compression (RFC 7692).
 	EnableCompression bool
 
 	// Jar specifies the cookie jar.
 	// If Jar is nil, cookies are not sent in requests and ignored
 	// in responses.
 	Jar http.CookieJar
+
+	// AllowClientContextTakeover specifies whether the server will negotiate client context
+	// takeover for per message compression. Context takeover improves compression at the
+	// the cost of using more memory.
+	AllowClientContextTakeover bool
 }
 
 // Dial creates a new client connection by calling DialContext with a background context.
@@ -235,8 +239,11 @@ func (d *Dialer) DialContext(ctx context.Context, urlStr string, requestHeader h
 		}
 	}
 
-	if d.EnableCompression {
-		req.Header["Sec-WebSocket-Extensions"] = []string{"permessage-deflate; server_no_context_takeover; client_no_context_takeover"}
+	switch {
+	case d.EnableCompression && d.AllowClientContextTakeover:
+		req.Header.Set("Sec-Websocket-Extensions", "permessage-deflate; server_max_window_bits=15; client_max_window_bits=15")
+	case d.EnableCompression:
+		req.Header.Set("Sec-Websocket-Extensions", "permessage-deflate; server_no_context_takeover; client_no_context_takeover")
 	}
 
 	if d.HandshakeTimeout != 0 {
@@ -408,13 +415,24 @@ func (d *Dialer) DialContext(ctx context.Context, urlStr string, requestHeader h
 		if ext[""] != "permessage-deflate" {
 			continue
 		}
-		_, snct := ext["server_no_context_takeover"]
-		_, cnct := ext["client_no_context_takeover"]
-		if !snct || !cnct {
-			return nil, resp, errInvalidCompression
+
+		_, cmwb := ext["client_max_window_bits"]
+		_, smwb := ext["server_max_window_bits"]
+
+		switch {
+		case cmwb && smwb:
+			var wf contextTakeoverWriterFactory
+			conn.newCompressionWriter = wf.newCompressionWriter
+
+			var rf contextTakeoverReaderFactory
+			fr := flate.NewReader(nil)
+			rf.fr = fr
+			conn.newDecompressionReader = rf.newDeCompressionReader
+		default:
+			conn.newCompressionWriter = compressNoContextTakeover
+			conn.newDecompressionReader = decompressNoContextTakeover
 		}
-		conn.newCompressionWriter = compressNoContextTakeover
-		conn.newDecompressionReader = decompressNoContextTakeover
+
 		break
 	}
 

client_server_test.go

@@ -45,7 +45,14 @@ var cstDialer = Dialer{
 	HandshakeTimeout: 30 * time.Second,
 }
 
-type cstHandler struct{ *testing.T }
+type cstHandlerConfig struct {
+	contextTakeover bool
+}
+
+type cstHandler struct {
+	*testing.T
+	cstHandlerConfig
+}
 
 type cstServer struct {
 	*httptest.Server
@@ -59,17 +66,17 @@ const (
 	cstRequestURI = cstPath + "?" + cstRawQuery
 )
 
-func newServer(t *testing.T) *cstServer {
+func newServer(t *testing.T, c cstHandlerConfig) *cstServer {
 	var s cstServer
-	s.Server = httptest.NewServer(cstHandler{t})
+	s.Server = httptest.NewServer(cstHandler{t, c})
 	s.Server.URL += cstRequestURI
 	s.URL = makeWsProto(s.Server.URL)
 	return &s
 }
 
-func newTLSServer(t *testing.T) *cstServer {
+func newTLSServer(t *testing.T, c cstHandlerConfig) *cstServer {
 	var s cstServer
-	s.Server = httptest.NewTLSServer(cstHandler{t})
+	s.Server = httptest.NewTLSServer(cstHandler{t, c})
 	s.Server.URL += cstRequestURI
 	s.URL = makeWsProto(s.Server.URL)
 	return &s
@@ -92,6 +99,9 @@ func (t cstHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, "bad protocol", http.StatusBadRequest)
 		return
 	}
+	if t.contextTakeover {
+		cstUpgrader.AllowServerContextTakeover = true
+	}
 	ws, err := cstUpgrader.Upgrade(w, r, http.Header{"Set-Cookie": {"sessionID=1234"}})
 	if err != nil {
 		t.Logf("Upgrade: %v", err)
@@ -122,6 +132,28 @@ func (t cstHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		t.Logf("Close: %v", err)
 		return
 	}
+
+	// for multipleSendRecv when context takeover.
+	if t.contextTakeover {
+		op, rd, err := ws.NextReader()
+		if err != nil {
+			t.Logf("NextReader: %v", err)
+			return
+		}
+		wr, err := ws.NextWriter(op)
+		if err != nil {
+			t.Logf("NextWriter: %v", err)
+			return
+		}
+		if _, err = io.Copy(wr, rd); err != nil {
+			t.Logf("NextWriter: %v", err)
+			return
+		}
+		if err := wr.Close(); err != nil {
+			t.Logf("Close: %v", err)
+			return
+		}
+	}
 }
 
 func makeWsProto(s string) string {
@@ -148,9 +180,30 @@ func sendRecv(t *testing.T, ws *Conn) {
 	}
 }
 
+func multipleSendRecv(t *testing.T, ws *Conn) {
+	for _, message := range []string{"Hello World", "Can you read message?"} {
+		if err := ws.SetWriteDeadline(time.Now().Add(time.Second)); err != nil {
+			t.Fatalf("SetWriteDeadline: %v", err)
+		}
+		if err := ws.WriteMessage(TextMessage, []byte(message)); err != nil {
+			t.Fatalf("WriteMessage: %v", err)
+		}
+		if err := ws.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
+			t.Fatalf("SetReadDeadline: %v", err)
+		}
+		_, p, err := ws.ReadMessage()
+		if err != nil {
+			t.Fatalf("ReadMessage: %v", err)
+		}
+		if string(p) != message {
+			t.Fatalf("message=%s, want %s", p, message)
+		}
+	}
+}
+
 func TestProxyDial(t *testing.T) {
 
-	s := newServer(t)
+	s := newServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	surl, _ := url.Parse(s.Server.URL)
@@ -187,7 +240,7 @@ func TestProxyDial(t *testing.T) {
 }
 
 func TestProxyAuthorizationDial(t *testing.T) {
-	s := newServer(t)
+	s := newServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	surl, _ := url.Parse(s.Server.URL)
@@ -227,7 +280,7 @@ func TestProxyAuthorizationDial(t *testing.T) {
 }
 
 func TestDial(t *testing.T) {
-	s := newServer(t)
+	s := newServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	ws, _, err := cstDialer.Dial(s.URL, nil)
@@ -239,7 +292,7 @@ func TestDial(t *testing.T) {
 }
 
 func TestDialCookieJar(t *testing.T) {
-	s := newServer(t)
+	s := newServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	jar, _ := cookiejar.New(nil)
@@ -301,7 +354,7 @@ func rootCAs(t *testing.T, s *httptest.Server) *x509.CertPool {
 }
 
 func TestDialTLS(t *testing.T) {
-	s := newTLSServer(t)
+	s := newTLSServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	d := cstDialer
@@ -315,7 +368,7 @@ func TestDialTLS(t *testing.T) {
 }
 
 func TestDialTimeout(t *testing.T) {
-	s := newServer(t)
+	s := newServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	d := cstDialer
@@ -371,7 +424,7 @@ func (c *requireDeadlineNetConn) LocalAddr() net.Addr  { return c.c.LocalAddr()
 func (c *requireDeadlineNetConn) RemoteAddr() net.Addr { return c.c.RemoteAddr() }
 
 func TestHandshakeTimeout(t *testing.T) {
-	s := newServer(t)
+	s := newServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	d := cstDialer
@@ -387,7 +440,7 @@ func TestHandshakeTimeout(t *testing.T) {
 }
 
 func TestHandshakeTimeoutInContext(t *testing.T) {
-	s := newServer(t)
+	s := newServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	d := cstDialer
@@ -408,7 +461,7 @@ func TestHandshakeTimeoutInContext(t *testing.T) {
 }
 
 func TestDialBadScheme(t *testing.T) {
-	s := newServer(t)
+	s := newServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	ws, _, err := cstDialer.Dial(s.Server.URL, nil)
@@ -419,7 +472,7 @@ func TestDialBadScheme(t *testing.T) {
 }
 
 func TestDialBadOrigin(t *testing.T) {
-	s := newServer(t)
+	s := newServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	ws, resp, err := cstDialer.Dial(s.URL, http.Header{"Origin": {"bad"}})
@@ -436,7 +489,7 @@ func TestDialBadOrigin(t *testing.T) {
 }
 
 func TestDialBadHeader(t *testing.T) {
-	s := newServer(t)
+	s := newServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	for _, k := range []string{"Upgrade",
@@ -500,7 +553,7 @@ func TestDialExtraTokensInRespHeaders(t *testing.T) {
 }
 
 func TestHandshake(t *testing.T) {
-	s := newServer(t)
+	s := newServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	ws, resp, err := cstDialer.Dial(s.URL, http.Header{"Origin": {s.URL}})
@@ -751,7 +804,7 @@ func TestHost(t *testing.T) {
 }
 
 func TestDialCompression(t *testing.T) {
-	s := newServer(t)
+	s := newServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	dialer := cstDialer
@@ -764,8 +817,24 @@ func TestDialCompression(t *testing.T) {
 	sendRecv(t, ws)
 }
 
+func TestDialCompressionOfContextTakeover(t *testing.T) {
+	s := newServer(t, cstHandlerConfig{true})
+	defer s.Close()
+
+	dialer := cstDialer
+	dialer.EnableCompression = true
+	dialer.AllowClientContextTakeover = true
+	ws, _, err := dialer.Dial(s.URL, nil)
+	if err != nil {
+		t.Fatalf("Dial: %v", err)
+	}
+	defer ws.Close()
+
+	multipleSendRecv(t, ws)
+}
+
 func TestSocksProxyDial(t *testing.T) {
-	s := newServer(t)
+	s := newServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	proxyListener, err := net.Listen("tcp", "127.0.0.1:0")
@@ -868,7 +937,7 @@ func TestTracingDialWithContext(t *testing.T) {
 	}
 	ctx := httptrace.WithClientTrace(context.Background(), trace)
 
-	s := newTLSServer(t)
+	s := newTLSServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	d := cstDialer
@@ -907,7 +976,7 @@ func TestEmptyTracingDialWithContext(t *testing.T) {
 	trace := &httptrace.ClientTrace{}
 	ctx := httptrace.WithClientTrace(context.Background(), trace)
 
-	s := newTLSServer(t)
+	s := newTLSServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	d := cstDialer