Add support for sending IPv6 flow labels to QuicGsoBatchWriter

Protected by quic_restart_flag_quic_support_flow_label (which is enabling blocked by).

PiperOrigin-RevId: 693535451
diff --git a/build/source_list.bzl b/build/source_list.bzl
index 360be8d..c2ad15f 100644
--- a/build/source_list.bzl
+++ b/build/source_list.bzl
@@ -217,6 +217,7 @@
     "quic/core/crypto/transport_parameters.h",
     "quic/core/crypto/web_transport_fingerprint_proof_verifier.h",
     "quic/core/deterministic_connection_id_generator.h",
+    "quic/core/flow_label.h",
     "quic/core/frames/quic_ack_frame.h",
     "quic/core/frames/quic_ack_frequency_frame.h",
     "quic/core/frames/quic_blocked_frame.h",
diff --git a/build/source_list.gni b/build/source_list.gni
index c5b8040..06418b7 100644
--- a/build/source_list.gni
+++ b/build/source_list.gni
@@ -217,6 +217,7 @@
     "src/quiche/quic/core/crypto/transport_parameters.h",
     "src/quiche/quic/core/crypto/web_transport_fingerprint_proof_verifier.h",
     "src/quiche/quic/core/deterministic_connection_id_generator.h",
+    "src/quiche/quic/core/flow_label.h",
     "src/quiche/quic/core/frames/quic_ack_frame.h",
     "src/quiche/quic/core/frames/quic_ack_frequency_frame.h",
     "src/quiche/quic/core/frames/quic_blocked_frame.h",
diff --git a/build/source_list.json b/build/source_list.json
index 0869e11..c6b0392 100644
--- a/build/source_list.json
+++ b/build/source_list.json
@@ -216,6 +216,7 @@
     "quiche/quic/core/crypto/transport_parameters.h",
     "quiche/quic/core/crypto/web_transport_fingerprint_proof_verifier.h",
     "quiche/quic/core/deterministic_connection_id_generator.h",
+    "quiche/quic/core/flow_label.h",
     "quiche/quic/core/frames/quic_ack_frame.h",
     "quiche/quic/core/frames/quic_ack_frequency_frame.h",
     "quiche/quic/core/frames/quic_blocked_frame.h",
diff --git a/quiche/quic/core/batch_writer/quic_gso_batch_writer.cc b/quiche/quic/core/batch_writer/quic_gso_batch_writer.cc
index 6213e86..3009d1d 100644
--- a/quiche/quic/core/batch_writer/quic_gso_batch_writer.cc
+++ b/quiche/quic/core/batch_writer/quic_gso_batch_writer.cc
@@ -10,6 +10,7 @@
 #include <memory>
 #include <utility>
 
+#include "quiche/quic/core/flow_label.h"
 #include "quiche/quic/core/quic_linux_socket_utils.h"
 #include "quiche/quic/platform/api/quic_server_stats.h"
 
@@ -140,7 +141,8 @@
 void QuicGsoBatchWriter::BuildCmsg(QuicMsgHdr* hdr,
                                    const QuicIpAddress& self_address,
                                    uint16_t gso_size, uint64_t release_time,
-                                   QuicEcnCodepoint ecn_codepoint) {
+                                   QuicEcnCodepoint ecn_codepoint,
+                                   uint32_t flow_label) {
   hdr->SetIpInNextCmsg(self_address);
   if (gso_size > 0) {
     *hdr->GetNextCmsgData<uint16_t>(SOL_UDP, UDP_SEGMENT) = gso_size;
@@ -158,6 +160,11 @@
           static_cast<int>(ecn_codepoint);
     }
   }
+
+  if (flow_label != 0) {
+    *hdr->GetNextCmsgData<uint32_t>(IPPROTO_IPV6, IPV6_FLOWINFO) =
+        htonl(flow_label & IPV6_FLOWINFO_FLOWLABEL);
+  }
 }
 
 QuicGsoBatchWriter::FlushImplResult QuicGsoBatchWriter::FlushImpl() {
diff --git a/quiche/quic/core/batch_writer/quic_gso_batch_writer.h b/quiche/quic/core/batch_writer/quic_gso_batch_writer.h
index d9f0cb6..96a46ec 100644
--- a/quiche/quic/core/batch_writer/quic_gso_batch_writer.h
+++ b/quiche/quic/core/batch_writer/quic_gso_batch_writer.h
@@ -8,6 +8,7 @@
 #include <cstddef>
 
 #include "quiche/quic/core/batch_writer/quic_batch_writer_base.h"
+#include "quiche/quic/core/flow_label.h"
 #include "quiche/quic/core/quic_linux_socket_utils.h"
 
 namespace quic {
@@ -60,10 +61,11 @@
   }
 
   static const int kCmsgSpace = kCmsgSpaceForIp + kCmsgSpaceForSegmentSize +
-                                kCmsgSpaceForTxTime + kCmsgSpaceForTOS;
+                                kCmsgSpaceForTxTime + kCmsgSpaceForTOS +
+                                kCmsgSpaceForFlowLabel;
   static void BuildCmsg(QuicMsgHdr* hdr, const QuicIpAddress& self_address,
                         uint16_t gso_size, uint64_t release_time,
-                        QuicEcnCodepoint ecn_codepoint);
+                        QuicEcnCodepoint ecn_codepoint, uint32_t flow_label);
 
   template <size_t CmsgSpace, typename CmsgBuilderT>
   FlushImplResult InternalFlushImpl(CmsgBuilderT cmsg_builder) {
@@ -83,7 +85,7 @@
 
     uint16_t gso_size = buffered_writes().size() > 1 ? first.buf_len : 0;
     cmsg_builder(&hdr, first.self_address, gso_size, first.release_time,
-                 first.params.ecn_codepoint);
+                 first.params.ecn_codepoint, first.params.flow_label);
 
     write_result = QuicLinuxSocketUtils::WritePacket(fd(), hdr);
     QUIC_DVLOG(1) << "Write GSO packet result: " << write_result
