Repository: grpc/grpc-swift (swift)

Issue:
Client doesn't handle DATA frames with end stream set gracefully
<!--
Please check the README, docs/ directory and other issues before filing an issue.

This form is for submitting BUG REPORTS only.

For troubleshooting or new feature requests please use one of the other templates.
-->

### Describe the bug

If the client receives a DATA frame with the end stream flag set then the RPC will fail with a "Transport became inactive" error; while that's true -- the transport _did_ become inactive -- it's not really representative of the actual problem: the remote peer is not conforming to the spec; it should be sending trailers with the end stream flag set.

### Expected behaviour

The RPC should fail with an accurate error message.

Pull request description:
Motivation:

If a peer doesn't follow spec and sends the client a DATA frame with end
stream set, the client will ignore the fact that end stream was set.
When the stream channel closes just after, the call is unaware of why it
was closed (because it hasn't received a final status) which leads to a
pretty non-descript error message.

Modifications:

Add a 'receiveEndOfResponseStream' to the client state machine which
produces a 'GRPCStatus' if the end stream is unexpected and should not
be ignored.

Result:

Calls will terminate with an 'internal error' status and a message about
a protocol violation if the server sends a data frame with end stream
set.

Resolves #1136

Patch:
diff --git a/Sources/GRPC/GRPCClientChannelHandler.swift b/Sources/GRPC/GRPCClientChannelHandler.swift
index b68d2c558..252e1f05e 100644
--- a/Sources/GRPC/GRPCClientChannelHandler.swift
+++ b/Sources/GRPC/GRPCClientChannelHandler.swift
@@ -425,6 +425,16 @@ extension GRPCClientChannelHandler: ChannelInboundHandler {
       MetadataKey.h2EndStream: "\(content.endStream)",
     ])
 
+    self.consumeBytes(from: &buffer, context: context)
+
+    // End stream is set; we don't usually expect this but can handle it in some situations.
+    if content.endStream, let status = self.stateMachine.receiveEndOfResponseStream() {
+      self.logger.warning("Unexpected end stream set on DATA frame")
+      context.fireChannelRead(self.wrapInboundOut(.status(status)))
+    }
+  }
+
+  private func consumeBytes(from buffer: inout ByteBuffer, context: ChannelHandlerContext) {
     // Do we have bytes to read? If there are no bytes to read then we can't do anything. This may
     // happen if the end-of-stream flag is not set on the trailing headers frame (i.e. the one
     // containing the gRPC status code) and an additional empty data frame is sent with the
diff --git a/Sources/GRPC/GRPCClientStateMachine.swift b/Sources/GRPC/GRPCClientStateMachine.swift
index cc367a48d..ac413c1d0 100644
--- a/Sources/GRPC/GRPCClientStateMachine.swift
+++ b/Sources/GRPC/GRPCClientStateMachine.swift
@@ -346,6 +346,21 @@ struct GRPCClientStateMachine {
     }
   }
 
+  /// Receive a DATA frame with the end stream flag set. Determines whether it is safe for the
+  /// caller to ignore the end stream flag or whether a synthesised status should be forwarded.
+  ///
+  /// Receiving a DATA frame with the end stream flag set is unexpected: the specification dictates
+  /// that an RPC should be ended by the server sending the client a HEADERS frame with end stream
+  /// set. However, we will tolerate end stream on a DATA frame if we believe the RPC has already
+  /// completed (i.e. we are in the 'clientClosedServerClosed' state). In cases where we don't
+  /// expect end of stream on a DATA frame we will emit a status with a message explaining
+  /// the protocol violation.
+  mutating func receiveEndOfResponseStream() -> GRPCStatus? {
+    return self.withStateAvoidingCoWs { state in
+      state.receiveEndOfResponseStream()
+    }
+  }
+
   /// Temporarily sets `self.state` to `.modifying` before calling the provided block and setting
   /// `self.state` to the `State` modified by the block.
   ///
@@ -555,6 +570,36 @@ extension GRPCClientStateMachine.State {
     return result
   }
 
+  /// See `GRPCClientStateMachine.receiveEndOfResponseStream()`.
+  mutating func receiveEndOfResponseStream() -> GRPCStatus? {
+    let status: GRPCStatus?
+
+    switch self {
+    case .clientIdleServerIdle:
+      // Can't see end stream before writing on it.
+      preconditionFailure()
+
+    case .clientActiveServerIdle,
+         .clientActiveServerActive,
+         .clientClosedServerIdle,
+         .clientClosedServerActive:
+      self = .clientClosedServerClosed
+      status = .init(
+        code: .internalError,
+        message: "Protocol violation: received DATA frame with end stream set"
+      )
+
+    case .clientClosedServerClosed:
+      // We've already closed. Ignore this.
+      status = nil
+
+    case .modifying:
+      preconditionFailure("State left as 'modifying'")
+    }
+
+    return status
+  }
+
   /// Makes the request headers (`Request-Headers` in the specification) used to initiate an RPC
   /// call.
   ///


