Making BHTTPSerializer compliant to known-length messages

README.md

@@ -47,8 +47,9 @@ dependencies: [
 ### Binary HTTP Encoding
 
 To serialise binary HTTP messages use `BHTTPSerializer.serialize(message, buffer)`.
+As defined in [RFC9292](https://www.rfc-editor.org/rfc/rfc9292),  you can choose to use either the known-length format or an indeterminate-length format. This choice can be configured during the initialization of the serializer by passing the desired type: `BHTTPSerializer(type: .knownLength)`.
 
-To deserialise binary HTTP messages use `BHTTPParser`, adding received data with `append()`, then calling `completeBodyRecieved()`.  The read the message parts received call `nextMessage()`.
+To deserialise binary HTTP messages use `BHTTPParser`, adding received data with `append()`, then calling `completeBodyReceived()`.  The read the message parts received call `nextMessage()`.
 
 ### Oblivious Encapsulation
 

Sources/ObliviousHTTP/BHTTPSerializer.swift

@@ -14,92 +14,310 @@
 import NIOCore
 import NIOHTTP1
 
-// For now this type is entirely stateless, which is achieved by using the indefinite-length encoding.
-// It also means it does not enforce correctness, and so can produce invalid encodings if a user holds
-// it wrong.
-//
-// Later optimizations can be made by adding more state into this type.
 /// Binary HTTP serialiser as described in [RFC9292](https://www.rfc-editor.org/rfc/rfc9292).
-/// Currently only indeterminate-length encoding is supported.
 public struct BHTTPSerializer {
+
+    private var fsm: BHTTPSerializerFSM
+    private var context: SerializerContext
+
     /// Initialise a Binary HTTP Serialiser.
-    public init() {}
+    /// - Parameters:
+    ///   - type: The type of BHTTPSerializer you want: either known or indeterminate length.
+    ///   - allocator: Byte buffer allocator used.
+    public init(
+        type: SerializerType = .indeterminateLength,
+        allocator: ByteBufferAllocator = ByteBufferAllocator()
+    ) {
+        self.fsm = BHTTPSerializerFSM(initialState: BHTTPSerializerState.start)
+        self.context = SerializerContext(type: type, allocator: allocator)
+    }
 
     /// Serialise a message into a buffer using binary HTTP encoding.
     /// - Parameters:
-    ///   - message: The message to serialise.  File regions are currently not supported.
+    ///   - message: The message to serialise. File regions are currently not supported.
     ///   - buffer: Destination buffer to serialise into.
-    public func serialize(_ message: Message, into buffer: inout ByteBuffer) {
+    public mutating func serialize(_ message: Message, into buffer: inout ByteBuffer) throws {
         switch message {
         case .request(.head(let requestHead)):
-            Self.serializeRequestHead(requestHead, into: &buffer)
+            try self.fsm.writeRequestHead(requestHead, into: &buffer, using: &self.context)
+
         case .response(.head(let responseHead)):
-            Self.serializeResponseHead(responseHead, into: &buffer)
+            try self.fsm.writeResponseHead(responseHead, into: &buffer, using: &self.context)
+
         case .request(.body(.byteBuffer(let body))), .response(.body(.byteBuffer(let body))):
-            Self.serializeContentChunk(body, into: &buffer)
+            try self.fsm.writeBodyChunk(body, into: &buffer, using: &self.context)
+
         case .request(.body(.fileRegion)), .response(.body(.fileRegion)):
-            fatalError("fileregion unsupported")
+            throw ObliviousHTTPError.unsupportedOption(reason: "fileregion unsupported")
+
         case .request(.end(.some(let trailers))), .response(.end(.some(let trailers))):
-            // Send a 0 to terminate the body, then a field section.
-            buffer.writeInteger(UInt8(0))
-            Self.serializeIndeterminateLengthFieldSection(trailers, into: &buffer)
+            try self.fsm.writeTrailers(trailers, into: &buffer, using: &self.context)
+
         case .request(.end(.none)), .response(.end(.none)):
-            // We can omit the trailers in this context, but we will always send a zero
-            // byte, either to communicate no trailers or no body.
-            buffer.writeInteger(UInt8(0))
+            try self.fsm.writeRequestEnd(into: &buffer, using: &self.context)
         }
     }
 
-    private static func serializeRequestHead(_ head: HTTPRequestHead, into buffer: inout ByteBuffer) {
-        // First, the framing indicator. 2 for indeterminate length request.
-        buffer.writeVarint(2)
 
-        let method = head.method
-        let scheme = "https"  // Hardcoded for now, but not really the right option.
-        let path = head.uri
-        let authority = head.headers["Host"].first ?? ""
+}
 
-        buffer.writeVarintPrefixedString(method.rawValue)
-        buffer.writeVarintPrefixedString(scheme)
-        buffer.writeVarintPrefixedString(authority)
-        buffer.writeVarintPrefixedString(path)
+// Enum definitions for message, states, and types.
+extension BHTTPSerializer {
+
+    private struct SerializerContext {
+        private var chunkBuffer: ByteBuffer
+        private var fieldSectionBuffer: ByteBuffer
+        private var type: SerializerType
+
+
+        public init(
+            type: SerializerType,
+            allocator: ByteBufferAllocator = ByteBufferAllocator()
+        ) {
+            self.chunkBuffer = allocator.buffer(capacity: 0)
+            self.fieldSectionBuffer = allocator.buffer(capacity: 0)
+            self.type = type
+        }
+        private var requestFramingIndicator: Int {
+            switch self.type {
+            case .knownLength:
+                return FramingIndicator.requestKnownLength
+            default:
+                return FramingIndicator.requestIndeterminateLength
+            }
+        }
 
-        Self.serializeIndeterminateLengthFieldSection(head.headers, into: &buffer)
-    }
+        private var responseFramingIndicator: Int {
+            switch self.type {
+            case .knownLength:
+                return FramingIndicator.responseKnownLength
+            default:
+                return FramingIndicator.responseIndeterminateLength
+            }
+        }
+
+
+        mutating func serializeRequestHead(_ head: HTTPRequestHead, into buffer: inout ByteBuffer) {
+            // First, the framing indicator
+            buffer.writeVarint(requestFramingIndicator)
 
-    private static func serializeResponseHead(_ head: HTTPResponseHead, into buffer: inout ByteBuffer) {
-        // First, the framing indicator. 3 for indeterminate length response.
-        buffer.writeVarint(3)
-        buffer.writeVarint(Int(head.status.code))
-        Self.serializeIndeterminateLengthFieldSection(head.headers, into: &buffer)
-    }
+            let method = head.method
+            let scheme = "https"  // Hardcoded for now, but not really the right option.
+            let path = head.uri
+            let authority = head.headers["Host"].first ?? ""
 
-    private static func serializeContentChunk(_ chunk: ByteBuffer, into buffer: inout ByteBuffer) {
-        // Omit zero-length chunks.
-        if chunk.readableBytes == 0 { return }
-        buffer.writeVarintPrefixedImmutableBuffer(chunk)
-    }
+            buffer.writeVarintPrefixedString(method.rawValue)
+            buffer.writeVarintPrefixedString(scheme)
+            buffer.writeVarintPrefixedString(authority)
+            buffer.writeVarintPrefixedString(path)
 
-    private static func serializeIndeterminateLengthFieldSection(
-        _ fields: HTTPHeaders,
-        into buffer: inout ByteBuffer
-    ) {
-        for (name, value) in fields {
-            buffer.writeVarintPrefixedString(name)
-            buffer.writeVarintPrefixedString(value)
+            switch self.type {
+            case .knownLength:
+                self.stackKnownLengthFieldSection(head.headers)
+                self.serializeKnownLengthFieldSection(into: &buffer)
+            case .indeterminateLength:
+                Self.serializeIndeterminateLengthFieldSection(head.headers, into: &buffer)
+            default: break
+            }
+        }
+
+        mutating func serializeResponseHead(_ head: HTTPResponseHead, into buffer: inout ByteBuffer) {
+            // First, the framing indicator
+            buffer.writeVarint(responseFramingIndicator)
+
+            buffer.writeVarint(Int(head.status.code))
+
+            switch self.type {
+            case .knownLength:
+                self.stackKnownLengthFieldSection(head.headers)
+                self.serializeKnownLengthFieldSection(into: &buffer)
+            case .indeterminateLength:
+                Self.serializeIndeterminateLengthFieldSection(head.headers, into: &buffer)
+            default: break
+            }
+        }
+
+        mutating func serializeChunk(_ chunk: ByteBuffer, into buffer: inout ByteBuffer) {
+            switch self.type {
+            case .knownLength:
+                self.stackContentChunk(chunk)
+            case .indeterminateLength:
+                Self.serializeContentChunk(chunk, into: &buffer)
+            default: break
+            }
+        }
+
+        static func serializeContentChunk(_ chunk: ByteBuffer, into buffer: inout ByteBuffer) {
+            if chunk.readableBytes == 0 { return }
+            buffer.writeVarintPrefixedImmutableBuffer(chunk)
+        }
+
+        mutating func serializeStackedContent(into buffer: inout ByteBuffer) {
+            if self.chunkBuffer.readableBytes == 0 { return }
+            buffer.writeVarintPrefixedImmutableBuffer(self.chunkBuffer)
+            self.chunkBuffer.clear()
+        }
+
+        mutating func stackContentChunk(_ chunk: ByteBuffer) {
+            self.chunkBuffer.writeImmutableBuffer(chunk)
+        }
+
+        static func serializeIndeterminateLengthFieldSection(
+            _ fields: HTTPHeaders,
+            into buffer: inout ByteBuffer
+        ) {
+            for (name, value) in fields {
+                buffer.writeVarintPrefixedString(name)
+                buffer.writeVarintPrefixedString(value)
+            }
+            buffer.writeInteger(UInt8(0))  // End of field section
+        }
+
+        mutating func serializeTrailers(_ trailers: HTTPHeaders, into buffer: inout ByteBuffer) {
+            switch self.type {
+            case .knownLength:
+                self.serializeStackedContent(into: &buffer)
+                self.stackKnownLengthFieldSection(trailers)
+            case .indeterminateLength:
+                // Send a 0 to terminate the body, then a field section.
+                buffer.writeInteger(UInt8(0))
+                Self.serializeIndeterminateLengthFieldSection(trailers, into: &buffer)
+            default: break
+            }
+        }
+
+
+        mutating func endRequest(into buffer: inout ByteBuffer) {
+            switch self.type {
+            case .knownLength:
+                self.serializeStackedContent(into: &buffer)
+                self.serializeKnownLengthFieldSection(into: &buffer)
+            case .indeterminateLength:
+                buffer.writeInteger(UInt8(0))
+            default: break
+            }
+        }
+
+        mutating func stackKnownLengthFieldSection(_ fields: HTTPHeaders) {
+            for (name, value) in fields {
+                self.fieldSectionBuffer.writeVarintPrefixedString(name)
+                self.fieldSectionBuffer.writeVarintPrefixedString(value)
+            }
+        }
+
+
+        mutating func serializeKnownLengthFieldSection(into buffer: inout ByteBuffer) {
+            buffer.writeVarintPrefixedImmutableBuffer(self.fieldSectionBuffer)
+            self.fieldSectionBuffer.clear()
         }
-        // This is technically a varint but we can skip the check there because we know it can always encode in one byte.
-        buffer.writeInteger(UInt8(0))
     }
+
+    // Finite State Machine for managing transitions in BHTTPSerializer.
+    private struct BHTTPSerializerFSM {
+        private var currentState: BHTTPSerializerState
 
-}
+        init(initialState: BHTTPSerializerState) {
+            self.currentState = initialState
+        }
+
+        mutating func writeRequestHead(
+            _ requestHead: HTTPRequestHead,
+            into buffer: inout ByteBuffer,
+            using context: inout SerializerContext
+        ) throws {
+            try self.transition(to: .header)
+            context.serializeRequestHead(requestHead, into: &buffer)
+        }
+
+        mutating func writeResponseHead(
+            _ responseHead: HTTPResponseHead,
+            into buffer: inout ByteBuffer,
+            using context: inout SerializerContext
+        ) throws {
+            try self.transition(to: .header)
+            context.serializeResponseHead(responseHead, into: &buffer)
+        }
+
+        mutating func writeRequestEnd(
+            into buffer: inout ByteBuffer,
+            using context: inout SerializerContext
+        ) throws {
+            context.endRequest(into: &buffer)
+            try self.transition(to: .end)
+        }
+        mutating func writeBodyChunk(
+            _ body: ByteBuffer,
+            into buffer: inout ByteBuffer,
+            using context: inout SerializerContext
+        ) throws {
+            context.serializeChunk(body, into: &buffer)
+            try self.transition(to: .chunk)
+        }
+
+        mutating func writeTrailers(
+            _ trailers: HTTPHeaders,
+            into buffer: inout ByteBuffer,
+            using context: inout SerializerContext
+        ) throws {
+            context.serializeTrailers(trailers, into: &buffer)
+            try self.transition(to: .trailers)
+        }
+
+        private mutating func transition(to state: BHTTPSerializerState) throws {
+            let allowedNextStates: Set<BHTTPSerializerState>
+            switch currentState {
+            case .start:
+                allowedNextStates = [.header]
+            case .header:
+                allowedNextStates = [.chunk, .trailers, .end]
+            case .chunk:
+                allowedNextStates = [.trailers, .end, .chunk]
+            case .trailers:
+                allowedNextStates = [.trailers, .end]
+            case .end:
+                allowedNextStates = []
+            }
+            guard allowedNextStates.contains(state) else {
+                throw ObliviousHTTPError.unexpectedHTTPMessageSection()
+            }
+            currentState = state
+        }
+    }
 
-extension BHTTPSerializer {
-    /// Types of message for binary http serilaisation
     public enum Message {
-        /// Part of an HTTP request.
         case request(HTTPClientRequestPart)
-        /// Part of an HTTP response.
         case response(HTTPServerResponsePart)
     }
+
+    public struct SerializerType: Hashable, Sendable {
+        private enum InternalType: Hashable {
+            case knownLength
+            case indeterminateLength
+        }
+
+        private let type: InternalType
+
+        public static let knownLength = SerializerType(type: .knownLength)
+        public static let indeterminateLength = SerializerType(type: .indeterminateLength)
+
+        private init(type: InternalType) {
+            self.type = type
+        }
+
+    }
+
+    internal struct FramingIndicator {
+        static var requestKnownLength: Int { 0 }
+        static var responseKnownLength: Int { 1 }
+        static var requestIndeterminateLength: Int { 2 }
+        static var responseIndeterminateLength: Int { 3 }
+    }
+
+    internal enum BHTTPSerializerState {
+        case start
+        case header
+        case chunk
+        case trailers
+        case end
+    }
 }