Frame datagrams on the way in via frameDatagram

Summary:
Every transport framed its egress datagrams in a different place: H3 wrapped
them in a `sendDatagram` override, while coro and qmux framed them in their
write loops. Because framing happened that late, a datagram too large to ever
fit on the wire was accepted by `sendDatagram` and then silently dropped by the
write loop. This adds a `frameDatagram` hook on `WtSessionBase` that each
transport overrides once, so what gets queued is already on-wire bytes. Coro
and qmux check their size limits while framing, so callers now get an error
back from `sendDatagram` instead of losing the datagram later. H2 has no
datagram egress path yet, so its override has no caller until the capsule
transports are wired up.

Reviewed By: sharmafb

Differential Revision: D115362262

fbshipit-source-id: c96e329f42aa3213ec193c2f1509681a4b229c9f

Author: afrind

proxygen/lib/http/coro/test/HttpWtUpstreamTests.cpp

@@ -777,6 +777,18 @@ CO_TEST_P_X(WtTest, MaxStreamsBidiUni) {
   wt->closeSession();
 }
 
+// A datagram too large to ever fit in a write chunk is rejected up front
+// rather than left for the write loop to rediscover on every pass.
+CO_TEST_P_X(WtTest, SendDatagramExceedsMaxWriteSize) {
+  XCHECK(wt);
+  auto res = wt->sendDatagram(makeBuf(65'535));
+  EXPECT_TRUE(res.hasError());
+  co_await rescheduleN(2);
+  EXPECT_TRUE(wtCodecCb.dgrams.empty());
+
+  wt->closeSession();
+}
+
 CO_TEST_P_X(WtTest, Datagrams) {
   XCHECK(wt);
   // tx datagram to peer

proxygen/lib/http/coro/util/CoroWtSession.cpp

@@ -50,11 +50,33 @@ WtExpected<folly::Unit>::Type CoroWtSession::closeSession(
 
 WtExpected<folly::Unit>::Type CoroWtSession::sendDatagram(
     IoBufPtr datagram) noexcept {
-  if (!writeLoopDone_) {
-    std::ignore = WtSessionBase::sendDatagram(std::move(datagram));
+  if (writeLoopDone_) {
+    return folly::unit;
+  }
+  auto res = WtSessionBase::sendDatagram(std::move(datagram));
+  if (res.hasValue()) {
     wtSmEgressCb.waitForEvent.signal(); // wake up write loop
   }
-  return folly::unit;
+  return res;
+}
+
+CoroWtSession::IoBufPtr CoroWtSession::frameDatagram(
+    IoBufPtr datagram) noexcept {
+  folly::IOBufQueue queue{folly::IOBufQueue::cacheChainLength()};
+  if (!writeDatagram(
+          queue,
+          DatagramCapsule{.httpDatagramPayload = std::move(datagram)},
+          FrameProtocol::WT_CAPSULE)) {
+    return nullptr;
+  }
+  // A datagram larger than a whole write chunk could never fit, so reject it
+  // here rather than letting the write loop discover it every pass.
+  if (queue.chainLength() > kMaxWriteSize) {
+    XLOG(ERR) << "datagram exceeds max write size; len=" << queue.chainLength()
+              << " kMaxWriteSize=" << kMaxWriteSize;
+    return nullptr;
+  }
+  return queue.move();
 }
 
 using WtCapsuleCallback = proxygen::detail::WtCapsuleCallback;
@@ -116,7 +138,7 @@ folly::coro::Task<void> CoroWtSession::writeLoop(Ptr self) {
     // TODO: currently datagrams can indefinitely preempt stream data
     auto datagrams = moveEgressDatagrams();
     for (auto& dgram : datagrams) {
-      writeDatagram(egressBuf, DatagramCapsule{std::move(dgram)});
+      egressBuf.append(std::move(dgram));
     }
 
     auto* wh = sm.nextWritable();

proxygen/lib/http/coro/util/CoroWtSession.h

@@ -122,6 +122,8 @@ class CoroWtSession
   void start(Ptr self);
 
  private:
+  IoBufPtr frameDatagram(IoBufPtr datagram) noexcept override;
+
   folly::coro::Task<void> readLoop(Ptr self);
   folly::coro::Task<void> writeLoop(Ptr self);
 

proxygen/lib/http/webtransport/QuicWtSession.cpp

@@ -201,7 +201,11 @@ QuicWtSessionBase::awaitBidiStreamCredit() noexcept {
 folly::Expected<folly::Unit, WebTransport::ErrorCode>
 QuicWtSessionBase::sendDatagram(IoBufPtr datagram) noexcept {
   XCHECK(quicSocket_);
-  auto writeRes = quicSocket_->writeDatagram(std::move(datagram));
+  auto framed = frameDatagram(std::move(datagram));
+  if (!framed) {
+    return folly::makeUnexpected(ErrorCode::GENERIC_ERROR);
+  }
+  auto writeRes = quicSocket_->writeDatagram(std::move(framed));
   if (writeRes.hasError()) {
     XLOG(ERR) << __func__ << "; err= " << writeRes.error();
     return folly::makeUnexpected(ErrorCode::GENERIC_ERROR);
@@ -511,8 +515,7 @@ folly::Expected<folly::Unit, WebTransport::ErrorCode> H3WtSession::closeSession(
   return folly::unit;
 }
 
-folly::Expected<folly::Unit, WebTransport::ErrorCode> H3WtSession::sendDatagram(
-    IoBufPtr datagram) noexcept {
+H3WtSession::IoBufPtr H3WtSession::frameDatagram(IoBufPtr datagram) noexcept {
   /**
    * RFC9297:
    * HTTP/3 Datagram {
@@ -528,7 +531,7 @@ folly::Expected<folly::Unit, WebTransport::ErrorCode> H3WtSession::sendDatagram(
       quarterStreamId, [&appender](auto val) { appender.writeBE(val); });
   XCHECK(encodeRes);
   queue.append(std::move(datagram));
-  return QuicWtSessionBase::sendDatagram(queue.move());
+  return queue.move();
 }
 
 folly::Expected<WebTransport::StreamWriteHandle*, WebTransport::ErrorCode>

proxygen/lib/http/webtransport/QuicWtSession.h

@@ -257,13 +257,12 @@ class H3WtSession final : public QuicWtSessionBase {
       override;
 
   /**
-   * Serializes the http/3 datagram, i.e. prefixes the datagram w/ quarter
-   * stream id (connectStreamId_ / 4) per RFC9297. Similarly to
+   * Prefixes the datagram w/ quarter stream id (connectStreamId_ / 4) per
+   * RFC9297. Applies to every sendDatagram overload. Similarly to
    * ::create(Uni|Bidi)Stream functions above, ::sendDatagram writes bypass
    * the http session and directly write to the QuicSocket.
    */
-  folly::Expected<folly::Unit, ErrorCode> sendDatagram(
-      IoBufPtr datagram) noexcept override;
+  IoBufPtr frameDatagram(IoBufPtr datagram) noexcept override;
 
   /**
    * Bidirectionally closes all associated webtransport streams. Notifies the

proxygen/lib/http/webtransport/WebTransportSession.cpp

@@ -124,6 +124,17 @@ H2WtSession::H2WtSession(folly::EventBase* evb,
       wtHandler_(wtHandler) {
 }
 
+H2WtSession::IoBufPtr H2WtSession::frameDatagram(IoBufPtr datagram) noexcept {
+  folly::IOBufQueue queue{folly::IOBufQueue::cacheChainLength()};
+  if (!writeDatagram(
+          queue,
+          DatagramCapsule{.httpDatagramPayload = std::move(datagram)},
+          FrameProtocol::WT_CAPSULE)) {
+    return nullptr;
+  }
+  return queue.move();
+}
+
 H2WtSession::~H2WtSession() noexcept {
   // abort txn and detach handler if applicable
   if (auto* txn = std::exchange(txnHandler_.txn_, nullptr)) {

proxygen/lib/http/webtransport/WebTransportSession.h

@@ -151,6 +151,8 @@ class H2WtSession
               WtStreamManager& sm,
               WebTransportHandler::Ptr& wtHandler) noexcept;
 
+  IoBufPtr frameDatagram(IoBufPtr datagram) noexcept override;
+
   WtStreamManager& sm_;
   WebTransportHandler::Ptr& wtHandler_;
 };

proxygen/lib/http/webtransport/WtUtils.cpp

@@ -501,7 +501,11 @@ WtExpected<folly::Unit>::Type WtSessionBase::stopSending(
 
 WtExpected<folly::Unit>::Type WtSessionBase::sendDatagram(
     IoBufPtr datagram) noexcept {
-  datagrams_.egress.push_back(std::move(datagram));
+  auto framed = frameDatagram(std::move(datagram));
+  if (!framed) {
+    return folly::makeUnexpected(WebTransport::ErrorCode::GENERIC_ERROR);
+  }
+  datagrams_.egress.push_back(std::move(framed));
   return folly::unit;
 }
 

proxygen/lib/http/webtransport/WtUtils.h

@@ -225,10 +225,17 @@ class WtSessionBase : public WebTransport {
   std::vector<IoBufPtr> moveIngressDatagrams() noexcept {
     return std::move(datagrams_.ingress);
   }
+  // Egress datagrams are framed on the way in, so what is queued is already
+  // on-wire bytes.
   std::vector<IoBufPtr> moveEgressDatagrams() noexcept {
     return std::move(datagrams_.egress);
   }
 
+  // Applies this transport's per-datagram framing. Returns null to reject.
+  virtual IoBufPtr frameDatagram(IoBufPtr datagram) noexcept {
+    return datagram;
+  }
+
  private:
   folly::Executor* executor_;
   WtStreamManager& sm_;

proxygen/lib/transport/qmux/QmuxSession.cpp

@@ -174,11 +174,34 @@ folly::AsyncTransport* QmuxSession::getUnderlyingTransport() const noexcept {
 
 proxygen::detail::WtExpected<folly::Unit>::Type QmuxSession::sendDatagram(
     std::unique_ptr<folly::IOBuf> datagram) noexcept {
-  pendingDatagrams_.emplace_back(std::move(datagram));
+  auto framed = frameDatagram(std::move(datagram));
+  if (!framed) {
+    return folly::makeUnexpected(WebTransport::ErrorCode::GENERIC_ERROR);
+  }
+  pendingDatagrams_.emplace_back(std::move(framed));
   wtSmEgressCb.waitForEvent.signal();
   return folly::unit;
 }
 
+QmuxSession::IoBufPtr QmuxSession::frameDatagram(IoBufPtr datagram) noexcept {
+  folly::IOBufQueue queue{folly::IOBufQueue::cacheChainLength()};
+  if (!writeDatagram(
+          queue,
+          DatagramCapsule{.httpDatagramPayload = std::move(datagram)},
+          FrameProtocol::QMUX)) {
+    return nullptr;
+  }
+  // A datagram that cannot fit in a single record is unsendable, so reject it
+  // here rather than letting the write loop discover it and drop it.
+  if (queue.chainLength() > peerMaxRecordSize_) {
+    XLOG(ERR)
+        << "datagram exceeds peer max_record_size; len=" << queue.chainLength()
+        << " peerMaxRecordSize=" << peerMaxRecordSize_;
+    return nullptr;
+  }
+  return queue.move();
+}
+
 void QmuxSession::start(Ptr self) {
   XLOG(DBG4) << "QmuxSession::start dir=" << (peerAddr_.describe());
   if (wtHandler_) {
@@ -366,15 +389,11 @@ folly::coro::Task<void> QmuxSession::writeLoop(Ptr self) {
 
     while (!pendingDatagrams_.empty()) {
       folly::IOBufQueue frameBuf{folly::IOBufQueue::cacheChainLength()};
-      writeDatagram(frameBuf,
-                    DatagramCapsule{.httpDatagramPayload =
-                                        std::move(pendingDatagrams_.front())},
-                    FrameProtocol::QMUX);
+      frameBuf.append(std::move(pendingDatagrams_.front()));
       pendingDatagrams_.pop_front();
-      // Returns false only for a datagram too large for any record, which it
-      // logs and drops.
-      std::ignore = appendFrameToRecord(
-          recordBuf, egressBuf, frameBuf, recordPayloadLimit);
+      // Cannot fail; frameDatagram rejects anything too large for a record.
+      XCHECK(appendFrameToRecord(
+          recordBuf, egressBuf, frameBuf, recordPayloadLimit));
     }
 
     // Flush all accumulated QMux records in one transport write.