diff --git a/Package.swift b/Package.swift index f61589f4d..4b194d25b 100644 --- a/Package.swift +++ b/Package.swift @@ -128,6 +128,7 @@ extension Target.Dependency { name: "NIOTransportServices", package: "swift-nio-transport-services" ) + static let nioTestUtils: Self = .product(name: "NIOTestUtils", package: "swift-nio") static let logging: Self = .product(name: "Logging", package: "swift-log") static let protobuf: Self = .product(name: "SwiftProtobuf", package: "swift-protobuf") static let protobufPluginLibrary: Self = .product( @@ -312,6 +313,7 @@ extension Target { .nioCore, .nioHTTP2, .nioEmbedded, + .nioTestUtils, ] ) diff --git a/Sources/GRPCHTTP2Core/GRPCMessageDeframer.swift b/Sources/GRPCHTTP2Core/GRPCMessageDeframer.swift new file mode 100644 index 000000000..af8282bb4 --- /dev/null +++ b/Sources/GRPCHTTP2Core/GRPCMessageDeframer.swift @@ -0,0 +1,103 @@ +/* + * Copyright 2024, gRPC Authors All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import GRPCCore +import NIOCore + +/// A ``GRPCMessageDeframer`` helps with the deframing of gRPC data frames: +/// - It reads the frame's metadata to know whether the message payload is compressed or not, and its length +/// - It reads and decompresses the payload, if compressed +/// - It helps put together frames that have been split across multiple `ByteBuffers` by the underlying transport +struct GRPCMessageDeframer: NIOSingleStepByteToMessageDecoder { + /// Length of the gRPC message header (1 compression byte, 4 bytes for the length). + static let metadataLength = 5 + static let defaultMaximumPayloadSize = Int.max + + typealias InboundOut = [UInt8] + + private let decompressor: Zlib.Decompressor? + private let maximumPayloadSize: Int + + /// Create a new ``GRPCMessageDeframer``. + /// - Parameters: + /// - maximumPayloadSize: The maximum size a message payload can be. + /// - decompressor: A `Zlib.Decompressor` to use when decompressing compressed gRPC messages. + /// - Important: You must call `end()` on the `decompressor` when you're done using it, to clean + /// up any resources allocated by `Zlib`. + init( + maximumPayloadSize: Int = Self.defaultMaximumPayloadSize, + decompressor: Zlib.Decompressor? = nil + ) { + self.maximumPayloadSize = maximumPayloadSize + self.decompressor = decompressor + } + + mutating func decode(buffer: inout ByteBuffer) throws -> InboundOut? { + guard buffer.readableBytes >= Self.metadataLength else { + // If we cannot read enough bytes to cover the metadata's length, then we + // need to wait for more bytes to become available to us. + return nil + } + + // Store the current reader index in case we don't yet have enough + // bytes in the buffer to decode a full frame, and need to reset it. + // The force-unwraps for the compression flag and message length are safe, + // because we've checked just above that we've got at least enough bytes to + // read all of the metadata. + let originalReaderIndex = buffer.readerIndex + let isMessageCompressed = buffer.readInteger(as: UInt8.self)! == 1 + let messageLength = buffer.readInteger(as: UInt32.self)! + + if messageLength > self.maximumPayloadSize { + throw RPCError( + code: .resourceExhausted, + message: """ + Message has exceeded the configured maximum payload size \ + (max: \(self.maximumPayloadSize), actual: \(messageLength)) + """ + ) + } + + guard var message = buffer.readSlice(length: Int(messageLength)) else { + // `ByteBuffer/readSlice(length:)` returns nil when there are not enough + // bytes to read the requested length. This can happen if we don't yet have + // enough bytes buffered to read the full message payload. + // By reading the metadata though, we have already moved the reader index, + // so we must reset it to its previous, original position for now, + // and return. We'll try decoding again, once more bytes become available + // in our buffer. + buffer.moveReaderIndex(to: originalReaderIndex) + return nil + } + + if isMessageCompressed { + guard let decompressor = self.decompressor else { + // We cannot decompress the payload - throw an error. + throw RPCError( + code: .internalError, + message: "Received a compressed message payload, but no decompressor has been configured." + ) + } + return try decompressor.decompress(&message, limit: self.maximumPayloadSize) + } else { + return Array(buffer: message) + } + } + + mutating func decodeLast(buffer: inout ByteBuffer, seenEOF: Bool) throws -> InboundOut? { + try self.decode(buffer: &buffer) + } +} diff --git a/Tests/GRPCHTTP2CoreTests/GRPCMessageDeframerTests.swift b/Tests/GRPCHTTP2CoreTests/GRPCMessageDeframerTests.swift new file mode 100644 index 000000000..98d5fc046 --- /dev/null +++ b/Tests/GRPCHTTP2CoreTests/GRPCMessageDeframerTests.swift @@ -0,0 +1,219 @@ +/* + * Copyright 2024, gRPC Authors All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import GRPCCore +import NIOCore +import NIOTestUtils +import XCTest + +@testable import GRPCHTTP2Core + +final class GRPCMessageDeframerTests: XCTestCase { + func testReadMultipleMessagesWithoutCompression() throws { + let firstMessage = { + var buffer = ByteBuffer() + buffer.writeInteger(UInt8(0)) + buffer.writeInteger(UInt32(16)) + buffer.writeRepeatingByte(42, count: 16) + return buffer + }() + + let secondMessage = { + var buffer = ByteBuffer() + buffer.writeInteger(UInt8(0)) + buffer.writeInteger(UInt32(8)) + buffer.writeRepeatingByte(43, count: 8) + return buffer + }() + + try ByteToMessageDecoderVerifier.verifyDecoder( + inputOutputPairs: [ + (firstMessage, [Array(repeating: UInt8(42), count: 16)]), + (secondMessage, [Array(repeating: UInt8(43), count: 8)]), + ]) { + GRPCMessageDeframer() + } + } + + func testReadMessageOverSizeLimitWithoutCompression() throws { + let deframer = GRPCMessageDeframer(maximumPayloadSize: 100) + let processor = NIOSingleStepByteToMessageProcessor(deframer) + + var buffer = ByteBuffer() + buffer.writeInteger(UInt8(0)) + buffer.writeInteger(UInt32(101)) + buffer.writeRepeatingByte(42, count: 101) + + XCTAssertThrowsError( + ofType: RPCError.self, + try processor.process(buffer: buffer) { _ in + XCTFail("No message should be produced.") + } + ) { error in + XCTAssertEqual(error.code, .resourceExhausted) + XCTAssertEqual( + error.message, + "Message has exceeded the configured maximum payload size (max: 100, actual: 101)" + ) + } + } + + func testReadMessageOverSizeLimitButWithoutActualMessageBytes() throws { + let deframer = GRPCMessageDeframer(maximumPayloadSize: 100) + let processor = NIOSingleStepByteToMessageProcessor(deframer) + + var buffer = ByteBuffer() + buffer.writeInteger(UInt8(0)) + // Set the message length field to be over the maximum payload size, but + // don't write the actual message bytes. This is to ensure that the payload + // size limit is enforced _before_ the payload is actually read. + buffer.writeInteger(UInt32(101)) + + XCTAssertThrowsError( + ofType: RPCError.self, + try processor.process(buffer: buffer) { _ in + XCTFail("No message should be produced.") + } + ) { error in + XCTAssertEqual(error.code, .resourceExhausted) + XCTAssertEqual( + error.message, + "Message has exceeded the configured maximum payload size (max: 100, actual: 101)" + ) + } + } + + func testCompressedMessageWithoutConfiguringDecompressor() throws { + let deframer = GRPCMessageDeframer(maximumPayloadSize: 100) + let processor = NIOSingleStepByteToMessageProcessor(deframer) + + var buffer = ByteBuffer() + buffer.writeInteger(UInt8(1)) + buffer.writeInteger(UInt32(10)) + buffer.writeRepeatingByte(42, count: 10) + + XCTAssertThrowsError( + ofType: RPCError.self, + try processor.process(buffer: buffer) { _ in + XCTFail("No message should be produced.") + } + ) { error in + XCTAssertEqual(error.code, .internalError) + XCTAssertEqual( + error.message, + "Received a compressed message payload, but no decompressor has been configured." + ) + } + } + + private func testReadMultipleMessagesWithCompression(method: Zlib.Method) throws { + let decompressor = Zlib.Decompressor(method: method) + let compressor = Zlib.Compressor(method: method) + var framer = GRPCMessageFramer() + defer { + decompressor.end() + compressor.end() + } + + let firstMessage = try { + framer.append(Array(repeating: 42, count: 100)) + return try framer.next(compressor: compressor)! + }() + + let secondMessage = try { + framer.append(Array(repeating: 43, count: 110)) + return try framer.next(compressor: compressor)! + }() + + try ByteToMessageDecoderVerifier.verifyDecoder( + inputOutputPairs: [ + (firstMessage, [Array(repeating: 42, count: 100)]), + (secondMessage, [Array(repeating: 43, count: 110)]), + ]) { + GRPCMessageDeframer(maximumPayloadSize: 1000, decompressor: decompressor) + } + } + + func testReadMultipleMessagesWithDeflateCompression() throws { + try self.testReadMultipleMessagesWithCompression(method: .deflate) + } + + func testReadMultipleMessagesWithGZIPCompression() throws { + try self.testReadMultipleMessagesWithCompression(method: .gzip) + } + + func testReadCompressedMessageOverSizeLimitBeforeDecompressing() throws { + let deframer = GRPCMessageDeframer(maximumPayloadSize: 1) + let processor = NIOSingleStepByteToMessageProcessor(deframer) + let compressor = Zlib.Compressor(method: .gzip) + var framer = GRPCMessageFramer() + defer { + compressor.end() + } + + framer.append(Array(repeating: 42, count: 100)) + let framedMessage = try framer.next(compressor: compressor)! + + XCTAssertThrowsError( + ofType: RPCError.self, + try processor.process(buffer: framedMessage) { _ in + XCTFail("No message should be produced.") + } + ) { error in + XCTAssertEqual(error.code, .resourceExhausted) + XCTAssertEqual( + error.message, + """ + Message has exceeded the configured maximum payload size \ + (max: 1, actual: \(framedMessage.readableBytes - GRPCMessageDeframer.metadataLength)) + """ + ) + } + } + + private func testReadDecompressedMessageOverSizeLimit(method: Zlib.Method) throws { + let decompressor = Zlib.Decompressor(method: method) + let deframer = GRPCMessageDeframer(maximumPayloadSize: 100, decompressor: decompressor) + let processor = NIOSingleStepByteToMessageProcessor(deframer) + let compressor = Zlib.Compressor(method: method) + var framer = GRPCMessageFramer() + defer { + decompressor.end() + compressor.end() + } + + framer.append(Array(repeating: 42, count: 101)) + let framedMessage = try framer.next(compressor: compressor)! + + XCTAssertThrowsError( + ofType: RPCError.self, + try processor.process(buffer: framedMessage) { _ in + XCTFail("No message should be produced.") + } + ) { error in + XCTAssertEqual(error.code, .resourceExhausted) + XCTAssertEqual(error.message, "Message is too large to decompress.") + } + } + + func testReadDecompressedMessageOverSizeLimitWithDeflateCompression() throws { + try self.testReadDecompressedMessageOverSizeLimit(method: .deflate) + } + + func testReadDecompressedMessageOverSizeLimitWithGZIPCompression() throws { + try self.testReadDecompressedMessageOverSizeLimit(method: .gzip) + } +}