Add reliable reset functionality to QuicWriteCodec

Summary: In this diff, I'm writing functionality to enable us to write RESET_STREAM_AT frames. You can find the format of the RST_STREAM_AT frame in the [RFC](https://datatracker.ietf.org/doc/html/draft-ietf-quic-reliable-stream-reset-06#name-reset_stream_at-frame). Reviewed By: jbeshay Differential Revision: D64907958 fbshipit-source-id: 55e82278ebd0649dc33c82b472be80893af02509
facebook · Nov 18, 2024 · 568f19c · 568f19c
1 parent 67794c5
commit 568f19c
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 7 deletions.
diff --git a/quic/codec/QuicWriteCodec.cpp b/quic/codec/QuicWriteCodec.cpp
@@ -814,18 +814,28 @@ size_t writeFrame(QuicWriteFrame&& frame, PacketBuilderInterface& builder) {
     }
     case QuicWriteFrame::Type::RstStreamFrame: {
       RstStreamFrame& rstStreamFrame = *frame.asRstStreamFrame();
-      QuicInteger intFrameType(static_cast<uint8_t>(FrameType::RST_STREAM));
+      QuicInteger intFrameType(static_cast<uint8_t>(
+          rstStreamFrame.reliableSize ? FrameType::RST_STREAM_AT
+                                      : FrameType::RST_STREAM));
       QuicInteger streamId(rstStreamFrame.streamId);
-      QuicInteger offset(rstStreamFrame.finalSize);
+      QuicInteger finalSize(rstStreamFrame.finalSize);
       QuicInteger errorCode(static_cast<uint64_t>(rstStreamFrame.errorCode));
+      folly::Optional<QuicInteger> maybeReliableSize = folly::none;
       size_t errorSize = errorCode.getSize();
       auto rstStreamFrameSize = intFrameType.getSize() + errorSize +
-          streamId.getSize() + offset.getSize();
+          streamId.getSize() + finalSize.getSize();
+      if (rstStreamFrame.reliableSize) {
+        maybeReliableSize = QuicInteger(*rstStreamFrame.reliableSize);
+        rstStreamFrameSize += maybeReliableSize->getSize();
+      }
       if (packetSpaceCheck(spaceLeft, rstStreamFrameSize)) {
         builder.write(intFrameType);
         builder.write(streamId);
         builder.write(errorCode);
-        builder.write(offset);
+        builder.write(finalSize);
+        if (maybeReliableSize) {
+          builder.write(*maybeReliableSize);
+        }
         builder.appendFrame(std::move(rstStreamFrame));
         return rstStreamFrameSize;
       }

diff --git a/quic/codec/test/QuicWriteCodecTest.cpp b/quic/codec/test/QuicWriteCodecTest.cpp
@@ -2122,6 +2122,38 @@ TEST_F(QuicWriteCodecTest, NoSpaceForRst) {
   EXPECT_EQ(0, writeFrame(rstStreamFrame, pktBuilder));
 }
 
+TEST_F(QuicWriteCodecTest, WriteRstStreamAt) {
+  MockQuicPacketBuilder pktBuilder;
+  setupCommonExpects(pktBuilder);
+  StreamId id = 0xBAAD;
+  ApplicationErrorCode errorCode = GenericApplicationErrorCode::UNKNOWN;
+  uint64_t finalSize = 0xF00D;
+  uint64_t reliableSize = 0xF00C;
+  RstStreamFrame rstStreamFrame(id, errorCode, finalSize, reliableSize);
+  auto rstStreamBytesWritten = writeFrame(rstStreamFrame, pktBuilder);
+
+  auto builtOut = std::move(pktBuilder).buildTestPacket();
+  auto regularPacket = builtOut.first;
+  EXPECT_EQ(21, rstStreamBytesWritten);
+  auto& resultRstStreamFrame = *regularPacket.frames[0].asRstStreamFrame();
+  EXPECT_EQ(errorCode, resultRstStreamFrame.errorCode);
+  EXPECT_EQ(id, resultRstStreamFrame.streamId);
+  EXPECT_EQ(finalSize, resultRstStreamFrame.finalSize);
+  EXPECT_EQ(reliableSize, resultRstStreamFrame.reliableSize);
+
+  auto wireBuf = std::move(builtOut.second);
+  BufQueue queue;
+  queue.append(wireBuf->clone());
+  QuicFrame decodedFrame = parseQuicFrame(queue);
+  auto& wireRstStreamFrame = *decodedFrame.asRstStreamFrame();
+  EXPECT_EQ(errorCode, wireRstStreamFrame.errorCode);
+  EXPECT_EQ(id, wireRstStreamFrame.streamId);
+  EXPECT_EQ(finalSize, wireRstStreamFrame.finalSize);
+  EXPECT_EQ(reliableSize, wireRstStreamFrame.reliableSize);
+  // At last, verify there is nothing left in the wire format bytes:
+  EXPECT_EQ(queue.chainLength(), 0);
+}
+
 TEST_F(QuicWriteCodecTest, WriteBlockedFrame) {
   MockQuicPacketBuilder pktBuilder;
   setupCommonExpects(pktBuilder);

diff --git a/quic/logging/QLoggerTypes.cpp b/quic/logging/QLoggerTypes.cpp
@@ -28,10 +28,14 @@ folly::dynamic NewTokenFrameLog::toDynamic() const {
 
 folly::dynamic RstStreamFrameLog::toDynamic() const {
   folly::dynamic d = folly::dynamic::object();
-  d["frame_type"] = toQlogString(FrameType::RST_STREAM);
+  d["frame_type"] = toQlogString(
+      reliableOffset ? FrameType::RST_STREAM_AT : FrameType::RST_STREAM);
   d["stream_id"] = streamId;
   d["error_code"] = errorCode;
   d["offset"] = offset;
+  if (reliableOffset) {
+    d["reliable_offset"] = reliableOffset.value();
+  }
   return d;
 }
 

diff --git a/quic/logging/QLoggerTypes.h b/quic/logging/QLoggerTypes.h
@@ -37,12 +37,17 @@ class RstStreamFrameLog : public QLogFrame {
   StreamId streamId;
   ApplicationErrorCode errorCode;
   uint64_t offset;
+  folly::Optional<uint64_t> reliableOffset;
 
   RstStreamFrameLog(
       StreamId streamIdIn,
       ApplicationErrorCode errorCodeIn,
-      uint64_t offsetIn)
-      : streamId{streamIdIn}, errorCode{errorCodeIn}, offset{offsetIn} {}
+      uint64_t offsetIn,
+      folly::Optional<uint64_t> reliableOffsetIn = folly::none)
+      : streamId{streamIdIn},
+        errorCode{errorCodeIn},
+        offset{offsetIn},
+        reliableOffset{reliableOffsetIn} {}
 
   ~RstStreamFrameLog() override = default;
   FOLLY_NODISCARD folly::dynamic toDynamic() const override;