From 02c50615a44da21d7dfd7f71307a731b3da22f05 Mon Sep 17 00:00:00 2001
From: Xintong Hu <huxintong@fb.com>
Date: Fri, 16 Jul 2021 13:34:39 -0700
Subject: [PATCH] Support setting cmsg per accepted socket

Summary:
For each `QuicServerTransport` created by `QuicServerWorker`, we are able to set the `cmsg`s to be sent whenever there's a write.

With this we can set things like TOS/SO_MARK per connection

Reviewed By: bschlinker

Differential Revision: D29313608

fbshipit-source-id: d60c97f65681086ae1079b6b95beade95158ec59
---
 quic/api/QuicTransportBase.cpp          |  8 ++++++++
 quic/api/QuicTransportBase.h            |  9 +++++++++
 quic/api/test/QuicTransportBaseTest.cpp | 11 +++++++++++
 3 files changed, 28 insertions(+)

diff --git a/quic/api/QuicTransportBase.cpp b/quic/api/QuicTransportBase.cpp
index eebda2af7..c849f0878 100644
--- a/quic/api/QuicTransportBase.cpp
+++ b/quic/api/QuicTransportBase.cpp
@@ -3298,4 +3298,12 @@ void QuicTransportBase::notifyAppRateLimited() {
   }
 }
 
+void QuicTransportBase::setCmsgs(const folly::SocketOptionMap& options) {
+  socket_->setCmsgs(options);
+}
+
+void QuicTransportBase::appendCmsgs(const folly::SocketOptionMap& options) {
+  socket_->appendCmsgs(options);
+}
+
 } // namespace quic
diff --git a/quic/api/QuicTransportBase.h b/quic/api/QuicTransportBase.h
index c164a1565..247761195 100644
--- a/quic/api/QuicTransportBase.h
+++ b/quic/api/QuicTransportBase.h
@@ -659,6 +659,15 @@ class QuicTransportBase : public QuicSocket {
   folly::Expected<std::vector<Buf>, LocalErrorCode> readDatagrams(
       size_t atMost = 0) override;
 
+  /**
+   * Set control messages to be sent for socket_ write, note that it's for this
+   * specific transport and does not change the other sockets sharing the same
+   * fd.
+   */
+  void setCmsgs(const folly::SocketOptionMap& options);
+
+  void appendCmsgs(const folly::SocketOptionMap& options);
+
  protected:
   void updateCongestionControlSettings(
       const TransportSettings& transportSettings);
diff --git a/quic/api/test/QuicTransportBaseTest.cpp b/quic/api/test/QuicTransportBaseTest.cpp
index 045f2eb1e..3387dcce0 100644
--- a/quic/api/test/QuicTransportBaseTest.cpp
+++ b/quic/api/test/QuicTransportBaseTest.cpp
@@ -3845,5 +3845,16 @@ TEST_F(QuicTransportImplTest, DatagramCallbackDatagramAvailable) {
   transport->driveReadCallbacks();
 }
 
+TEST_F(QuicTransportImplTest, Cmsgs) {
+  transport->setServerConnectionId();
+  folly::SocketOptionMap cmsgs;
+  cmsgs[{IPPROTO_IP, IP_TOS}] = 123;
+  EXPECT_CALL(*socketPtr, setCmsgs(_)).Times(1);
+  transport->setCmsgs(cmsgs);
+
+  EXPECT_CALL(*socketPtr, appendCmsgs(_)).Times(1);
+  transport->appendCmsgs(cmsgs);
+}
+
 } // namespace test
 } // namespace quic