Support draft-ietf-masque-h3-datagram-05

This revision makes datagram contexts optional to implement and negotiated via the Sec-Use-Datagram-Contexts header. See the PR for full details: https://github.com/ietf-wg-masque/draft-ietf-masque-h3-datagram/pull/105

This CL retains backwards compatibility with draft-ietf-masque-h3-datagram-04 by keeping support for the previous DATAGRAM capsule, which is now renamed to LEGACY_DATAGRAM. In this CL, our WebTransport implementation does not request support for contexts, so they will not be used for WebTransport.

This code isn't used in our servers and therefore does not require flag protection.

PiperOrigin-RevId: 406461234
diff --git a/quic/core/http/capsule.cc b/quic/core/http/capsule.cc
index c811bb8..21df166 100644
--- a/quic/core/http/capsule.cc
+++ b/quic/core/http/capsule.cc
@@ -24,8 +24,12 @@
       return "REGISTER_DATAGRAM_CONTEXT";
     case CapsuleType::CLOSE_DATAGRAM_CONTEXT:
       return "CLOSE_DATAGRAM_CONTEXT";
-    case CapsuleType::DATAGRAM:
-      return "DATAGRAM";
+    case CapsuleType::LEGACY_DATAGRAM:
+      return "LEGACY_DATAGRAM";
+    case CapsuleType::DATAGRAM_WITH_CONTEXT:
+      return "DATAGRAM_WITH_CONTEXT";
+    case CapsuleType::DATAGRAM_WITHOUT_CONTEXT:
+      return "DATAGRAM_WITHOUT_CONTEXT";
     case CapsuleType::REGISTER_DATAGRAM_NO_CONTEXT:
       return "REGISTER_DATAGRAM_NO_CONTEXT";
     case CapsuleType::CLOSE_WEBTRANSPORT_SESSION:
