From 1ff5e1062598b7d70991f65263a453af868d8be5 Mon Sep 17 00:00:00 2001
From: Aman Sharma <amsharma@meta.com>
Date: Thu, 20 Feb 2025 14:46:05 -0800
Subject: [PATCH] Implement Delivery callbacks for QuicWebTransport

Summary: I'm modifying `QuicWebTransport::sendWebTransportStreamData` so that delivery callbacks are conveyed to the passed-in `DeliveryCallback*` if it is non-`nullptr`.

Reviewed By: hanidamlaj

Differential Revision: D69618959

fbshipit-source-id: 280c8581ed00bf1364f24cfb8f269d7883e872ea
---
 .../http/webtransport/QuicWebTransport.cpp    | 12 +++++++--
 .../test/QuicWebTransportTest.cpp             | 26 +++++++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/proxygen/lib/http/webtransport/QuicWebTransport.cpp b/proxygen/lib/http/webtransport/QuicWebTransport.cpp
index 160879ef46..d97223e7be 100644
--- a/proxygen/lib/http/webtransport/QuicWebTransport.cpp
+++ b/proxygen/lib/http/webtransport/QuicWebTransport.cpp
@@ -7,6 +7,7 @@
  */
 
 #include <proxygen/lib/http/webtransport/QuicWebTransport.h>
+#include <proxygen/lib/http/webtransport/QuicWtDeliveryCallbackWrapper.h>
 
 using FCState = proxygen::WebTransport::FCState;
 
@@ -97,9 +98,16 @@ QuicWebTransport::sendWebTransportStreamData(
     HTTPCodec::StreamID id,
     std::unique_ptr<folly::IOBuf> data,
     bool eof,
-    WebTransport::DeliveryCallback* /* deliveryCallback */) {
+    DeliveryCallback* deliveryCallback) {
   XCHECK(quicSocket_);
-  auto res = quicSocket_->writeChain(id, std::move(data), eof);
+  std::unique_ptr<QuicWtDeliveryCallbackWrapper> deliveryCallbackWrapper =
+      nullptr;
+  if (deliveryCallback) {
+    deliveryCallbackWrapper =
+        std::make_unique<QuicWtDeliveryCallbackWrapper>(deliveryCallback, 0);
+  }
+  auto res = quicSocket_->writeChain(
+      id, std::move(data), eof, deliveryCallbackWrapper.release());
   if (!res) {
     return folly::makeUnexpected(WebTransport::ErrorCode::GENERIC_ERROR);
   }
diff --git a/proxygen/lib/http/webtransport/test/QuicWebTransportTest.cpp b/proxygen/lib/http/webtransport/test/QuicWebTransportTest.cpp
index e6159c0487..7f1adb467e 100644
--- a/proxygen/lib/http/webtransport/test/QuicWebTransportTest.cpp
+++ b/proxygen/lib/http/webtransport/test/QuicWebTransportTest.cpp
@@ -17,6 +17,13 @@ using namespace proxygen::test;
 using namespace testing;
 using quic::MockQuicSocketDriver;
 
+class MockDeliveryCallback : public WebTransport::DeliveryCallback {
+ public:
+  MOCK_METHOD(void, onDelivery, (uint64_t, uint32_t), (noexcept));
+
+  MOCK_METHOD(void, onDeliveryCancelled, (uint64_t, uint32_t), (noexcept));
+};
+
 namespace {
 const uint32_t WT_ERROR_1 = 19;
 const uint32_t WT_ERROR_2 = 77;
@@ -149,6 +156,25 @@ TEST_F(QuicWebTransportTest, SetPriority) {
   eventBase_.loop();
 }
 
+TEST_F(QuicWebTransportTest, WriteWithDeliveryCallback) {
+  auto handle = webTransport()->createUniStream();
+  EXPECT_TRUE(handle.hasValue());
+  auto mockCallback = std::make_unique<StrictMock<MockDeliveryCallback>>();
+
+  folly::StringPiece data = "test data";
+
+  uint64_t expectedStreamId = handle.value()->getID();
+  uint32_t expectedOffset = data.size();
+  EXPECT_CALL(*mockCallback, onDelivery(expectedStreamId, expectedOffset))
+      .Times(1);
+
+  handle.value()->writeStreamData(
+      folly::IOBuf::copyBuffer(data), true, mockCallback.get());
+  // The MockQuicSocketDriver automatically simulates the delivery of data
+  // written to the QUIC socket.
+  eventBase_.loop();
+}
+
 // TODO:
 //
 // new streams with no handler