rfc7692 context-takeover implemention

client.go

@@ -6,6 +6,7 @@ package websocket
 
 import (
 	"bytes"
+	"compress/flate"
 	"crypto/tls"
 	"errors"
 	"io"
@@ -72,16 +73,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
 }
 
 var errMalformedURL = errors.New("malformed ws or wss URL")
@@ -200,8 +204,11 @@ func (d *Dialer) Dial(urlStr string, requestHeader http.Header) (*Conn, *http.Re
 		}
 	}
 
-	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")
 	}
 
 	var deadline time.Time
@@ -311,13 +318,24 @@ func (d *Dialer) Dial(urlStr string, requestHeader http.Header) (*Conn, *http.Re
 		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

@@ -40,7 +40,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
@@ -53,17 +60,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
@@ -86,6 +93,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)
@@ -116,6 +126,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 {
@@ -142,9 +174,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)
@@ -181,7 +234,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)
@@ -221,7 +274,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)
@@ -233,7 +286,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)
@@ -281,7 +334,7 @@ func TestDialCookieJar(t *testing.T) {
 }
 
 func TestDialTLS(t *testing.T) {
-	s := newTLSServer(t)
+	s := newTLSServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	certs := x509.NewCertPool()
@@ -307,7 +360,7 @@ func TestDialTLS(t *testing.T) {
 
 func xTestDialTLSBadCert(t *testing.T) {
 	// This test is deactivated because of noisy logging from the net/http package.
-	s := newTLSServer(t)
+	s := newTLSServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	ws, _, err := cstDialer.Dial(s.URL, nil)
@@ -318,7 +371,7 @@ func xTestDialTLSBadCert(t *testing.T) {
 }
 
 func TestDialTLSNoVerify(t *testing.T) {
-	s := newTLSServer(t)
+	s := newTLSServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	d := cstDialer
@@ -332,7 +385,7 @@ func TestDialTLSNoVerify(t *testing.T) {
 }
 
 func TestDialTimeout(t *testing.T) {
-	s := newServer(t)
+	s := newServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	d := cstDialer
@@ -345,7 +398,7 @@ func TestDialTimeout(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)
@@ -356,7 +409,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"}})
@@ -373,7 +426,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",
@@ -420,7 +473,7 @@ func TestBadMethod(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}})
@@ -482,7 +535,7 @@ func TestRespOnBadHandshake(t *testing.T) {
 // TestHostHeader confirms that the host header provided in the call to Dial is
 // sent to the server.
 func TestHostHeader(t *testing.T) {
-	s := newServer(t)
+	s := newServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	specifiedHost := make(chan string, 1)
@@ -509,7 +562,7 @@ func TestHostHeader(t *testing.T) {
 }
 
 func TestDialCompression(t *testing.T) {
-	s := newServer(t)
+	s := newServer(t, cstHandlerConfig{})
 	defer s.Close()
 
 	dialer := cstDialer
@@ -522,8 +575,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")