permessage-deflate: implement context takeover (#26)

Pushkar N Kulkarni · web-flow · commit 8bde34ce7227 · 2019-07-04T11:34:30.000+05:30
* permessage-deflate: implement context takeover * Use negative windowBits values for the deflater The compression format of the output produced by zlibs deflater is decided by the windowBits values it is configured with. We used positive windowBits values which, according to the zlib manual, emits the zlib stream starting with a zlib header. WebSockets use raw deflate streams, and not zlib streams. Consequently, we had to strip the zlib header. The deflater can be configured to emit a raw deflate format by using negative values for the windowBits. With this in place, there's no zlib header in the output stream. * Handle the special case of LZ77 window size of 256 The server_max_window_bits value sent by the client in the negotiation offer is used to configure the deflater on the server. This acceptance is also notified to the client by returning the server_max_window_bits header in the negotiation response. A special case arises when we receive server_max_window_bits = 8. There's an open zlib issue with the window size of 256 (windoBits=8). For zlib streams, the library silently changes the windowBits to 9 and informs the inflater via the zlib header. However, this apparent hack is not feasible for raw deflate streams, which are used in WebSockets. To take care of this fact, zlib has been patched to reject a window size of 256, via a deflateInit2() failure. Details here: madler/zlib#171 As a result, we need to handle a window size of 256 (windowBits=8) as a special case. We silently change it to 9 and inform the client via a suitable header in the negotiation response. * Close deflater and inflater streams on connection close Without context takeover, the deflater and inflater are initialized and closed on every message. However, with context takeover, the deflater and inflater persist for the entire lifetime of the connection. They need to be closed on a connection close. The deflater is on a ChannelOutboundHandler and must be closed when the WebSocketConnections issues a channel.close(). The inflater is on a ChannelInboundHandler, it must be closed when we receive a `closeConnection` frame from the remote peer.
diff --git a/Sources/KituraWebSocket/PermessageDeflate.swift b/Sources/KituraWebSocket/PermessageDeflate.swift
@@ -26,35 +26,44 @@ class PermessageDeflate: WebSocketProtocolExtension {
         guard header.hasPrefix("permessage-deflate") else { return [] }
         var deflaterMaxWindowBits: Int32 = 15
         var inflaterMaxWindowBits: Int32 = 15
-        //TODO: change these defaults to false after implementing context takeover
-        var clientNoContextTakeover = true
-        var serverNoContextTakeover = true
+        var clientNoContextTakeover = false
+        var serverNoContextTakeover = false
 
         // Four parameters to handle:
         // * server_max_window_bits: the LZ77 sliding window size used by the server for compression
         // * client_max_window_bits: the LZ77 sliding window size used by the server for decompression
         // * server_no_context_takeover: prevent the server from using context-takeover
         // * client_no_context_takeover: prevent the client from using context-takeover
         for parameter in header.components(separatedBy: "; ") {
-            // If we receieved a valid value for server_max_window_bits, configure the deflater to use it
+            // If we receieved a valid value for server_max_window_bits, use it to configure the deflater
             if parameter.hasPrefix("server_max_window_bits") {
                 let maxWindowBits = parameter.components(separatedBy: "=")
                 guard maxWindowBits.count == 2 else { continue }
-                if let mwBits = Int32(maxWindowBits[1]) {
-                    if mwBits >= 8 && mwBits <= 15 {
-                        deflaterMaxWindowBits = mwBits
-                    }
+                guard let mwBits = Int32(maxWindowBits[1]) else { continue }
+                if mwBits >= 8 && mwBits <= 15 {
+                    // We received a valid value. However there's a special case here:
+                    //
+                    // There's an open zlib issue which does not set the window size
+                    // to 256 (windowBits=8). For windowBits=8, zlib silently changes the
+                    // value to 9. However, this apparent hack works only with zlib streams.
+                    // WebSockets use raw deflate streams. For raw deflate streams, zlib has been
+                    // patched to ignore the windowBits value 8.
+                    // More details here: https://github.com/madler/zlib/issues/171
+                    //
+                    // So, if the server requested for server_max_window_bits=8, we are
+                    // going to use server_max_window_bits=9 instead and notify this in
+                    // our negotiation response too.
+                    deflaterMaxWindowBits = mwBits == 8 ? 9 : mwBits
                 }
             }
 
-            // If we receieved a valid value for server_max_window_bits, configure the inflater to use it
+            // If we received a valid client_max_window_bits value, use it to configure the inflater
             if parameter.hasPrefix("client_max_window_bits") {
                 let maxWindowBits = parameter.components(separatedBy: "=")
                 guard maxWindowBits.count == 2 else { continue }
-                if let mwBits = Int32(maxWindowBits[1]) {
-                    if mwBits >= 8 && mwBits <= 15  {
-                        inflaterMaxWindowBits = mwBits
-                    }
+                guard let mwBits = Int32(maxWindowBits[1]) else { continue }
+                if mwBits >= 8 && mwBits <= 15  {
+                    inflaterMaxWindowBits = mwBits
                 }
             }
 
@@ -66,7 +75,6 @@ class PermessageDeflate: WebSocketProtocolExtension {
                 serverNoContextTakeover = true
             }
         }
-
         return [PermessageDeflateCompressor(maxWindowBits: deflaterMaxWindowBits, noContextTakeOver: serverNoContextTakeover),
                    PermessageDeflateDecompressor(maxWindowBits: inflaterMaxWindowBits, noContextTakeOver: clientNoContextTakeover)]
     }
@@ -81,16 +89,44 @@ class PermessageDeflate: WebSocketProtocolExtension {
 
         for parameter in header.components(separatedBy: "; ") {
             if parameter == "client_no_context_takeover" {
-                //TODO: include client_no_context_takeover in the response
+                response.append("; client_no_context_takeover")
             }
 
             if parameter == "server_no_context_takeover" {
-                //TODO: include server_no_context_takeover in the response
+                response.append("; server_no_context_takeover")
+            }
+
+            // If we receive a valid value for server_max_window_bits, we accept it and return if
+            // in the response. If we receive an invalid value, we default to 15 and return the
+            // same in the response. If we receive no value, we ignore this header.
+            if parameter.hasPrefix("server_max_window_bits") {
+                let maxWindowBits = parameter.components(separatedBy: "=")
+                guard maxWindowBits.count == 2 else { continue }
+                guard let mwBits = Int32(maxWindowBits[1]) else { continue }
+                if mwBits >= 8 && mwBits <= 15 {
+                    // We received a valid value. However there's a special case here:
+                    //
+                    // There's an open zlib issue which does not set the window size
+                    // to 256 (windowBits=8). For windowBits=8, zlib silently changes the
+                    // value to 9. However, this apparent hack works only with zlib streams.
+                    // WebSockets use raw deflate streams. For raw deflate streams, zlib has been
+                    // patched to ignore the windowBits value 8.
+                    // More details here: https://github.com/madler/zlib/issues/171
+                    //
+                    // So, if the server requested for server_max_window_bits=8, we are
+                    // going to use server_max_window_bits=9 instead and notify this in
+                    // our negotiation response too.
+                    if mwBits == 8 {
+                        response.append("; server_max_window_bits=9")
+                    } else {
+                        response.append("; \(parameter)")
+                    }
+                } else {
+                    // we received an invalid value
+                    response.append("; server_max_window_bits=15")
+                }
             }
         }
-        //TODO: remove this after we have implemented context takeover
-        response.append("; server_no_context_takeover")
-        response.append("; client_no_context_takeover")
         return response
     }
 }
diff --git a/Sources/KituraWebSocket/PermessageDeflateCompressor.swift b/Sources/KituraWebSocket/PermessageDeflateCompressor.swift
@@ -44,6 +44,9 @@ class PermessageDeflateCompressor : ChannelOutboundHandler {
     // The zlib stream
     private var stream: z_stream = z_stream()
 
+    // Initialize the z_stream only once if context takeover is enabled
+    private var streamInitialized = false
+
     // PermessageDeflateCompressor is an outbound handler, this function gets called when a frame is written to the channel by WebSocketConnection.
     func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
         var frame = unwrapOutboundIn(data)
@@ -77,19 +80,38 @@ class PermessageDeflateCompressor : ChannelOutboundHandler {
         _ = context.writeAndFlush(self.wrapOutboundOut(deflatedFrame))
     }
 
+    func close(context: ChannelHandlerContext, mode: CloseMode, promise: EventLoopPromise<Void>?) {
+        // PermessageDeflateCompressor is an outbound handler. If the underlying
+        // WebSocketConnection decides to close the connection, the close message
+        // needs to be intercepted and the deflater closed while we're using context takeover.
+        if noContextTakeOver == false {
+            deflateEnd(&stream)
+        }
+        context.close(mode: mode, promise: promise)
+    }
+
     func deflatePayload(in buffer: ByteBuffer, allocator: ByteBufferAllocator, dropFourTrailingOctets: Bool = false) -> ByteBuffer {
         // Initialize the deflater as per https://www.zlib.net/zlib_how.html
-        stream.zalloc = nil 
-        stream.zfree = nil 
-        stream.opaque = nil 
-
-        let rc = deflateInit2_(&stream, Z_DEFAULT_COMPRESSION, Z_DEFLATED, self.maxWindowBits, 8,
-                     Z_DEFAULT_STRATEGY, ZLIB_VERSION, Int32(MemoryLayout<z_stream>.size))
-        precondition(rc == Z_OK, "Unexpected return from zlib init: \(rc)")
+        if noContextTakeOver || streamInitialized == false {
+            stream.zalloc = nil
+            stream.zfree = nil
+            stream.opaque = nil
+            stream.next_in = nil
+            stream.avail_in = 0
+            // The zlib manual asks us to provide a negative windowBits value for raw deflate
+            let rc = deflateInit2_(&stream, Z_DEFAULT_COMPRESSION, Z_DEFLATED, -self.maxWindowBits, 8,
+                                   Z_DEFAULT_STRATEGY, ZLIB_VERSION, Int32(MemoryLayout<z_stream>.size))
+            precondition(rc == Z_OK, "Unexpected return from zlib init: \(rc)")
+            self.streamInitialized = true
+        }
 
         defer {
-            // Deinitialize the deflater before returning
-            deflateEnd(&stream)
+            if noContextTakeOver {
+                // We aren't doing a context takeover.
+                // This means the deflater is to be used on a per-message basis.
+                // So, we deinitialize the deflater before returning.
+                deflateEnd(&stream)
+            }
         }
 
         // Deflate/compress the payload
@@ -114,10 +136,6 @@ class PermessageDeflateCompressor : ChannelOutboundHandler {
         precondition(inputBuffer.readableBytes == 0)
         precondition(outputBuffer.readableBytes > 0)
 
-        // Remove the 0x78 0x9C zlib header added by zlib
-        _ = outputBuffer.readBytes(length: 2)
-        outputBuffer.discardReadBytes()
-
         // Ignore the 0, 0, 0xff, 0xff trailer added by zlib
         if dropFourTrailingOctets {
             outputBuffer = outputBuffer.getSlice(at: 0, length: outputBuffer.readableBytes-4) ?? outputBuffer
diff --git a/Sources/KituraWebSocket/PermessageDeflateDecompressor.swift b/Sources/KituraWebSocket/PermessageDeflateDecompressor.swift
@@ -36,6 +36,8 @@ class PermessageDeflateDecompressor : ChannelInboundHandler {
     // The zlib stream
     private var stream: z_stream = z_stream()
 
+    private var streamInitialized = false
+
     // A buffer to accumulate payload across multiple frames
     var payload: ByteBuffer?
 
@@ -56,6 +58,12 @@ class PermessageDeflateDecompressor : ChannelInboundHandler {
         let frame = unwrapInboundIn(data)
         // We should either have a data frame with rsv1 set, or a continuation frame of a compressed message. There's nothing to do otherwise.
         guard frame.isCompressedDataFrame || (frame.isContinuationFrame && self.receivingCompressedMessage) else {
+            // If we are using context takeover, this is a good time to free the zstream!
+            if streamInitialized && frame.opcode == .connectionClose && !noContextTakeOver {
+                deflateEnd(&stream)
+                streamInitialized = false
+            }
+
             context.fireChannelRead(self.wrapInboundOut(frame))
             return
         }
@@ -90,17 +98,22 @@ class PermessageDeflateDecompressor : ChannelInboundHandler {
 
     func inflatePayload(in buffer: ByteBuffer, allocator: ByteBufferAllocator) -> ByteBuffer {
         // Initialize the inflater as per https://www.zlib.net/zlib_how.html
-        stream.zalloc = nil 
-        stream.zfree = nil 
-        stream.opaque = nil 
-        stream.avail_in = 0 
-        stream.next_in = nil 
-        let rc = inflateInit2_(&stream, -self.maxWindowBits, ZLIB_VERSION, Int32(MemoryLayout<z_stream>.size))
-        precondition(rc == Z_OK, "Unexpected return from zlib init: \(rc)")
+        if noContextTakeOver || streamInitialized == false {
+            stream.zalloc = nil
+            stream.zfree = nil
+            stream.opaque = nil
+            stream.avail_in = 0
+            stream.next_in = nil
+            let rc = inflateInit2_(&stream, -self.maxWindowBits, ZLIB_VERSION, Int32(MemoryLayout<z_stream>.size))
+            precondition(rc == Z_OK, "Unexpected return from zlib init: \(rc)")
+            self.streamInitialized = true
+        }
 
         defer {
-            // Deinitialize before returning
-            inflateEnd(&stream)
+            if noContextTakeOver {
+                // Deinitialize before returning
+                inflateEnd(&stream)
+            }
         }
 
         // Inflate/decompress the payload
diff --git a/Sources/KituraWebSocket/WebSocketConnection.swift b/Sources/KituraWebSocket/WebSocketConnection.swift
@@ -335,10 +335,10 @@ extension WebSocketConnection {
          let frame = WebSocketFrame(fin: true, opcode: .connectionClose, data: data)
          let promise = context.eventLoop.makePromise(of: Void.self)
          context.writeAndFlush(self.wrapOutboundOut(frame), promise: promise)
-         promise.futureResult.whenComplete { _ in
-             if hard {
-                 _ = context.close(mode: .output)
-             }
+         if hard {
+            promise.futureResult.flatMap { _ in
+                context.close(mode: .output)
+            }.whenComplete { _ in }
          }
          awaitClose = true
     }
diff --git a/Tests/KituraWebSocketTests/ComplexTests.swift b/Tests/KituraWebSocketTests/ComplexTests.swift
@@ -29,7 +29,11 @@ class ComplexTests: KituraTest {
             ("testPingBetweenBinaryFrames", testPingBetweenBinaryFrames),
             ("testPingBetweenTextFrames", testPingBetweenTextFrames),
             ("testTextShortAndMediumFrames", testTextShortAndMediumFrames),
-            ("testTextTwoShortFrames", testTextTwoShortFrames)
+            ("testTextTwoShortFrames", testTextTwoShortFrames),
+            ("testTwoMessagesWithContextTakeover", testTwoMessagesWithContextTakeover),
+            ("testTwoMessagesWithClientContextTakeover", testTwoMessagesWithClientContextTakeover),
+            ("testTwoMessagesWithServerContextTakeover", testTwoMessagesWithServerContextTakeover),
+            ("testTwoMessagesWithNoContextTakeover", testTwoMessagesWithNoContextTakeover),
         ]
     }
 
@@ -58,7 +62,7 @@ class ComplexTests: KituraTest {
                 self.performTest(framesToSend: [(false, self.opcodeBinary, shortBinaryPayload), (true, self.opcodeContinuation, mediumBinaryPayload)],
                                  expectedFrames: [(true, self.opcodeBinary, expectedBinaryPayload)],
                                  expectation: expectation, negotiateCompression: true, compressed: true)
-            },  { expectation in
+            }, { expectation in
                 self.performTest(framesToSend: [(false, self.opcodeBinary, shortBinaryPayload), (true, self.opcodeContinuation, mediumBinaryPayload)],
                                  expectedFrames: [(true, self.opcodeBinary, expectedBinaryPayload)],
                                  expectation: expectation, negotiateCompression: true, compressed: false)
@@ -183,4 +187,33 @@ class ComplexTests: KituraTest {
                                  expectation: expectation, negotiateCompression: true, compressed: false)
         })
     }
+
+    func testTwoMessages(contextTakeover: ContextTakeover = .both) {
+        register(closeReason: .noReasonCodeSent)
+
+        let text = "RFC7692 specifies a framework for adding compression functionality to the WebSocket Protocol"
+        let textPayload = self.payload(text: text)
+
+        performServerTest(asyncTasks: { expectation in
+            self.performTest(framesToSend: [(true, self.opcodeText, textPayload), (true, self.opcodeText, textPayload)],
+                             expectedFrames: [(true, self.opcodeText, textPayload), (true, self.opcodeText, textPayload)],
+                             expectation: expectation, negotiateCompression: true, compressed: true, contextTakeover: contextTakeover)
+        })
+    }
+
+    func testTwoMessagesWithContextTakeover() {
+        testTwoMessages(contextTakeover: .both)
+    }
+
+    func testTwoMessagesWithClientContextTakeover() {
+        testTwoMessages(contextTakeover: .client)
+    }
+
+    func testTwoMessagesWithServerContextTakeover() {
+        testTwoMessages(contextTakeover: .server)
+    }
+
+    func testTwoMessagesWithNoContextTakeover() {
+        testTwoMessages(contextTakeover: .none)
+    }
 }
diff --git a/Tests/KituraWebSocketTests/KituraTest+Frames.swift b/Tests/KituraWebSocketTests/KituraTest+Frames.swift
@@ -65,7 +65,7 @@ extension KituraTest {
     }
 
     //Sometimes, we may have a non-final frame as the last frame
-    func sendFrame(final: Bool, withOpcode: Int, withMasking: Bool=true, withPayload: NSData, on channel: Channel, lastFrame: Bool = false, compressed: Bool = false) {
+    func sendFrame(final: Bool, withOpcode: Int, withMasking: Bool=true, withPayload: NSData, on channel: Channel, lastFrame: Bool = false, compressed: Bool = false, contextTakeover: ContextTakeover? = nil) {
         var buffer = channel.allocator.buffer(capacity: 8)
         var payloadLength = withPayload.length
 
@@ -76,7 +76,7 @@ extension KituraTest {
         }
 
         if compressed {
-            payloadBuffer = PermessageDeflateCompressor().deflatePayload(in: payloadBuffer, allocator: ByteBufferAllocator(), dropFourTrailingOctets: final)
+            payloadBuffer = self.compressor.deflatePayload(in: payloadBuffer, allocator: ByteBufferAllocator(), dropFourTrailingOctets: final)
             payloadLength = payloadBuffer.readableBytes
         }
 
@@ -192,11 +192,14 @@ class WebSocketClientHandler: ChannelInboundHandler {
 
     var compressed: Bool = false
 
-    init(expectedFrames: [(Bool, Int, NSData)], expectation: XCTestExpectation, compressed: Bool = false) {
+    var decompressor: PermessageDeflateDecompressor
+
+    init(expectedFrames: [(Bool, Int, NSData)], expectation: XCTestExpectation, compressed: Bool = false, decompressor: PermessageDeflateDecompressor) {
         self.numberOfFramesExpected = expectedFrames.count
         self.expectedFrames = expectedFrames
         self.expectation = expectation
         self.compressed = compressed
+        self.decompressor = decompressor
     }
 
     func channelRead(context: ChannelHandlerContext, data: NIOAny) {
@@ -225,7 +228,7 @@ class WebSocketClientHandler: ChannelInboundHandler {
                 currentFramePayload += [0, 0, 0xff, 0xff]
                 var payloadBuffer = ByteBufferAllocator().buffer(capacity: 8)
                 payloadBuffer.writeBytes(currentFramePayload)
-                let inflatedBuffer = PermessageDeflateDecompressor().inflatePayload(in: payloadBuffer, allocator: ByteBufferAllocator())
+                let inflatedBuffer = self.decompressor.inflatePayload(in: payloadBuffer, allocator: ByteBufferAllocator())
                 currentFramePayload = inflatedBuffer.getBytes(at: 0, length: inflatedBuffer.readableBytes) ?? []
             }
             let currentFramePayloadPtr = UnsafeBufferPointer(start: &currentFramePayload, count: currentFramePayload.count)
diff --git a/Tests/KituraWebSocketTests/KituraTest.swift b/Tests/KituraWebSocketTests/KituraTest.swift
