CHECK on adding STUN attributes after signature is applied.

This is a programming error, and needs to be fixed in the codebase
before we can start ignoring such attributes.

Bug: chromium:504567957
Change-Id: I1107b8099d7c6c09afc98147159870e825350034
Reviewed-on: https://webrtc-review.googlesource.com/c/src/+/466000
Reviewed-by: Danil Chapovalov <danilchap@webrtc.org>
Auto-Submit: Harald Alvestrand <hta@webrtc.org>
Commit-Queue: Harald Alvestrand <hta@webrtc.org>
Cr-Commit-Position: refs/heads/main@{#47604}
diff --git a/api/transport/BUILD.gn b/api/transport/BUILD.gn
index e0f7e52..407a99e 100644
--- a/api/transport/BUILD.gn
+++ b/api/transport/BUILD.gn
@@ -173,12 +173,14 @@
       ":stun_types",
       "../../rtc_base:byte_buffer",
       "../../rtc_base:byte_order",
+      "../../rtc_base:checks",
       "../../rtc_base:ip_address",
       "../../rtc_base:socket_address",
       "../../rtc_base:span_helpers",
       "../../system_wrappers:metrics",
       "../../test:test_support",
       "//testing/gtest",
+      "//third_party/abseil-cpp/absl/strings:string_view",
     ]
   }
 }
diff --git a/api/transport/stun.cc b/api/transport/stun.cc
index 47010ba..8802b4b 100644
--- a/api/transport/stun.cc
+++ b/api/transport/stun.cc
@@ -142,6 +142,15 @@
 }
 
 void StunMessage::AddAttribute(std::unique_ptr<StunAttribute> attr) {
+  // Once a message is signed, adding attributes is a programming error.
+  // Exceptions are the Integrity attributes and Fingerprint.
+  if (attr->type() != STUN_ATTR_MESSAGE_INTEGRITY &&
+      attr->type() != STUN_ATTR_GOOG_MESSAGE_INTEGRITY_32 &&
+      attr->type() != STUN_ATTR_FINGERPRINT) {
+    RTC_CHECK(integrity_ != IntegrityStatus::kIntegrityOk)
+        << "You cannot add attributes to a signed STUN message."
+        << " Message type was 0x" << ToHex(attr->type());
+  }
   // Fail any attributes that aren't valid for this type of message,
   // but allow any type for the range that in the RFC is reserved for
   // the "designated experts".
@@ -160,6 +169,12 @@
 }
 
 std::unique_ptr<StunAttribute> StunMessage::RemoveAttribute(int type) {
+  if (type != STUN_ATTR_MESSAGE_INTEGRITY &&
+      type != STUN_ATTR_GOOG_MESSAGE_INTEGRITY_32) {
+    // Removing attributes will break integrity, so don't allow
+    // removing attributes from a signed message.
+    RTC_CHECK(integrity_ != IntegrityStatus::kIntegrityOk);
+  }
   std::unique_ptr<StunAttribute> attribute;
   for (auto it = attrs_.rbegin(); it != attrs_.rend(); ++it) {
     if ((*it)->type() == type) {
@@ -175,6 +190,12 @@
       attr_length += (4 - (attr_length % 4));
     }
     length_ -= static_cast<uint16_t>(attr_length + 4);
+    if (type == STUN_ATTR_MESSAGE_INTEGRITY ||
+        type == STUN_ATTR_GOOG_MESSAGE_INTEGRITY_32) {
+      // If the message was signed, it is now unsigned, and adding more
+      // attributes is possible.
+      integrity_ = IntegrityStatus::kNotSet;
+    }
   }
   return attribute;
 }
@@ -185,6 +206,7 @@
   }
   attrs_.clear();
   length_ = 0;
+  integrity_ = IntegrityStatus::kNotSet;
 }
 
 std::vector<uint16_t> StunMessage::GetNonComprehendedAttributes() const {
@@ -620,7 +642,6 @@
   attrs_.resize(0);
 
   size_t rest = buf->Length() - length_;
-  bool have_seen_integrity = false;
   while (buf->Length() > rest) {
     uint16_t attr_type, attr_length;
     if (!buf->ReadUInt16(&attr_type))
@@ -642,20 +663,7 @@
       if (!attr->Read(buf)) {
         return false;
       }
-      // If the message is integrity-protected, ignore all attributes
-      // after the integrity attributes except for FINGERPRINT.
-      // RFC 8489 implicity enforces this; RFC 5389 section 15.4
-      // had explicit text saying it.
-      if (!have_seen_integrity || attr_type == STUN_ATTR_FINGERPRINT) {
-        attrs_.push_back(std::move(attr));
-      } else {
-        RTC_LOG(LS_WARNING) << "Attribute found after INTEGRITY: 0x"
-                            << ToHex(attr_type) << ", ignored";
-      }
-      if (attr_type == STUN_ATTR_MESSAGE_INTEGRITY ||
-          attr_type == STUN_ATTR_GOOG_MESSAGE_INTEGRITY_32) {
-        have_seen_integrity = true;
-      }
+      attrs_.push_back(std::move(attr));
     }
   }
 
diff --git a/api/transport/stun_unittest.cc b/api/transport/stun_unittest.cc
index 3e10500..0145966 100644
--- a/api/transport/stun_unittest.cc
+++ b/api/transport/stun_unittest.cc
@@ -21,8 +21,10 @@
 #include <utility>
 #include <vector>
 
+#include "absl/strings/string_view.h"
 #include "rtc_base/byte_buffer.h"
 #include "rtc_base/byte_order.h"
