Add GRPCClientStreamHandler

Sources/GRPCHTTP2Core/Client/GRPCClientStreamHandler.swift

@@ -0,0 +1,238 @@
+/*
+ * 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 NIOHTTP2
+
+@available(macOS 13.0, iOS 16.0, watchOS 9.0, tvOS 16.0, *)
+final class GRPCClientStreamHandler: ChannelDuplexHandler {
+  typealias InboundIn = HTTP2Frame.FramePayload
+  typealias InboundOut = RPCResponsePart
+
+  typealias OutboundIn = RPCRequestPart
+  typealias OutboundOut = HTTP2Frame.FramePayload
+
+  private var stateMachine: GRPCStreamStateMachine
+
+  private var isReading = false
+  private var flushPending = false
+
+  init(
+    methodDescriptor: MethodDescriptor,
+    scheme: Scheme,
+    outboundEncoding: CompressionAlgorithm,
+    acceptedEncodings: [CompressionAlgorithm],
+    maximumPayloadSize: Int,
+    skipStateMachineAssertions: Bool = false
+  ) {
+    self.stateMachine = .init(
+      configuration: .client(
+        .init(
+          methodDescriptor: methodDescriptor,
+          scheme: scheme,
+          outboundEncoding: outboundEncoding,
+          acceptedEncodings: acceptedEncodings
+        )
+      ),
+      maximumPayloadSize: maximumPayloadSize,
+      skipAssertions: skipStateMachineAssertions
+    )
+  }
+}
+
+// - MARK: ChannelInboundHandler
+
+@available(macOS 13.0, iOS 16.0, watchOS 9.0, tvOS 16.0, *)
+extension GRPCClientStreamHandler {
+  func channelRead(context: ChannelHandlerContext, data: NIOAny) {
+    self.isReading = true
+    let frame = self.unwrapInboundIn(data)
+    switch frame {
+    case .data(let frameData):
+      let endStream = frameData.endStream
+      switch frameData.data {
+      case .byteBuffer(let buffer):
+        do {
+          try self.stateMachine.receive(buffer: buffer, endStream: endStream)
+          loop: while true {
+            switch self.stateMachine.nextInboundMessage() {
+            case .receiveMessage(let message):
+              context.fireChannelRead(self.wrapInboundOut(.message(message)))
+            case .awaitMoreMessages:
+              break loop
+            case .noMoreMessages:
+              context.fireUserInboundEventTriggered(ChannelEvent.inputClosed)
+              break loop
+            }
+          }
+        } catch {
+          context.fireErrorCaught(error)
+        }
+
+      case .fileRegion:
+        preconditionFailure("Unexpected IOData.fileRegion")
+      }
+
+    case .headers(let headers):
+      do {
+        let action = try self.stateMachine.receive(
+          headers: headers.headers,
+          endStream: headers.endStream
+        )
+        switch action {
+        case .receivedMetadata(let metadata):
+          context.fireChannelRead(self.wrapInboundOut(.metadata(metadata)))
+
+        case .rejectRPC:
+          throw RPCError(
+            code: .internalError,
+            message: "Client cannot get rejectRPC."
+          )
+
+        case .receivedStatusAndMetadata(let status, let metadata):
+          context.fireChannelRead(self.wrapInboundOut(.status(status, metadata)))
+
+        case .doNothing:
+          ()
+        }
+      } catch {
+        context.fireErrorCaught(error)
+      }
+
+    case .ping, .goAway, .priority, .rstStream, .settings, .pushPromise, .windowUpdate,
+      .alternativeService, .origin:
+      ()
+    }
+  }
+
+  func channelReadComplete(context: ChannelHandlerContext) {
+    self.isReading = false
+    if self.flushPending {
+      self.flushPending = false
+      self.flush(context: context)
+    }
+    context.fireChannelReadComplete()
+  }
+
+  func handlerRemoved(context: ChannelHandlerContext) {
+    self.stateMachine.tearDown()
+  }
+}
+
+// - MARK: ChannelOutboundHandler
+
+@available(macOS 13.0, iOS 16.0, watchOS 9.0, tvOS 16.0, *)
+extension GRPCClientStreamHandler {
+  func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
+    switch self.unwrapOutboundIn(data) {
+    case .metadata(let metadata):
+      do {
+        self.flushPending = true
+        let headers = try self.stateMachine.send(metadata: metadata)
+        context.write(self.wrapOutboundOut(.headers(.init(headers: headers))), promise: nil)
+        // TODO: move the promise handling into the state machine
+        promise?.succeed()
+      } catch {
+        context.fireErrorCaught(error)
+        // TODO: move the promise handling into the state machine
+        promise?.fail(error)
+      }
+
+    case .message(let message):
+      do {
+        try self.stateMachine.send(message: message)
+        // TODO: move the promise handling into the state machine
+        promise?.succeed()
+      } catch {
+        context.fireErrorCaught(error)
+        // TODO: move the promise handling into the state machine
+        promise?.fail(error)
+      }
+    }
+  }
+
+  func close(context: ChannelHandlerContext, mode: CloseMode, promise: EventLoopPromise<Void>?) {
+    switch mode {
+    case .output, .all:
+      do {
+        try self.stateMachine.closeOutbound()
+        // Force a flush by calling _flush
+        // (otherwise, we'd skip flushing if we're in a read loop)
+        self._flush(context: context)
+        context.close(mode: mode, promise: promise)
+      } catch {
+        promise?.fail(error)
+        context.fireErrorCaught(error)
+      }
+
+    case .input:
+      context.close(mode: .input, promise: promise)
+    }
+  }
+
+  func flush(context: ChannelHandlerContext) {
+    if self.isReading {
+      // We don't want to flush yet if we're still in a read loop.
+      self.flushPending = true
+      return
+    }
+
+    self._flush(context: context)
+  }
+
+  private func _flush(context: ChannelHandlerContext) {
+    do {
+      loop: while true {
+        switch try self.stateMachine.nextOutboundMessage() {
+        case .sendMessage(let byteBuffer):
+          self.flushPending = true
+          context.write(
+            self.wrapOutboundOut(.data(.init(data: .byteBuffer(byteBuffer)))),
+            promise: nil
+          )
+
+        case .noMoreMessages:
+          // Write an empty data frame with the EOS flag set, to signal the RPC
+          // request is now finished.
+          context.write(
+            self.wrapOutboundOut(
+              HTTP2Frame.FramePayload.data(
+                .init(
+                  data: .byteBuffer(.init()),
+                  endStream: true
+                )
+              )
+            ),
+            promise: nil
+          )
+
+          context.flush()
+          break loop
+
+        case .awaitMoreMessages:
+          if self.flushPending {
+            self.flushPending = false
+            context.flush()
+          }
+          break loop
+        }
+      }
+    } catch {
+      context.fireErrorCaught(error)
+    }
+  }
+}

Sources/GRPCHTTP2Core/GRPCStreamStateMachine.swift

@@ -373,8 +373,12 @@ struct GRPCStreamStateMachine {
 
   mutating func receive(headers: HPACKHeaders, endStream: Bool) throws -> OnMetadataReceived {
     switch self.configuration {
-    case .client:
-      return try self.clientReceive(headers: headers, endStream: endStream)
+    case .client(let clientConfiguration):
+      return try self.clientReceive(
+        headers: headers,
+        endStream: endStream,
+        configuration: clientConfiguration
+      )
     case .server(let serverConfiguration):
       return try self.serverReceive(
         headers: headers,
@@ -567,9 +571,7 @@ extension GRPCStreamStateMachine {
     case .clientOpenServerClosed(let state):
       self.state = .clientClosedServerClosed(.init(previousState: state))
     case .clientClosedServerIdle, .clientClosedServerOpen, .clientClosedServerClosed:
-      try self.invalidState(
-        "Client is closed, cannot send a message."
-      )
+      try self.invalidState("Client is already closed.")
     }
   }
 
@@ -665,7 +667,7 @@ extension GRPCStreamStateMachine {
         .receivedStatusAndMetadata(
           status: .init(
             code: .internalError,
-            message: "Missing \(GRPCHTTP2Keys.contentType) header"
+            message: "Missing \(GRPCHTTP2Keys.contentType.rawValue) header"
           ),
           metadata: Metadata(headers: metadata)
         )
@@ -680,18 +682,23 @@ extension GRPCStreamStateMachine {
     case success(CompressionAlgorithm)
   }
 
-  private func processInboundEncoding(_ metadata: HPACKHeaders) -> ProcessInboundEncodingResult {
+  private func processInboundEncoding(
+    headers: HPACKHeaders,
+    configuration: GRPCStreamStateMachineConfiguration.ClientConfiguration
+  ) -> ProcessInboundEncodingResult {
     let inboundEncoding: CompressionAlgorithm
-    if let serverEncoding = metadata.first(name: GRPCHTTP2Keys.encoding.rawValue) {
-      guard let parsedEncoding = CompressionAlgorithm(rawValue: serverEncoding) else {
+    if let serverEncoding = headers.first(name: GRPCHTTP2Keys.encoding.rawValue) {
+      guard let parsedEncoding = CompressionAlgorithm(rawValue: serverEncoding),
+        configuration.acceptedEncodings.contains(parsedEncoding)
+      else {
         return .error(
           .receivedStatusAndMetadata(
             status: .init(
               code: .internalError,
               message:
                 "The server picked a compression algorithm ('\(serverEncoding)') the client does not know about."
             ),
-            metadata: Metadata(headers: metadata)
+            metadata: Metadata(headers: headers)
           )
         )
       }
@@ -732,7 +739,8 @@ extension GRPCStreamStateMachine {
 
   private mutating func clientReceive(
     headers: HPACKHeaders,
-    endStream: Bool
+    endStream: Bool,
+    configuration: GRPCStreamStateMachineConfiguration.ClientConfiguration
   ) throws -> OnMetadataReceived {
     switch self.state {
     case .clientOpenServerIdle(let state):
@@ -750,7 +758,7 @@ extension GRPCStreamStateMachine {
         self.state = .clientOpenServerClosed(.init(previousState: state))
         return try self.validateAndReturnStatusAndMetadata(headers)
       case (.valid, false):
-        switch self.processInboundEncoding(headers) {
+        switch self.processInboundEncoding(headers: headers, configuration: configuration) {
         case .error(let failure):
           return failure
         case .success(let inboundEncoding):
@@ -798,7 +806,7 @@ extension GRPCStreamStateMachine {
         self.state = .clientClosedServerClosed(.init(previousState: state))
         return try self.validateAndReturnStatusAndMetadata(headers)
       case (.valid, false):
-        switch self.processInboundEncoding(headers) {
+        switch self.processInboundEncoding(headers: headers, configuration: configuration) {
         case .error(let failure):
           return failure
         case .success(let inboundEncoding):