Test patch:
diff --git a/Tests/GRPCTests/GRPCClientChannelHandlerTests.swift b/Tests/GRPCTests/GRPCClientChannelHandlerTests.swift
new file mode 100644
index 000000000..5708809af
--- /dev/null
+++ b/Tests/GRPCTests/GRPCClientChannelHandlerTests.swift
@@ -0,0 +1,94 @@
+/*
+ * Copyright 2021, 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.
+ */
+@testable import GRPC
+import NIO
+import NIOHPACK
+import NIOHTTP2
+import XCTest
+
+class GRPCClientChannelHandlerTests: GRPCTestCase {
+  private func makeRequestHead() -> _GRPCRequestHead {
+    return _GRPCRequestHead(
+      method: "POST",
+      scheme: "https",
+      path: "/foo/bar",
+      host: "localhost",
+      deadline: .distantFuture,
+      customMetadata: [:],
+      encoding: .disabled
+    )
+  }
+
+  func doTestDataFrameWithEndStream(dataContainsMessage: Bool) throws {
+    let handler = GRPCClientChannelHandler(
+      callType: .unary,
+      logger: self.clientLogger
+    )
+
+    let channel = EmbeddedChannel(handler: handler)
+
+    // Write request head.
+    let head = self.makeRequestHead()
+    XCTAssertNoThrow(try channel.writeOutbound(_RawGRPCClientRequestPart.head(head)))
+    // Read out a frame payload.
+    XCTAssertNotNil(try channel.readOutbound(as: HTTP2Frame.FramePayload.self))
+
+    // Respond with headers.
+    let headers: HPACKHeaders = [":status": "200", "content-type": "application/grpc"]
+    let headersPayload = HTTP2Frame.FramePayload.headers(.init(headers: headers))
+    XCTAssertNoThrow(try channel.writeInbound(headersPayload))
+    // Read them out the other side.
+    XCTAssertNotNil(try channel.readInbound(as: _RawGRPCClientResponsePart.self))
+
+    // Respond with DATA and end stream.
+    var buffer = ByteBuffer()
+
+    // Write a message, if we need to.
+    if dataContainsMessage {
+      buffer.writeInteger(UInt8(0)) // not compressed
+      buffer.writeInteger(UInt32(42)) // message length
+      buffer.writeRepeatingByte(0, count: 42) // message
+    }
+
+    let dataPayload = HTTP2Frame.FramePayload.Data(data: .byteBuffer(buffer), endStream: true)
+    XCTAssertNoThrow(try channel.writeInbound(HTTP2Frame.FramePayload.data(dataPayload)))
+
+    if dataContainsMessage {
+      // Read the message out the other side.
+      XCTAssertNotNil(try channel.readInbound(as: _RawGRPCClientResponsePart.self))
+    }
+
+    // We should also generate a status since end stream was set.
+    if let part = try channel.readInbound(as: _RawGRPCClientResponsePart.self) {
+      switch part {
+      case .initialMetadata, .message, .trailingMetadata:
+        XCTFail("Unexpected response part")
+      case .status:
+        () // Expected
+      }
+    } else {
+      XCTFail("Expected to read another response part")
+    }
+  }
+
+  func testDataFrameWithEndStream() throws {
+    try self.doTestDataFrameWithEndStream(dataContainsMessage: true)
+  }
+
+  func testEmptyDataFrameWithEndStream() throws {
+    try self.doTestDataFrameWithEndStream(dataContainsMessage: false)
+  }
+}
diff --git a/Tests/GRPCTests/GRPCClientStateMachineTests.swift b/Tests/GRPCTests/GRPCClientStateMachineTests.swift
index 44d953461..c2737a6fc 100644
--- a/Tests/GRPCTests/GRPCClientStateMachineTests.swift
+++ b/Tests/GRPCTests/GRPCClientStateMachineTests.swift
@@ -471,6 +471,45 @@ extension GRPCClientStateMachineTests {
       expected: .invalidState
     )
   }
+
+  private func doTestReceiveEndStreamOnDataWhenActive(_ state: StateMachine.State) throws {
+    var stateMachine = self.makeStateMachine(state)
+    let status = try assertNotNil(stateMachine.receiveEndOfResponseStream())
+    XCTAssertEqual(status.code, .internalError)
+  }
+
+  func testReceiveEndStreamOnDataClientActiveServerIdle() throws {
+    try self.doTestReceiveEndStreamOnDataWhenActive(
+      .clientActiveServerIdle(
+        writeState: .one(),
+        pendingReadState: .init(arity: .one, messageEncoding: .disabled)
+      )
+    )
+  }
+
+  func testReceiveEndStreamOnDataClientClosedServerIdle() throws {
+    try self.doTestReceiveEndStreamOnDataWhenActive(
+      .clientClosedServerIdle(pendingReadState: .init(arity: .one, messageEncoding: .disabled))
+    )
+  }
+
+  func testReceiveEndStreamOnDataClientActiveServerActive() throws {
+    try self.doTestReceiveEndStreamOnDataWhenActive(
+      .clientActiveServerActive(writeState: .one(), readState: .one())
+    )
+  }
+
+  func testReceiveEndStreamOnDataClientClosedServerActive() throws {
+    try self.doTestReceiveEndStreamOnDataWhenActive(
+      .clientClosedServerActive(readState: .one())
+    )
+  }
+
+  func testReceiveEndStreamOnDataWhenClosed() {
+    var stateMachine = self.makeStateMachine(.clientClosedServerClosed)
+    // Already closed, end stream is ignored.
+    XCTAssertNil(stateMachine.receiveEndOfResponseStream())
+  }
 }
 
 // MARK: - Basic RPC flow.