+#include "rtc_base/checks.h"
 #include "rtc_base/ip_address.h"
 #include "rtc_base/socket_address.h"
 #include "rtc_base/span_helpers.h"
@@ -35,6 +37,7 @@
 namespace {
 
 using ::testing::ElementsAreArray;
+using ::testing::NotNull;
 
 // Sample STUN packets with various attributes
 // Gathered by wiresharking pjproject's pjnath test programs
@@ -1686,49 +1689,51 @@
   EXPECT_THAT(*integrity_vector, ElementsAreArray(expected_integrity_vector));
 }
 
-TEST_F(StunTest, AttributesAfterMessageIntegrityAreIgnored) {
-  // 1. Create a legitimate STUN message.
-  IceMessage msg(STUN_BINDING_REQUEST, "abcdefghijkl");
-  auto username = StunAttribute::CreateByteString(STUN_ATTR_USERNAME);
-  username->CopyBytes("user");
-  msg.AddAttribute(std::move(username));
+TEST_F(StunTest, RemovingSigningAttributesAllowsAddingNewAttributes) {
+  IceMessage msg;
+  ByteBufferReader buf(kRfc5769SampleRequestWithoutMI);
+  EXPECT_TRUE(msg.Read(&buf));
+  EXPECT_TRUE(msg.AddMessageIntegrity(kRfc5769SampleMsgPassword));
+  ASSERT_TRUE(msg.IntegrityOk());
 
-  // 2. Add MESSAGE-INTEGRITY.
-  const std::string password = "password";
-  EXPECT_TRUE(msg.AddMessageIntegrity(password));
+  // Removing the MI attribute should make the message unsigned.
+  auto removed_attr = msg.RemoveAttribute(STUN_ATTR_MESSAGE_INTEGRITY);
+  ASSERT_THAT(removed_attr, NotNull());
+  EXPECT_FALSE(msg.IntegrityOk());
 
-  // 3. Serialize to buffer.
-  ByteBufferWriter buf;
-  EXPECT_TRUE(msg.Write(&buf));
+  // Now, removing and adding a new attribute should be possible.
+  msg.RemoveAttribute(STUN_ATTR_USERNAME);
+  auto attr = StunAttribute::CreateByteString(STUN_ATTR_USERNAME);
+  attr->CopyBytes(AsUint8Span(absl::string_view("newuser")));
+  msg.AddAttribute(std::move(attr));
 
-  // 4. Manually append an "attacker-injected" attribute AFTER the buffer.
-  // We'll use NOMINATION (0xC001) which is 4 bytes header + 4 bytes value.
-  buf.WriteUInt16(STUN_ATTR_NOMINATION);
-  buf.WriteUInt16(4);  // length
-  buf.WriteUInt32(0xDEADBEEF);
-
-  // 5. Update the STUN header length to include the new attribute.
-  std::span<const uint8_t> raw = buf.DataView();
-  std::vector<uint8_t> packet(raw.begin(), raw.end());
-  size_t new_total_size = packet.size();
-  uint16_t new_stun_length =
-      static_cast<uint16_t>(new_total_size - kStunHeaderSize);
-  SetBE16(std::span(packet).subspan(2, 2), new_stun_length);
-
-  // 6. Parse the tampered message.
-  IceMessage parsed_msg;
-  ByteBufferReader reader(packet);
-  EXPECT_TRUE(parsed_msg.Read(&reader));
-
-  // 7. Verify integrity still passes (this confirms MI check ignores trailing
-  // data).
-  EXPECT_EQ(parsed_msg.ValidateMessageIntegrity(password),
-            StunMessage::IntegrityStatus::kIntegrityOk);
-
-  // 8. According to RFC 8489, attributes following MESSAGE-INTEGRITY (except
-  // FINGERPRINT) MUST be ignored. This test case expects that the injected
-  // attribute is NOT present.
-  EXPECT_EQ(parsed_msg.GetUInt32(STUN_ATTR_NOMINATION), nullptr);
+  const StunByteStringAttribute* username =
+      msg.GetByteString(STUN_ATTR_USERNAME);
+  ASSERT_THAT(username, NotNull());
+  EXPECT_EQ("newuser", username->string_view());
 }
 
+#if RTC_DCHECK_IS_ON && GTEST_HAS_DEATH_TEST && !defined(WEBRTC_ANDROID)
+TEST_F(StunTest, AddingAttributesAfterSigningCrashesDeathTest) {
+  IceMessage msg;
+  ByteBufferReader buf(kRfc5769SampleRequestWithoutMI);
+  EXPECT_TRUE(msg.Read(&buf));
+  EXPECT_TRUE(msg.AddMessageIntegrity(kRfc5769SampleMsgPassword));
+  ASSERT_TRUE(msg.IntegrityOk());
+  auto attr = StunAttribute::CreateByteString(STUN_ATTR_USERNAME);
+  attr->CopyBytes(AsUint8Span(absl::string_view("keso")));
+
+  EXPECT_DEATH(msg.AddAttribute(std::move(attr)), "");
+}
+
+TEST_F(StunTest, RemovingAttributesAfterSigningCrashesDeathTest) {
+  IceMessage msg;
+  ByteBufferReader buf(kRfc5769SampleRequestWithoutMI);
+  EXPECT_TRUE(msg.Read(&buf));
+  EXPECT_TRUE(msg.AddMessageIntegrity(kRfc5769SampleMsgPassword));
+  ASSERT_TRUE(msg.IntegrityOk());
+  EXPECT_DEATH(msg.RemoveAttribute(STUN_ATTR_USERNAME), "");
+}
+#endif
+
 }  // namespace webrtc