diff --git a/quiche/quic/core/batch_writer/quic_gso_batch_writer_test.cc b/quiche/quic/core/batch_writer/quic_gso_batch_writer_test.cc
index 88aa891..f26ad28 100644
--- a/quiche/quic/core/batch_writer/quic_gso_batch_writer_test.cc
+++ b/quiche/quic/core/batch_writer/quic_gso_batch_writer_test.cc
@@ -12,6 +12,7 @@
 #include <utility>
 #include <vector>
 
+#include "quiche/quic/core/flow_label.h"
 #include "quiche/quic/platform/api/quic_ip_address.h"
 #include "quiche/quic/platform/api/quic_test.h"
 #include "quiche/quic/test_tools/quic_mock_syscall_wrapper.h"
@@ -541,6 +542,52 @@
   ASSERT_EQ(WriteResult(WRITE_STATUS_OK, 2700), result);
 }
 
+TEST_F(QuicGsoBatchWriterTest, FlowLabelIPv6) {
+  const WriteResult write_buffered(WRITE_STATUS_OK, 0);
+
+  self_address_ = QuicIpAddress::Any6();
+  peer_address_ = QuicSocketAddress(QuicIpAddress::Any6(), 443);
+  auto writer = TestQuicGsoBatchWriter::NewInstanceWithReleaseTimeSupport();
+
+  QuicPacketWriterParams params;
+  EXPECT_TRUE(params.release_time_delay.IsZero());
+  EXPECT_FALSE(params.allow_burst);
+
+  for (uint32_t i = 1; i < 5; ++i) {
+    // Generate flow label which are on both side of zero to test
+    // coverage when the in-memory label is larger than 20 bits.
+    params.flow_label = i - 2;
+    WriteResult result = WritePacketWithParams(writer.get(), params);
+    ASSERT_EQ(write_buffered, result);
+
+    EXPECT_CALL(mock_syscalls_, Sendmsg(_, _, _))
+        .WillOnce(
+            Invoke([&params](int /*sockfd*/, const msghdr* msg, int /*flags*/) {
+              EXPECT_EQ(1350u, PacketLength(msg));
+              msghdr mutable_msg;
+              memcpy(&mutable_msg, msg, sizeof(*msg));
+              bool found_flow_label = false;
+              for (struct cmsghdr* cmsg = CMSG_FIRSTHDR(&mutable_msg);
+                   cmsg != NULL; cmsg = CMSG_NXTHDR(&mutable_msg, cmsg)) {
+                if (cmsg->cmsg_level == IPPROTO_IPV6 &&
+                    cmsg->cmsg_type == IPV6_FLOWINFO) {
+                  found_flow_label = true;
+                  uint32_t cmsg_flow_label =
+                      ntohl(*reinterpret_cast<uint32_t*> CMSG_DATA(cmsg));
+                  EXPECT_EQ(params.flow_label & 0xFFFFF, cmsg_flow_label);
+                  break;
+                }
+              }
+              // As long as the flow label is not zero, it should be present.
+              EXPECT_EQ(params.flow_label != 0, found_flow_label);
+              errno = 0;
+              return 0;
+            }));
+    WriteResult error_result = writer->Flush();
+    ASSERT_EQ(WriteResult(WRITE_STATUS_OK, 1350), error_result);
+  }
+}
+
 }  // namespace
 }  // namespace test
 }  // namespace quic
diff --git a/quiche/quic/core/flow_label.h b/quiche/quic/core/flow_label.h
new file mode 100644
index 0000000..34d024e
--- /dev/null
+++ b/quiche/quic/core/flow_label.h
@@ -0,0 +1,23 @@
+// Copyright (c) 2024 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef QUICHE_QUIC_CORE_QUIC_FLOW_LABEL_H_
+#define QUICHE_QUIC_CORE_QUIC_FLOW_LABEL_H_
+
+#include <cstdint>
+
+#if defined(__linux__)
+#include <linux/in6.h>
+#include <sys/socket.h>
+
+#ifndef IPV6_FLOWLABEL
+#define IPV6_FLOWINFO 11
+#define IPV6_FLOWINFO_FLOWLABEL 0x000fffff
+#endif
+
+static constexpr int kCmsgSpaceForFlowLabel = CMSG_SPACE(sizeof(uint32_t));
+
+#endif
+
+#endif  // QUICHE_QUIC_CORE_QUIC_FLOW_LABEL_H_
diff --git a/quiche/quic/core/quic_udp_socket_posix.inc b/quiche/quic/core/quic_udp_socket_posix.inc
index f0c28b6..eb7413d 100644
--- a/quiche/quic/core/quic_udp_socket_posix.inc
+++ b/quiche/quic/core/quic_udp_socket_posix.inc
@@ -10,16 +10,8 @@
 #include <sys/socket.h>
 #include <sys/types.h>
 
-#if defined(__linux__)
-#include <linux/in6.h>
-#ifndef IPV6_FLOWLABEL
-#define IPV6_FLOWINFO 11
-#define IPV6_FLOWINFO_SEND 33
-#define IPV6_FLOWINFO_FLOWLABEL 0x000fffff
-#endif
-#endif
-
 #include "absl/base/optimization.h"
+#include "quiche/quic/core/flow_label.h"
 #include "quiche/quic/core/io/socket.h"
 #include "quiche/quic/core/quic_udp_socket.h"
 #include "quiche/quic/platform/api/quic_bug_tracker.h"