@@ -80,11 +84,27 @@
 
 Capsule::Capsule(CapsuleType capsule_type) : capsule_type_(capsule_type) {
   switch (capsule_type) {
-    case CapsuleType::DATAGRAM:
-      static_assert(std::is_standard_layout<DatagramCapsule>::value &&
-                        std::is_trivially_destructible<DatagramCapsule>::value,
-                    "All capsule structs must have these properties");
-      datagram_capsule_ = DatagramCapsule();
+    case CapsuleType::LEGACY_DATAGRAM:
+      static_assert(
+          std::is_standard_layout<LegacyDatagramCapsule>::value &&
+              std::is_trivially_destructible<LegacyDatagramCapsule>::value,
+          "All capsule structs must have these properties");
+      legacy_datagram_capsule_ = LegacyDatagramCapsule();
+      break;
+    case CapsuleType::DATAGRAM_WITH_CONTEXT:
+      static_assert(
+          std::is_standard_layout<DatagramWithContextCapsule>::value &&
+              std::is_trivially_destructible<DatagramWithContextCapsule>::value,
+          "All capsule structs must have these properties");
+      datagram_with_context_capsule_ = DatagramWithContextCapsule();
+      break;
+    case CapsuleType::DATAGRAM_WITHOUT_CONTEXT:
+      static_assert(
+          std::is_standard_layout<DatagramWithoutContextCapsule>::value &&
+              std::is_trivially_destructible<
+                  DatagramWithoutContextCapsule>::value,
+          "All capsule structs must have these properties");
+      datagram_without_context_capsule_ = DatagramWithoutContextCapsule();
       break;
     case CapsuleType::REGISTER_DATAGRAM_CONTEXT:
       static_assert(
@@ -126,11 +146,32 @@
 }
 
 // static
-Capsule Capsule::Datagram(absl::optional<QuicDatagramContextId> context_id,
-                          absl::string_view http_datagram_payload) {
-  Capsule capsule(CapsuleType::DATAGRAM);
-  capsule.datagram_capsule().context_id = context_id;
-  capsule.datagram_capsule().http_datagram_payload = http_datagram_payload;
+Capsule Capsule::LegacyDatagram(
+    absl::optional<QuicDatagramContextId> context_id,
+    absl::string_view http_datagram_payload) {
+  Capsule capsule(CapsuleType::LEGACY_DATAGRAM);
+  capsule.legacy_datagram_capsule().context_id = context_id;
+  capsule.legacy_datagram_capsule().http_datagram_payload =
+      http_datagram_payload;
+  return capsule;
+}
+
+// static
+Capsule Capsule::DatagramWithContext(QuicDatagramContextId context_id,
+                                     absl::string_view http_datagram_payload) {
+  Capsule capsule(CapsuleType::DATAGRAM_WITH_CONTEXT);
+  capsule.datagram_with_context_capsule().context_id = context_id;
+  capsule.datagram_with_context_capsule().http_datagram_payload =
+      http_datagram_payload;
+  return capsule;
+}
+
+// static
+Capsule Capsule::DatagramWithoutContext(
+    absl::string_view http_datagram_payload) {
+  Capsule capsule(CapsuleType::DATAGRAM_WITHOUT_CONTEXT);
+  capsule.datagram_without_context_capsule().http_datagram_payload =
+      http_datagram_payload;
   return capsule;
 }
 
@@ -187,8 +228,15 @@
 Capsule& Capsule::operator=(const Capsule& other) {
   capsule_type_ = other.capsule_type_;
   switch (capsule_type_) {
-    case CapsuleType::DATAGRAM:
-      datagram_capsule_ = other.datagram_capsule_;
+    case CapsuleType::LEGACY_DATAGRAM:
+      legacy_datagram_capsule_ = other.legacy_datagram_capsule_;
+      break;
+    case CapsuleType::DATAGRAM_WITH_CONTEXT:
+      datagram_with_context_capsule_ = other.datagram_with_context_capsule_;
+      break;
+    case CapsuleType::DATAGRAM_WITHOUT_CONTEXT:
+      datagram_without_context_capsule_ =
+          other.datagram_without_context_capsule_;
       break;
     case CapsuleType::REGISTER_DATAGRAM_CONTEXT:
       register_datagram_context_capsule_ =
@@ -221,11 +269,19 @@
     return false;
   }
   switch (capsule_type_) {
-    case CapsuleType::DATAGRAM:
-      return datagram_capsule_.context_id ==
-                 other.datagram_capsule_.context_id &&
-             datagram_capsule_.http_datagram_payload ==
-                 other.datagram_capsule_.http_datagram_payload;
+    case CapsuleType::LEGACY_DATAGRAM:
+      return legacy_datagram_capsule_.context_id ==
+                 other.legacy_datagram_capsule_.context_id &&
+             legacy_datagram_capsule_.http_datagram_payload ==
+                 other.legacy_datagram_capsule_.http_datagram_payload;
+    case CapsuleType::DATAGRAM_WITH_CONTEXT:
+      return datagram_with_context_capsule_.context_id ==
+                 other.datagram_with_context_capsule_.context_id &&
+             datagram_with_context_capsule_.http_datagram_payload ==
+                 other.datagram_with_context_capsule_.http_datagram_payload;
+    case CapsuleType::DATAGRAM_WITHOUT_CONTEXT:
+      return datagram_without_context_capsule_.http_datagram_payload ==
+             other.datagram_without_context_capsule_.http_datagram_payload;
     case CapsuleType::REGISTER_DATAGRAM_CONTEXT:
       return register_datagram_context_capsule_.context_id ==
                  other.register_datagram_context_capsule_.context_id &&
@@ -260,13 +316,28 @@
 std::string Capsule::ToString() const {
   std::string rv = CapsuleTypeToString(capsule_type_);
   switch (capsule_type_) {
-    case CapsuleType::DATAGRAM:
-      if (datagram_capsule_.context_id.has_value()) {
-        absl::StrAppend(&rv, "(", datagram_capsule_.context_id.value(), ")");
+    case CapsuleType::LEGACY_DATAGRAM:
+      if (legacy_datagram_capsule_.context_id.has_value()) {
+        absl::StrAppend(&rv, "(", legacy_datagram_capsule_.context_id.value(),
+                        ")");
       }
+      absl::StrAppend(&rv, "[",
+                      absl::BytesToHexString(
+                          legacy_datagram_capsule_.http_datagram_payload),
+                      "]");
+      break;
+    case CapsuleType::DATAGRAM_WITH_CONTEXT:
+      absl::StrAppend(&rv, "(", datagram_with_context_capsule_.context_id, ")[",
+                      absl::BytesToHexString(
+                          datagram_with_context_capsule_.http_datagram_payload),
+                      "]");
+      break;
+    case CapsuleType::DATAGRAM_WITHOUT_CONTEXT:
       absl::StrAppend(
           &rv, "[",
-          absl::BytesToHexString(datagram_capsule_.http_datagram_payload), "]");
+          absl::BytesToHexString(
+              datagram_without_context_capsule_.http_datagram_payload),
+          "]");
       break;
     case CapsuleType::REGISTER_DATAGRAM_CONTEXT:
       absl::StrAppend(
@@ -327,14 +398,25 @@
       static_cast<uint64_t>(capsule.capsule_type()));
   QuicByteCount capsule_data_length;
   switch (capsule.capsule_type()) {
-    case CapsuleType::DATAGRAM:
+    case CapsuleType::LEGACY_DATAGRAM:
       capsule_data_length =
-          capsule.datagram_capsule().http_datagram_payload.length();
-      if (capsule.datagram_capsule().context_id.has_value()) {
+          capsule.legacy_datagram_capsule().http_datagram_payload.length();
+      if (capsule.legacy_datagram_capsule().context_id.has_value()) {
         capsule_data_length += QuicDataWriter::GetVarInt62Len(
-            capsule.datagram_capsule().context_id.value());
+            capsule.legacy_datagram_capsule().context_id.value());
       }
       break;
+    case CapsuleType::DATAGRAM_WITH_CONTEXT:
+      capsule_data_length =
+          QuicDataWriter::GetVarInt62Len(
+              capsule.datagram_with_context_capsule().context_id) +
+          capsule.datagram_with_context_capsule()
+              .http_datagram_payload.length();
+      break;
+    case CapsuleType::DATAGRAM_WITHOUT_CONTEXT:
+      capsule_data_length = capsule.datagram_without_context_capsule()
+                                .http_datagram_payload.length();
+      break;
     case CapsuleType::REGISTER_DATAGRAM_CONTEXT:
       capsule_data_length =
           QuicDataWriter::GetVarInt62Len(
@@ -383,19 +465,41 @@
     return QuicBuffer();
   }
   switch (capsule.capsule_type()) {
-    case CapsuleType::DATAGRAM:
-      if (capsule.datagram_capsule().context_id.has_value()) {
+    case CapsuleType::LEGACY_DATAGRAM:
+      if (capsule.legacy_datagram_capsule().context_id.has_value()) {
         if (!writer.WriteVarInt62(
-                capsule.datagram_capsule().context_id.value())) {
+                capsule.legacy_datagram_capsule().context_id.value())) {
           QUIC_BUG(datagram capsule context ID write fail)
-              << "Failed to write DATAGRAM CAPSULE context ID";
+              << "Failed to write LEGACY_DATAGRAM CAPSULE context ID";
           return QuicBuffer();
         }
       }
       if (!writer.WriteStringPiece(
-              capsule.datagram_capsule().http_datagram_payload)) {
+              capsule.legacy_datagram_capsule().http_datagram_payload)) {
         QUIC_BUG(datagram capsule payload write fail)
-            << "Failed to write DATAGRAM CAPSULE payload";
+            << "Failed to write LEGACY_DATAGRAM CAPSULE payload";
+        return QuicBuffer();
+      }
+      break;
+    case CapsuleType::DATAGRAM_WITH_CONTEXT:
+      if (!writer.WriteVarInt62(
+              capsule.datagram_with_context_capsule().context_id)) {
+        QUIC_BUG(datagram capsule context ID write fail)
+            << "Failed to write DATAGRAM_WITH_CONTEXT CAPSULE context ID";
+        return QuicBuffer();
+      }
+      if (!writer.WriteStringPiece(
+              capsule.datagram_with_context_capsule().http_datagram_payload)) {
+        QUIC_BUG(datagram capsule payload write fail)
+            << "Failed to write DATAGRAM_WITH_CONTEXT CAPSULE payload";
+        return QuicBuffer();
+      }
+      break;
+    case CapsuleType::DATAGRAM_WITHOUT_CONTEXT:
+      if (!writer.WriteStringPiece(capsule.datagram_without_context_capsule()
+                                       .http_datagram_payload)) {
+        QUIC_BUG(datagram capsule payload write fail)
+            << "Failed to write DATAGRAM_WITHOUT_CONTEXT CAPSULE payload";
         return QuicBuffer();
       }
       break;
@@ -533,16 +637,32 @@
   QuicDataReader capsule_data_reader(capsule_data);
   Capsule capsule(static_cast<CapsuleType>(capsule_type64));
   switch (capsule.capsule_type()) {
-    case CapsuleType::DATAGRAM:
+    case CapsuleType::LEGACY_DATAGRAM:
       if (datagram_context_id_present_) {
         uint64_t context_id;
         if (!capsule_data_reader.ReadVarInt62(&context_id)) {
-          ReportParseFailure("Unable to parse capsule DATAGRAM context ID");
+          ReportParseFailure(
+              "Unable to parse capsule LEGACY_DATAGRAM context ID");
           return 0;
         }
-        capsule.datagram_capsule().context_id = context_id;
+        capsule.legacy_datagram_capsule().context_id = context_id;
       }
-      capsule.datagram_capsule().http_datagram_payload =
+      capsule.legacy_datagram_capsule().http_datagram_payload =
+          capsule_data_reader.ReadRemainingPayload();
+      break;
+    case CapsuleType::DATAGRAM_WITH_CONTEXT:
+      uint64_t context_id;
+      if (!capsule_data_reader.ReadVarInt62(&context_id)) {
+        ReportParseFailure(
+            "Unable to parse capsule DATAGRAM_WITH_CONTEXT context ID");
+        return 0;
+      }
+      capsule.datagram_with_context_capsule().context_id = context_id;
+      capsule.datagram_with_context_capsule().http_datagram_payload =
+          capsule_data_reader.ReadRemainingPayload();
+      break;
+    case CapsuleType::DATAGRAM_WITHOUT_CONTEXT:
+      capsule.datagram_without_context_capsule().http_datagram_payload =
           capsule_data_reader.ReadRemainingPayload();
       break;
     case CapsuleType::REGISTER_DATAGRAM_CONTEXT:
diff --git a/quic/core/http/capsule.h b/quic/core/http/capsule.h
index 5198b46..42190bb 100644
--- a/quic/core/http/capsule.h
+++ b/quic/core/http/capsule.h
@@ -20,10 +20,12 @@
 
 enum class CapsuleType : uint64_t {
   // Casing in this enum matches the IETF specification.
-  DATAGRAM = 0xff37a0,
+  LEGACY_DATAGRAM = 0xff37a0,  // draft-ietf-masque-h3-datagram-04
   REGISTER_DATAGRAM_CONTEXT = 0xff37a1,
   REGISTER_DATAGRAM_NO_CONTEXT = 0xff37a2,
   CLOSE_DATAGRAM_CONTEXT = 0xff37a3,
+  DATAGRAM_WITH_CONTEXT = 0xff37a4,
+  DATAGRAM_WITHOUT_CONTEXT = 0xff37a5,
   CLOSE_WEBTRANSPORT_SESSION = 0x2843,
 };
 
@@ -55,10 +57,17 @@
 QUIC_EXPORT_PRIVATE std::ostream& operator<<(
     std::ostream& os, const ContextCloseCode& context_close_code);
 
-struct QUIC_EXPORT_PRIVATE DatagramCapsule {
+struct QUIC_EXPORT_PRIVATE LegacyDatagramCapsule {
   absl::optional<QuicDatagramContextId> context_id;
   absl::string_view http_datagram_payload;
 };
+struct QUIC_EXPORT_PRIVATE DatagramWithContextCapsule {
+  QuicDatagramContextId context_id;
+  absl::string_view http_datagram_payload;
+};
+struct QUIC_EXPORT_PRIVATE DatagramWithoutContextCapsule {
+  absl::string_view http_datagram_payload;
+};
 struct QUIC_EXPORT_PRIVATE RegisterDatagramContextCapsule {
   QuicDatagramContextId context_id;
   DatagramFormatType format_type;
@@ -85,9 +94,14 @@
 // or perform its own deep copy.
 class QUIC_EXPORT_PRIVATE Capsule {
  public:
-  static Capsule Datagram(
+  static Capsule LegacyDatagram(
       absl::optional<QuicDatagramContextId> context_id = absl::nullopt,
       absl::string_view http_datagram_payload = absl::string_view());
+  static Capsule DatagramWithContext(
+      QuicDatagramContextId context_id,
+      absl::string_view http_datagram_payload = absl::string_view());
+  static Capsule DatagramWithoutContext(
+      absl::string_view http_datagram_payload = absl::string_view());
   static Capsule RegisterDatagramContext(
       QuicDatagramContextId context_id, DatagramFormatType format_type,
       absl::string_view format_additional_data = absl::string_view());
@@ -116,13 +130,30 @@
                                                       const Capsule& capsule);
 
   CapsuleType capsule_type() const { return capsule_type_; }
-  DatagramCapsule& datagram_capsule() {
-    QUICHE_DCHECK_EQ(capsule_type_, CapsuleType::DATAGRAM);
-    return datagram_capsule_;
+  LegacyDatagramCapsule& legacy_datagram_capsule() {
+    QUICHE_DCHECK_EQ(capsule_type_, CapsuleType::LEGACY_DATAGRAM);
+    return legacy_datagram_capsule_;
   }
-  const DatagramCapsule& datagram_capsule() const {
-    QUICHE_DCHECK_EQ(capsule_type_, CapsuleType::DATAGRAM);
-    return datagram_capsule_;
+  const LegacyDatagramCapsule& legacy_datagram_capsule() const {
+    QUICHE_DCHECK_EQ(capsule_type_, CapsuleType::LEGACY_DATAGRAM);
+    return legacy_datagram_capsule_;
+  }
+  DatagramWithContextCapsule& datagram_with_context_capsule() {
+    QUICHE_DCHECK_EQ(capsule_type_, CapsuleType::DATAGRAM_WITH_CONTEXT);
+    return datagram_with_context_capsule_;
+  }
+  const DatagramWithContextCapsule& datagram_with_context_capsule() const {
+    QUICHE_DCHECK_EQ(capsule_type_, CapsuleType::DATAGRAM_WITH_CONTEXT);
+    return datagram_with_context_capsule_;
+  }
+  DatagramWithoutContextCapsule& datagram_without_context_capsule() {
+    QUICHE_DCHECK_EQ(capsule_type_, CapsuleType::DATAGRAM_WITHOUT_CONTEXT);
+    return datagram_without_context_capsule_;
+  }
+  const DatagramWithoutContextCapsule& datagram_without_context_capsule()
+      const {
+    QUICHE_DCHECK_EQ(capsule_type_, CapsuleType::DATAGRAM_WITHOUT_CONTEXT);
+    return datagram_without_context_capsule_;
   }
   RegisterDatagramContextCapsule& register_datagram_context_capsule() {
     QUICHE_DCHECK_EQ(capsule_type_, CapsuleType::REGISTER_DATAGRAM_CONTEXT);
@@ -160,7 +191,9 @@
     return close_web_transport_session_capsule_;
   }
   absl::string_view& unknown_capsule_data() {
-    QUICHE_DCHECK(capsule_type_ != CapsuleType::DATAGRAM &&
+    QUICHE_DCHECK(capsule_type_ != CapsuleType::LEGACY_DATAGRAM &&
+                  capsule_type_ != CapsuleType::DATAGRAM_WITH_CONTEXT &&
+                  capsule_type_ != CapsuleType::DATAGRAM_WITHOUT_CONTEXT &&
                   capsule_type_ != CapsuleType::REGISTER_DATAGRAM_CONTEXT &&
                   capsule_type_ != CapsuleType::REGISTER_DATAGRAM_NO_CONTEXT &&
                   capsule_type_ != CapsuleType::CLOSE_DATAGRAM_CONTEXT &&
@@ -169,7 +202,9 @@
     return unknown_capsule_data_;
   }
   const absl::string_view& unknown_capsule_data() const {
-    QUICHE_DCHECK(capsule_type_ != CapsuleType::DATAGRAM &&
+    QUICHE_DCHECK(capsule_type_ != CapsuleType::LEGACY_DATAGRAM &&
+                  capsule_type_ != CapsuleType::DATAGRAM_WITH_CONTEXT &&
+                  capsule_type_ != CapsuleType::DATAGRAM_WITHOUT_CONTEXT &&
                   capsule_type_ != CapsuleType::REGISTER_DATAGRAM_CONTEXT &&
                   capsule_type_ != CapsuleType::REGISTER_DATAGRAM_NO_CONTEXT &&
                   capsule_type_ != CapsuleType::CLOSE_DATAGRAM_CONTEXT &&
@@ -181,7 +216,9 @@
  private:
   CapsuleType capsule_type_;
   union {
-    DatagramCapsule datagram_capsule_;
+    LegacyDatagramCapsule legacy_datagram_capsule_;
+    DatagramWithContextCapsule datagram_with_context_capsule_;
+    DatagramWithoutContextCapsule datagram_without_context_capsule_;
     RegisterDatagramContextCapsule register_datagram_context_capsule_;
     RegisterDatagramNoContextCapsule register_datagram_no_context_capsule_;
     CloseDatagramContextCapsule close_datagram_context_capsule_;
diff --git a/quic/core/http/capsule_test.cc b/quic/core/http/capsule_test.cc
index 66cb3d0..2dfdef4 100644
--- a/quic/core/http/capsule_test.cc
+++ b/quic/core/http/capsule_test.cc
@@ -72,15 +72,15 @@
   CapsuleParser capsule_parser_;
 };
 
-TEST_F(CapsuleTest, DatagramCapsule) {
+TEST_F(CapsuleTest, LegacyDatagramCapsule) {
   std::string capsule_fragment = absl::HexStringToBytes(
-      "80ff37a0"          // DATAGRAM capsule type
+      "80ff37a0"          // LEGACY_DATAGRAM capsule type
       "08"                // capsule length
       "a1a2a3a4a5a6a7a8"  // HTTP Datagram payload
   );
   std::string datagram_payload = absl::HexStringToBytes("a1a2a3a4a5a6a7a8");
   Capsule expected_capsule =
-      Capsule::Datagram(/*context_id=*/absl::nullopt, datagram_payload);
+      Capsule::LegacyDatagram(/*context_id=*/absl::nullopt, datagram_payload);
   {
     EXPECT_CALL(visitor_, OnCapsule(expected_capsule));
     ASSERT_TRUE(capsule_parser_.IngestCapsuleFragment(capsule_fragment));
@@ -89,9 +89,9 @@
   TestSerialization(expected_capsule, capsule_fragment);
 }
 
-TEST_F(CapsuleTest, DatagramCapsuleWithContext) {
+TEST_F(CapsuleTest, LegacyDatagramCapsuleWithContext) {
   std::string capsule_fragment = absl::HexStringToBytes(
-      "80ff37a0"          // DATAGRAM capsule type
+      "80ff37a0"          // LEGACY_DATAGRAM capsule type
       "09"                // capsule length
       "04"                // context ID
       "a1a2a3a4a5a6a7a8"  // HTTP Datagram payload
@@ -99,7 +99,41 @@
   capsule_parser_.set_datagram_context_id_present(true);
   std::string datagram_payload = absl::HexStringToBytes("a1a2a3a4a5a6a7a8");
   Capsule expected_capsule =
-      Capsule::Datagram(/*context_id=*/4, datagram_payload);
+      Capsule::LegacyDatagram(/*context_id=*/4, datagram_payload);
+  {
+    EXPECT_CALL(visitor_, OnCapsule(expected_capsule));
+    ASSERT_TRUE(capsule_parser_.IngestCapsuleFragment(capsule_fragment));
+  }
+  ValidateParserIsEmpty();
+  TestSerialization(expected_capsule, capsule_fragment);
+}
+
+TEST_F(CapsuleTest, DatagramWithoutContextCapsule) {
+  std::string capsule_fragment = absl::HexStringToBytes(
+      "80ff37a5"          // DATAGRAM_WITHOUT_CONTEXT capsule type
+      "08"                // capsule length
+      "a1a2a3a4a5a6a7a8"  // HTTP Datagram payload
+  );
+  std::string datagram_payload = absl::HexStringToBytes("a1a2a3a4a5a6a7a8");
+  Capsule expected_capsule = Capsule::DatagramWithoutContext(datagram_payload);
+  {
+    EXPECT_CALL(visitor_, OnCapsule(expected_capsule));
+    ASSERT_TRUE(capsule_parser_.IngestCapsuleFragment(capsule_fragment));
+  }
+  ValidateParserIsEmpty();
+  TestSerialization(expected_capsule, capsule_fragment);
+}
+
+TEST_F(CapsuleTest, DatagramWithContextCapsule) {
+  std::string capsule_fragment = absl::HexStringToBytes(
+      "80ff37a4"          // DATAGRAM_WITH_CONTEXT capsule type
+      "09"                // capsule length
+      "04"                // context ID
+      "a1a2a3a4a5a6a7a8"  // HTTP Datagram payload
+  );
+  std::string datagram_payload = absl::HexStringToBytes("a1a2a3a4a5a6a7a8");
+  Capsule expected_capsule =
+      Capsule::DatagramWithContext(/*context_id=*/4, datagram_payload);
   {
     EXPECT_CALL(visitor_, OnCapsule(expected_capsule));
     ASSERT_TRUE(capsule_parser_.IngestCapsuleFragment(capsule_fragment));
@@ -199,21 +233,21 @@
   TestSerialization(expected_capsule, capsule_fragment);
 }
 
-TEST_F(CapsuleTest, TwoDatagramCapsules) {
+TEST_F(CapsuleTest, TwoCapsules) {
   std::string capsule_fragment = absl::HexStringToBytes(
-      "80ff37a0"          // DATAGRAM capsule type
+      "80ff37a5"          // DATAGRAM_WITHOUT_CONTEXT capsule type
       "08"                // capsule length
       "a1a2a3a4a5a6a7a8"  // HTTP Datagram payload
-      "80ff37a0"          // DATAGRAM capsule type
+      "80ff37a5"          // DATAGRAM_WITHOUT_CONTEXT capsule type
       "08"                // capsule length
       "b1b2b3b4b5b6b7b8"  // HTTP Datagram payload
   );
   std::string datagram_payload1 = absl::HexStringToBytes("a1a2a3a4a5a6a7a8");
   std::string datagram_payload2 = absl::HexStringToBytes("b1b2b3b4b5b6b7b8");
   Capsule expected_capsule1 =
-      Capsule::Datagram(/*context_id=*/absl::nullopt, datagram_payload1);
+      Capsule::DatagramWithoutContext(datagram_payload1);
   Capsule expected_capsule2 =
-      Capsule::Datagram(/*context_id=*/absl::nullopt, datagram_payload2);
+      Capsule::DatagramWithoutContext(datagram_payload2);
   {
     InSequence s;
     EXPECT_CALL(visitor_, OnCapsule(expected_capsule1));
@@ -223,15 +257,15 @@
   ValidateParserIsEmpty();
 }
 
-TEST_F(CapsuleTest, TwoDatagramCapsulesPartialReads) {
+TEST_F(CapsuleTest, TwoCapsulesPartialReads) {
   std::string capsule_fragment1 = absl::HexStringToBytes(
-      "80ff37a0"  // first capsule DATAGRAM capsule type
+      "80ff37a5"  // first capsule DATAGRAM_WITHOUT_CONTEXT capsule type
       "08"        // frist capsule length
       "a1a2a3a4"  // first half of HTTP Datagram payload of first capsule
   );
   std::string capsule_fragment2 = absl::HexStringToBytes(
       "a5a6a7a8"  // second half of HTTP Datagram payload 1
-      "80ff37a0"  // second capsule DATAGRAM capsule type
+      "80ff37a5"  // second capsule DATAGRAM_WITHOUT_CONTEXT capsule type
   );
   std::string capsule_fragment3 = absl::HexStringToBytes(
       "08"                // second capsule length
@@ -241,9 +275,9 @@
   std::string datagram_payload1 = absl::HexStringToBytes("a1a2a3a4a5a6a7a8");
   std::string datagram_payload2 = absl::HexStringToBytes("b1b2b3b4b5b6b7b8");
   Capsule expected_capsule1 =
-      Capsule::Datagram(/*context_id=*/absl::nullopt, datagram_payload1);
+      Capsule::DatagramWithoutContext(datagram_payload1);
   Capsule expected_capsule2 =
-      Capsule::Datagram(/*context_id=*/absl::nullopt, datagram_payload2);
+      Capsule::DatagramWithoutContext(datagram_payload2);
   {
     InSequence s;
     EXPECT_CALL(visitor_, OnCapsule(expected_capsule1));
@@ -255,21 +289,21 @@
   ValidateParserIsEmpty();
 }
 
-TEST_F(CapsuleTest, TwoDatagramCapsulesOneByteAtATime) {
+TEST_F(CapsuleTest, TwoCapsulesOneByteAtATime) {
   std::string capsule_fragment = absl::HexStringToBytes(
-      "80ff37a0"          // DATAGRAM capsule type
+      "80ff37a5"          // DATAGRAM_WITHOUT_CONTEXT capsule type
       "08"                // capsule length
       "a1a2a3a4a5a6a7a8"  // HTTP Datagram payload
-      "80ff37a0"          // DATAGRAM capsule type
+      "80ff37a5"          // DATAGRAM_WITHOUT_CONTEXT capsule type
       "08"                // capsule length
       "b1b2b3b4b5b6b7b8"  // HTTP Datagram payload
   );
   std::string datagram_payload1 = absl::HexStringToBytes("a1a2a3a4a5a6a7a8");
   std::string datagram_payload2 = absl::HexStringToBytes("b1b2b3b4b5b6b7b8");
   Capsule expected_capsule1 =
-      Capsule::Datagram(/*context_id=*/absl::nullopt, datagram_payload1);
+      Capsule::DatagramWithoutContext(datagram_payload1);
   Capsule expected_capsule2 =
-      Capsule::Datagram(/*context_id=*/absl::nullopt, datagram_payload2);
+      Capsule::DatagramWithoutContext(datagram_payload2);
   for (size_t i = 0; i < capsule_fragment.size(); i++) {
     if (i < capsule_fragment.size() / 2 - 1) {
       EXPECT_CALL(visitor_, OnCapsule(_)).Times(0);
@@ -297,7 +331,7 @@
 
 TEST_F(CapsuleTest, PartialCapsuleThenError) {
   std::string capsule_fragment = absl::HexStringToBytes(
-      "80ff37a0"  // DATAGRAM capsule type
+      "80ff37a5"  // DATAGRAM_WITHOUT_CONTEXT capsule type
       "08"        // capsule length
       "a1a2a3a4"  // first half of HTTP Datagram payload
   );
diff --git a/quic/core/http/end_to_end_test.cc b/quic/core/http/end_to_end_test.cc
index 3b073af..5f8675d 100644
--- a/quic/core/http/end_to_end_test.cc
+++ b/quic/core/http/end_to_end_test.cc
@@ -54,6 +54,7 @@
 #include "quic/test_tools/quic_server_peer.h"
 #include "quic/test_tools/quic_session_peer.h"
 #include "quic/test_tools/quic_spdy_session_peer.h"
+#include "quic/test_tools/quic_spdy_stream_peer.h"
 #include "quic/test_tools/quic_stream_id_manager_peer.h"
 #include "quic/test_tools/quic_stream_peer.h"
 #include "quic/test_tools/quic_stream_sequencer_peer.h"
@@ -229,6 +230,7 @@
     client->UseClientConnectionIdLength(override_client_connection_id_length_);
     client->client()->set_connection_debug_visitor(connection_debug_visitor_);
     client->client()->set_enable_web_transport(enable_web_transport_);
+    client->client()->set_use_datagram_contexts(use_datagram_contexts_);
     client->Connect();
     return client;
   }
@@ -368,6 +370,9 @@
     if (enable_web_transport_) {
       memory_cache_backend_.set_enable_webtransport(true);
     }
+    if (use_datagram_contexts_) {
+      memory_cache_backend_.set_use_datagram_contexts(true);
+    }
 
     QuicTagVector copt;
     server_config_.SetConnectionOptionsToSend(copt);
@@ -689,8 +694,9 @@
     return WaitForFooResponseAndCheckIt(client_.get());
   }
 
-  WebTransportHttp3* CreateWebTransportSession(const std::string& path,
-                                               bool wait_for_server_response) {
+  WebTransportHttp3* CreateWebTransportSession(
+      const std::string& path, bool wait_for_server_response,
+      QuicSpdyStream** connect_stream_out = nullptr) {
     // Wait until we receive the settings from the server indicating
     // WebTransport support.
     client_->WaitUntil(
@@ -722,6 +728,9 @@
                          [stream]() { return stream->headers_decompressed(); });
       EXPECT_TRUE(session->ready());
     }
+    if (connect_stream_out != nullptr) {
+      *connect_stream_out = stream;
+    }
     return session;
   }
 
@@ -829,6 +838,7 @@
   int override_client_connection_id_length_ = -1;
   uint8_t expected_server_connection_id_length_;
   bool enable_web_transport_ = false;
+  bool use_datagram_contexts_ = false;
   std::vector<std::string> received_webtransport_unidirectional_streams_;
 };
 
@@ -6347,11 +6357,7 @@
 
   SimpleBufferAllocator allocator;
   for (int i = 0; i < 10; i++) {
-    absl::string_view datagram = "test";
-    auto buffer = MakeUniqueBuffer(&allocator, datagram.size());
-    memcpy(buffer.get(), datagram.data(), datagram.size());
-    QuicMemSlice slice(std::move(buffer), datagram.size());
-    session->SendOrQueueDatagram(std::move(slice));
+    session->SendOrQueueDatagram(MemSliceFromString("test"));
   }
 
   int received = 0;
@@ -6362,6 +6368,37 @@
   EXPECT_GT(received, 0);
 }
 
+TEST_P(EndToEndTest, WebTransportDatagramsWithContexts) {
+  enable_web_transport_ = true;
+  use_datagram_contexts_ = true;
+  SetPacketLossPercentage(30);
+  ASSERT_TRUE(Initialize());
+
+  if (!version_.UsesHttp3()) {
+    return;
+  }
+
+  QuicSpdyStream* connect_stream = nullptr;
+  WebTransportHttp3* session = CreateWebTransportSession(
+      "/echo", /*wait_for_server_response=*/true, &connect_stream);
+  ASSERT_TRUE(session != nullptr);
+  ASSERT_TRUE(connect_stream != nullptr);
+  NiceMock<MockClientVisitor>& visitor = SetupWebTransportVisitor(session);
+
+  SimpleBufferAllocator allocator;
+  for (int i = 0; i < 10; i++) {
+    session->SendOrQueueDatagram(MemSliceFromString("test"));
+  }
+
+  int received = 0;
+  EXPECT_CALL(visitor, OnDatagramReceived(_)).WillRepeatedly([&received]() {
+    received++;
+  });
+  client_->WaitUntil(5000, [&received]() { return received > 0; });
+  EXPECT_GT(received, 0);
+  EXPECT_TRUE(QuicSpdyStreamPeer::use_datagram_contexts(connect_stream));
+}
+
 TEST_P(EndToEndTest, WebTransportSessionClose) {
   enable_web_transport_ = true;
   ASSERT_TRUE(Initialize());
diff --git a/quic/core/http/quic_spdy_session.cc b/quic/core/http/quic_spdy_session.cc
index 929e369..af54b06 100644
--- a/quic/core/http/quic_spdy_session.cc
+++ b/quic/core/http/quic_spdy_session.cc
@@ -875,6 +875,8 @@
 
 bool QuicSpdySession::ShouldNegotiateWebTransport() { return false; }
 
+bool QuicSpdySession::ShouldNegotiateDatagramContexts() { return false; }
+
 bool QuicSpdySession::WillNegotiateWebTransport() {
   return LocalHttpDatagramSupport() != HttpDatagramSupport::kNone &&
          version().UsesHttp3() && ShouldNegotiateWebTransport();
diff --git a/quic/core/http/quic_spdy_session.h b/quic/core/http/quic_spdy_session.h
index b84f5a1..9f95075 100644
--- a/quic/core/http/quic_spdy_session.h
+++ b/quic/core/http/quic_spdy_session.h
@@ -467,6 +467,10 @@
 
   QuicSpdyStream* GetOrCreateSpdyDataStream(const QuicStreamId stream_id);
 
+  // Indicates whether we will try to negotiate datagram contexts on newly
+  // created WebTransport sessions over HTTP/3.
+  virtual bool ShouldNegotiateDatagramContexts();
+
  protected:
   // Override CreateIncomingStream(), CreateOutgoingBidirectionalStream() and
   // CreateOutgoingUnidirectionalStream() with QuicSpdyStream return type to
diff --git a/quic/core/http/quic_spdy_stream.cc b/quic/core/http/quic_spdy_stream.cc
index 284a180..55265de 100644
--- a/quic/core/http/quic_spdy_stream.cc
+++ b/quic/core/http/quic_spdy_stream.cc
@@ -248,6 +248,13 @@
 
 QuicSpdyStream::~QuicSpdyStream() {}
 
+bool QuicSpdyStream::ShouldUseDatagramContexts() const {
+  return spdy_session_->SupportsH3Datagram() &&
+         spdy_session_->http_datagram_support() !=
+             HttpDatagramSupport::kDraft00 &&
+         use_datagram_contexts_;
+}
+
 size_t QuicSpdyStream::WriteHeaders(
     SpdyHeaderBlock header_block, bool fin,
     QuicReferenceCountedPointer<QuicAckListenerInterface> ack_listener) {
@@ -275,6 +282,12 @@
 
   MaybeProcessSentWebTransportHeaders(header_block);
 
+  if (ShouldUseDatagramContexts()) {
+    // RegisterHttp3DatagramRegistrationVisitor caller wishes to use contexts,
+    // inform the peer.
+    header_block["sec-use-datagram-contexts"] = "?1";
+  }
+
   size_t bytes_written =
       WriteHeadersImpl(std::move(header_block), fin, std::move(ack_listener));
   if (!VersionUsesHttp3(transport_version()) && fin) {
@@ -621,6 +634,19 @@
   if (!GetQuicReloadableFlag(quic_verify_request_headers) ||
       !header_too_large) {
     MaybeProcessReceivedWebTransportHeaders();
+    if (ShouldUseDatagramContexts()) {
+      bool peer_wishes_to_use_datagram_contexts = false;
+      for (const auto& header : header_list_) {
+        if (header.first == "sec-use-datagram-contexts" &&
+            header.second == "?1") {
+          peer_wishes_to_use_datagram_contexts = true;
+          break;
+        }
+      }
+      if (!peer_wishes_to_use_datagram_contexts) {
+        use_datagram_contexts_ = false;
+      }
+    }
   }
 
   if (VersionUsesHttp3(transport_version())) {
@@ -1255,8 +1281,9 @@
     RegisterHttp3DatagramFlowId(*flow_id);
   }
 
-  web_transport_ =
-      std::make_unique<WebTransportHttp3>(spdy_session_, this, id());
+  web_transport_ = std::make_unique<WebTransportHttp3>(
+      spdy_session_, this, id(),
+      spdy_session_->ShouldNegotiateDatagramContexts());
 
   if (spdy_session_->http_datagram_support() != HttpDatagramSupport::kDraft00) {
     return;
@@ -1293,8 +1320,9 @@
     headers["datagram-flow-id"] = absl::StrCat(id());
   }
 
-  web_transport_ =
-      std::make_unique<WebTransportHttp3>(spdy_session_, this, id());
+  web_transport_ = std::make_unique<WebTransportHttp3>(
+      spdy_session_, this, id(),
+      spdy_session_->ShouldNegotiateDatagramContexts());
 }
 
 void QuicSpdyStream::OnCanWriteNewData() {
@@ -1394,11 +1422,28 @@
     return false;
   }
   switch (capsule.capsule_type()) {
-    case CapsuleType::DATAGRAM: {
-      HandleReceivedDatagram(capsule.datagram_capsule().context_id,
-                             capsule.datagram_capsule().http_datagram_payload);
+    case CapsuleType::LEGACY_DATAGRAM: {
+      HandleReceivedDatagram(
+          capsule.legacy_datagram_capsule().context_id,
+          capsule.legacy_datagram_capsule().http_datagram_payload);
     } break;
-    case CapsuleType::REGISTER_DATAGRAM_CONTEXT:
+    case CapsuleType::DATAGRAM_WITH_CONTEXT: {
+      HandleReceivedDatagram(
+          capsule.datagram_with_context_capsule().context_id,
+          capsule.datagram_with_context_capsule().http_datagram_payload);
+    } break;
+    case CapsuleType::DATAGRAM_WITHOUT_CONTEXT: {
+      absl::optional<QuicDatagramContextId> context_id;
+      if (use_datagram_contexts_) {
+        // draft-ietf-masque-h3-datagram-05 encodes context ID 0 using
+        // DATAGRAM_WITHOUT_CONTEXT.
+        context_id = 0;
+      }
+      HandleReceivedDatagram(
+          context_id,
+          capsule.datagram_without_context_capsule().http_datagram_payload);
+    } break;
+    case CapsuleType::REGISTER_DATAGRAM_CONTEXT: {
       if (datagram_registration_visitor_ == nullptr) {
         QUIC_DLOG(ERROR) << ENDPOINT << "Received capsule " << capsule
                          << " without any registration visitor";
@@ -1408,20 +1453,26 @@
           id(), capsule.register_datagram_context_capsule().context_id,
           capsule.register_datagram_context_capsule().format_type,
           capsule.register_datagram_context_capsule().format_additional_data);
-      break;
-    case CapsuleType::REGISTER_DATAGRAM_NO_CONTEXT:
+    } break;
+    case CapsuleType::REGISTER_DATAGRAM_NO_CONTEXT: {
       if (datagram_registration_visitor_ == nullptr) {
         QUIC_DLOG(ERROR) << ENDPOINT << "Received capsule " << capsule
                          << " without any registration visitor";
         return false;
       }
+      absl::optional<QuicDatagramContextId> context_id;
+      if (use_datagram_contexts_) {
+        // draft-ietf-masque-h3-datagram-05 encodes context ID 0 using
+        // REGISTER_DATAGRAM_NO_CONTEXT.
+        context_id = 0;
+      }
       datagram_registration_visitor_->OnContextReceived(
-          id(), /*context_id=*/absl::nullopt,
+          id(), context_id,
           capsule.register_datagram_no_context_capsule().format_type,
           capsule.register_datagram_no_context_capsule()
               .format_additional_data);
-      break;
-    case CapsuleType::CLOSE_DATAGRAM_CONTEXT:
+    } break;
+    case CapsuleType::CLOSE_DATAGRAM_CONTEXT: {
       if (datagram_registration_visitor_ == nullptr) {
         QUIC_DLOG(ERROR) << ENDPOINT << "Received capsule " << capsule
                          << " without any registration visitor";
@@ -1431,8 +1482,8 @@
           id(), capsule.close_datagram_context_capsule().context_id,
           capsule.close_datagram_context_capsule().close_code,
           capsule.close_datagram_context_capsule().close_details);
-      break;
-    case CapsuleType::CLOSE_WEBTRANSPORT_SESSION:
+    } break;
+    case CapsuleType::CLOSE_WEBTRANSPORT_SESSION: {
       if (web_transport_ == nullptr) {
         QUIC_DLOG(ERROR) << ENDPOINT << "Received capsule " << capsule
                          << " for a non-WebTransport stream.";
@@ -1441,7 +1492,7 @@
       web_transport_->OnCloseReceived(
           capsule.close_web_transport_session_capsule().error_code,
           capsule.close_web_transport_session_capsule().error_message);
-      break;
+    } break;
   }
   return true;
 }
@@ -1485,7 +1536,7 @@
 }
 
 void QuicSpdyStream::RegisterHttp3DatagramRegistrationVisitor(
-    Http3DatagramRegistrationVisitor* visitor) {
+    Http3DatagramRegistrationVisitor* visitor, bool use_datagram_contexts) {
   if (visitor == nullptr) {
     QUIC_BUG(null datagram registration visitor)
         << ENDPOINT << "Null datagram registration visitor for" << id();
@@ -1496,7 +1547,10 @@
         << ENDPOINT << "Double datagram registration visitor for" << id();
     return;
   }
-  QUIC_DLOG(INFO) << ENDPOINT << "Registering datagram stream ID " << id();
+  use_datagram_contexts_ = use_datagram_contexts;
+  QUIC_DLOG(INFO) << ENDPOINT << "Registering datagram stream ID " << id()
+                  << " with" << (use_datagram_contexts_ ? "" : "out")
+                  << " contexts";
   datagram_registration_visitor_ = visitor;
   QUICHE_DCHECK(!capsule_parser_);
   capsule_parser_.reset(new CapsuleParser(this));
@@ -1589,8 +1643,15 @@
         QuicConnection::ScopedPacketFlusher flusher(
             spdy_session_->connection());
         WriteGreaseCapsule();
-        WriteCapsule(Capsule::RegisterDatagramContext(
-            context_id.value(), format_type, format_additional_data));
+        if (context_id.value() != 0) {
+          WriteCapsule(Capsule::RegisterDatagramContext(
+              context_id.value(), format_type, format_additional_data));
+        } else {
+          // draft-ietf-masque-h3-datagram-05 encodes context ID 0 using
+          // REGISTER_DATAGRAM_NO_CONTEXT.
+          WriteCapsule(Capsule::RegisterDatagramNoContext(
+              format_type, format_additional_data));
+        }
         WriteGreaseCapsule();
       }
     } else if (is_client) {
@@ -1676,9 +1737,13 @@
 }
 
 void QuicSpdyStream::OnDatagramReceived(QuicDataReader* reader) {
+  if (!headers_decompressed_) {
+    QUIC_DLOG(INFO) << "Dropping datagram received before headers on stream ID "
+                    << id();
+    return;
+  }
   absl::optional<QuicDatagramContextId> context_id;
-  const bool context_id_present = !datagram_context_visitors_.empty();
-  if (context_id_present) {
+  if (use_datagram_contexts_) {
     QuicDatagramContextId parsed_context_id;
     if (!reader->ReadVarInt62(&parsed_context_id)) {
       QUIC_DLOG(ERROR) << "Failed to parse context ID in received HTTP/3 "
diff --git a/quic/core/http/quic_spdy_stream.h b/quic/core/http/quic_spdy_stream.h
index cd9d99b..8ec56d5 100644
--- a/quic/core/http/quic_spdy_stream.h
+++ b/quic/core/http/quic_spdy_stream.h
@@ -296,7 +296,8 @@
   // UnregisterHttp3DatagramRegistrationVisitor. |visitor| must be valid until a
   // corresponding call to UnregisterHttp3DatagramRegistrationVisitor.
   void RegisterHttp3DatagramRegistrationVisitor(
-      Http3DatagramRegistrationVisitor* visitor);
+      Http3DatagramRegistrationVisitor* visitor,
+      bool use_datagram_contexts = false);
 
   // Unregisters for HTTP/3 datagram context registrations. Must not be called
   // unless previously registered.
@@ -427,6 +428,9 @@
   void HandleReceivedDatagram(absl::optional<QuicDatagramContextId> context_id,
                               absl::string_view payload);
 
+  // Whether datagram contexts should be used on this stream.
+  bool ShouldUseDatagramContexts() const;
+
   QuicSpdySession* spdy_session_;
 
   bool on_body_available_called_because_sequencer_is_closed_;
@@ -505,6 +509,7 @@
   QuicDatagramContextId datagram_next_available_context_id_;
   absl::flat_hash_map<QuicDatagramContextId, Http3DatagramVisitor*>
       datagram_context_visitors_;
+  bool use_datagram_contexts_ = false;
 };
 
 }  // namespace quic
diff --git a/quic/core/http/quic_spdy_stream_test.cc b/quic/core/http/quic_spdy_stream_test.cc
index b5905e7..50dc9ea 100644
--- a/quic/core/http/quic_spdy_stream_test.cc
+++ b/quic/core/http/quic_spdy_stream_test.cc
@@ -3208,6 +3208,10 @@
   session_->set_local_http_datagram_support(HttpDatagramSupport::kDraft00And04);
   QuicSpdySessionPeer::SetHttpDatagramSupport(session_.get(),
                                               HttpDatagramSupport::kDraft00);
+  headers_[":method"] = "CONNECT";
+  headers_[":protocol"] = "webtransport";
+  headers_["datagram-flow-id"] = absl::StrCat(stream_->id());
+  ProcessHeaders(false, headers_);
   session_->RegisterHttp3DatagramFlowId(stream_->id(), stream_->id());
   ::testing::NiceMock<MockHttp3DatagramRegistrationVisitor>
       h3_datagram_registration_visitor;
@@ -3259,6 +3263,9 @@
   session_->set_local_http_datagram_support(HttpDatagramSupport::kDraft04);
   QuicSpdySessionPeer::SetHttpDatagramSupport(session_.get(),
                                               HttpDatagramSupport::kDraft04);
+  headers_[":method"] = "CONNECT";
+  headers_[":protocol"] = "webtransport";
+  ProcessHeaders(false, headers_);
   ::testing::NiceMock<MockHttp3DatagramRegistrationVisitor>
       h3_datagram_registration_visitor;
   SavingHttp3DatagramVisitor h3_datagram_visitor;
@@ -3326,7 +3333,11 @@
     datagram[i] = i;
   }
   stream_->RegisterHttp3DatagramRegistrationVisitor(
-      &h3_datagram_registration_visitor);
+      &h3_datagram_registration_visitor, /*use_datagram_contexts=*/true);
+  headers_[":method"] = "CONNECT";
+  headers_[":protocol"] = "webtransport";
+  headers_["sec-use-datagram-contexts"] = "?1";
+  ProcessHeaders(false, headers_);
 
   // Expect us to send a REGISTER_DATAGRAM_CONTEXT capsule.
   EXPECT_CALL(*session_, WritevData(stream_->id(), _, _, _, _, _))
diff --git a/quic/core/http/web_transport_http3.cc b/quic/core/http/web_transport_http3.cc
index a54d210..128c44e 100644
--- a/quic/core/http/web_transport_http3.cc
+++ b/quic/core/http/web_transport_http3.cc
@@ -42,7 +42,8 @@
 
 WebTransportHttp3::WebTransportHttp3(QuicSpdySession* session,
                                      QuicSpdyStream* connect_stream,
-                                     WebTransportSessionId id)
+                                     WebTransportSessionId id,
+                                     bool attempt_to_use_datagram_contexts)
     : session_(session),
       connect_stream_(connect_stream),
       id_(id),
@@ -50,10 +51,14 @@
   QUICHE_DCHECK(session_->SupportsWebTransport());
   QUICHE_DCHECK(IsValidWebTransportSessionId(id, session_->version()));
   QUICHE_DCHECK_EQ(connect_stream_->id(), id);
-  connect_stream_->RegisterHttp3DatagramRegistrationVisitor(this);
+  connect_stream_->RegisterHttp3DatagramRegistrationVisitor(
+      this, attempt_to_use_datagram_contexts);
   if (session_->perspective() == Perspective::IS_CLIENT) {
     context_is_known_ = true;
     context_currently_registered_ = true;
+    if (attempt_to_use_datagram_contexts) {
+      context_id_ = connect_stream_->GetNextDatagramContextId();
+    }
   }
 }
 
diff --git a/quic/core/http/web_transport_http3.h b/quic/core/http/web_transport_http3.h
index 7a72931..cd446c1 100644
--- a/quic/core/http/web_transport_http3.h
+++ b/quic/core/http/web_transport_http3.h
@@ -34,7 +34,8 @@
       public QuicSpdyStream::Http3DatagramVisitor {
  public:
   WebTransportHttp3(QuicSpdySession* session, QuicSpdyStream* connect_stream,
-                    WebTransportSessionId id);
+                    WebTransportSessionId id,
+                    bool attempt_to_use_datagram_contexts);
 
   void HeadersReceived(const spdy::SpdyHeaderBlock& headers);
   void SetVisitor(std::unique_ptr<WebTransportVisitor> visitor) {
@@ -109,7 +110,7 @@
   bool ready_ = false;
   // Whether we know which |context_id_| to use. On the client this is always
   // true, and on the server it becomes true when we receive a context
-  // registeration capsule.
+  // registration capsule.
   bool context_is_known_ = false;
   // Whether |context_id_| is currently registered with |connect_stream_|.
   bool context_currently_registered_ = false;
diff --git a/quic/test_tools/quic_spdy_stream_peer.cc b/quic/test_tools/quic_spdy_stream_peer.cc
index 896a5b0..6db9a4a 100644
--- a/quic/test_tools/quic_spdy_stream_peer.cc
+++ b/quic/test_tools/quic_spdy_stream_peer.cc
@@ -23,5 +23,10 @@
   return stream->unacked_frame_headers_offsets_;
 }
 
+// static
+bool QuicSpdyStreamPeer::use_datagram_contexts(QuicSpdyStream* stream) {
+  return stream->use_datagram_contexts_;
+}
+
 }  // namespace test
 }  // namespace quic
diff --git a/quic/test_tools/quic_spdy_stream_peer.h b/quic/test_tools/quic_spdy_stream_peer.h
index 0f04338..8a50426 100644
--- a/quic/test_tools/quic_spdy_stream_peer.h
+++ b/quic/test_tools/quic_spdy_stream_peer.h
@@ -23,6 +23,7 @@
       QuicReferenceCountedPointer<QuicAckListenerInterface> ack_listener);
   static const QuicIntervalSet<QuicStreamOffset>& unacked_frame_headers_offsets(
       QuicSpdyStream* stream);
+  static bool use_datagram_contexts(QuicSpdyStream* stream);
 };
 
 }  // namespace test
diff --git a/quic/test_tools/quic_test_backend.h b/quic/test_tools/quic_test_backend.h
index 6aa175e..b0c0ce2 100644
--- a/quic/test_tools/quic_test_backend.h
+++ b/quic/test_tools/quic_test_backend.h
@@ -26,6 +26,12 @@
     enable_webtransport_ = enable_webtransport;
   }
 
+  bool UsesDatagramContexts() override { return use_datagram_contexts_; }
+
+  void set_use_datagram_contexts(bool use_datagram_contexts) {
+    use_datagram_contexts_ = use_datagram_contexts;
+  }
+
   bool SupportsExtendedConnect() override { return enable_extended_connect_; }
 
   void set_enable_extended_connect(bool enable_extended_connect) {
@@ -34,6 +40,7 @@
 
  private:
   bool enable_webtransport_ = false;
+  bool use_datagram_contexts_ = false;
   bool enable_extended_connect_ = true;
 };
 
diff --git a/quic/tools/quic_client.cc b/quic/tools/quic_client.cc
index e659b4a..cfc435e 100644
--- a/quic/tools/quic_client.cc
+++ b/quic/tools/quic_client.cc
@@ -164,7 +164,8 @@
     QuicConnection* connection) {
   return std::make_unique<QuicSimpleClientSession>(
       *config(), supported_versions, connection, server_id(), crypto_config(),
-      push_promise_index(), drop_response_body(), enable_web_transport());
+      push_promise_index(), drop_response_body(), enable_web_transport(),
+      use_datagram_contexts());
 }
 
 QuicClientEpollNetworkHelper* QuicClient::epoll_network_helper() {
diff --git a/quic/tools/quic_simple_client_session.cc b/quic/tools/quic_simple_client_session.cc
index 47c521b..2e70919 100644
--- a/quic/tools/quic_simple_client_session.cc
+++ b/quic/tools/quic_simple_client_session.cc
@@ -9,39 +9,27 @@
 namespace quic {
 
 QuicSimpleClientSession::QuicSimpleClientSession(
-    const QuicConfig& config,
-    const ParsedQuicVersionVector& supported_versions,
-    QuicConnection* connection,
-    const QuicServerId& server_id,
+    const QuicConfig& config, const ParsedQuicVersionVector& supported_versions,
+    QuicConnection* connection, const QuicServerId& server_id,
     QuicCryptoClientConfig* crypto_config,
-    QuicClientPushPromiseIndex* push_promise_index,
-    bool drop_response_body)
-    : QuicSimpleClientSession(config,
-                              supported_versions,
-                              connection,
-                              server_id,
-                              crypto_config,
-                              push_promise_index,
+    QuicClientPushPromiseIndex* push_promise_index, bool drop_response_body)
+    : QuicSimpleClientSession(config, supported_versions, connection, server_id,
+                              crypto_config, push_promise_index,
                               drop_response_body,
-                              /*enable_web_transport=*/false) {}
+                              /*enable_web_transport=*/false,
+                              /*use_datagram_contexts=*/false) {}
 
 QuicSimpleClientSession::QuicSimpleClientSession(
-    const QuicConfig& config,
-    const ParsedQuicVersionVector& supported_versions,
-    QuicConnection* connection,
-    const QuicServerId& server_id,
+    const QuicConfig& config, const ParsedQuicVersionVector& supported_versions,
+    QuicConnection* connection, const QuicServerId& server_id,
     QuicCryptoClientConfig* crypto_config,
-    QuicClientPushPromiseIndex* push_promise_index,
-    bool drop_response_body,
-    bool enable_web_transport)
-    : QuicSpdyClientSession(config,
-                            supported_versions,
-                            connection,
-                            server_id,
-                            crypto_config,
-                            push_promise_index),
+    QuicClientPushPromiseIndex* push_promise_index, bool drop_response_body,
+    bool enable_web_transport, bool use_datagram_contexts)
+    : QuicSpdyClientSession(config, supported_versions, connection, server_id,
+                            crypto_config, push_promise_index),
       drop_response_body_(drop_response_body),
-      enable_web_transport_(enable_web_transport) {}
+      enable_web_transport_(enable_web_transport),
+      use_datagram_contexts_(use_datagram_contexts) {}
 
 std::unique_ptr<QuicSpdyClientStream>
 QuicSimpleClientSession::CreateClientStream() {
@@ -54,6 +42,10 @@
   return enable_web_transport_;
 }
 
+bool QuicSimpleClientSession::ShouldNegotiateDatagramContexts() {
+  return use_datagram_contexts_;
+}
+
 HttpDatagramSupport QuicSimpleClientSession::LocalHttpDatagramSupport() {
   return enable_web_transport_ ? HttpDatagramSupport::kDraft04
                                : HttpDatagramSupport::kNone;
diff --git a/quic/tools/quic_simple_client_session.h b/quic/tools/quic_simple_client_session.h
index 1a6e694..6371448 100644
--- a/quic/tools/quic_simple_client_session.h
+++ b/quic/tools/quic_simple_client_session.h
@@ -25,16 +25,18 @@
                           const QuicServerId& server_id,
                           QuicCryptoClientConfig* crypto_config,
                           QuicClientPushPromiseIndex* push_promise_index,
-                          bool drop_response_body,
-                          bool enable_web_transport);
+                          bool drop_response_body, bool enable_web_transport,
+                          bool use_datagram_contexts);
 
   std::unique_ptr<QuicSpdyClientStream> CreateClientStream() override;
   bool ShouldNegotiateWebTransport() override;
+  bool ShouldNegotiateDatagramContexts() override;
   HttpDatagramSupport LocalHttpDatagramSupport() override;
 
  private:
   const bool drop_response_body_;
   const bool enable_web_transport_;
+  const bool use_datagram_contexts_;
 };
 
 }  // namespace quic
diff --git a/quic/tools/quic_simple_server_backend.h b/quic/tools/quic_simple_server_backend.h
index 67aeac9..25204ac 100644
--- a/quic/tools/quic_simple_server_backend.h
+++ b/quic/tools/quic_simple_server_backend.h
@@ -67,6 +67,7 @@
     return response;
   }
   virtual bool SupportsWebTransport() { return false; }
+  virtual bool UsesDatagramContexts() { return false; }
   virtual bool SupportsExtendedConnect() { return true; }
 };
 
diff --git a/quic/tools/quic_simple_server_session.h b/quic/tools/quic_simple_server_session.h
index 9746f26..393a2d8 100644
--- a/quic/tools/quic_simple_server_session.h
+++ b/quic/tools/quic_simple_server_session.h
@@ -70,6 +70,10 @@
 
   void OnCanCreateNewOutgoingStream(bool unidirectional) override;
 
+  bool ShouldNegotiateDatagramContexts() override {
+    return quic_simple_server_backend_->UsesDatagramContexts();
+  }
+
  protected:
   // QuicSession methods:
   QuicSpdyStream* CreateIncomingStream(QuicStreamId id) override;
diff --git a/quic/tools/quic_spdy_client_base.h b/quic/tools/quic_spdy_client_base.h
index c49d0ed..3e41725 100644
--- a/quic/tools/quic_spdy_client_base.h
+++ b/quic/tools/quic_spdy_client_base.h
@@ -144,6 +144,11 @@
   }
   bool enable_web_transport() const { return enable_web_transport_; }
 
+  void set_use_datagram_contexts(bool use_datagram_contexts) {
+    use_datagram_contexts_ = use_datagram_contexts;
+  }
+  bool use_datagram_contexts() const { return use_datagram_contexts_; }
+
   // QuicClientBase methods.
   bool goaway_received() const override;
   bool EarlyDataAccepted() override;
@@ -227,6 +232,7 @@
 
   bool drop_response_body_ = false;
   bool enable_web_transport_ = false;
+  bool use_datagram_contexts_ = false;
   // If not zero, used to set client's max inbound header size before session
   // initialize.
   size_t max_inbound_header_list_size_ = 0;