Implement DTLS-STUN piggybacking controller

which implements the handshaking logic of the DTLS-STUN piggybacking.

Not wired up yet, split from
  https://webrtc-review.googlesource.com/c/src/+/362480

BUG=webrtc:367395350

Change-Id: I9ee8ff17af4ec96fb891d9852ac50825155735a8
Reviewed-on: https://webrtc-review.googlesource.com/c/src/+/370679
Reviewed-by: Harald Alvestrand <hta@webrtc.org>
Commit-Queue: Jonas Oreland <jonaso@webrtc.org>
Reviewed-by: Jonas Oreland <jonaso@webrtc.org>
Cr-Commit-Position: refs/heads/main@{#43606}

api/transport/stun.h

@@ -722,6 +722,9 @@ enum IceAttributeType {
   STUN_ATTR_GOOG_DELTA_SYNC_REQ = 0xC05E,  // Not yet implemented.
   // MESSAGE-INTEGRITY truncated to 32-bit.
   STUN_ATTR_GOOG_MESSAGE_INTEGRITY_32 = 0xC060,
+  // Experimental: piggybacking the DTLS handshake in STUN.
+  STUN_ATTR_META_DTLS_IN_STUN = 0xC070,
+  STUN_ATTR_META_DTLS_IN_STUN_ACK = 0xC071,
 };
 
 // When adding new attributes to STUN_ATTR_GOOG_MISC_INFO

p2p/BUILD.gn

@@ -61,6 +61,7 @@ rtc_library("rtc_p2p") {
     ":candidate_pair_interface",
     ":connection",
     ":connection_info",
+    ":dtls_stun_piggyback_controller",
     ":dtls_transport",
     ":dtls_transport_internal",
     ":ice_agent_interface",
@@ -797,6 +798,28 @@ rtc_library("dtls_utils") {
   ]
 }
 
+rtc_library("dtls_stun_piggyback_controller") {
+  sources = [
+    "dtls/dtls_stun_piggyback_controller.cc",
+    "dtls/dtls_stun_piggyback_controller.h",
+  ]
+  deps = [
+    ":dtls_utils",
+    "../api:array_view",
+    "../api:sequence_checker",
+    "../api/transport:stun_types",
+    "../rtc_base:buffer",
+    "../rtc_base:byte_buffer",
+    "../rtc_base:checks",
+    "../rtc_base:logging",
+    "../rtc_base:macromagic",
+    "../rtc_base:stringutils",
+    "../rtc_base/system:no_unique_address",
+    "//third_party/abseil-cpp/absl/functional:any_invocable",
+    "//third_party/abseil-cpp/absl/strings:string_view",
+  ]
+}
+
 rtc_library("stun_port") {
   sources = [
     "base/stun_port.cc",
@@ -1142,6 +1165,7 @@ if (rtc_include_tests) {
       "base/turn_server_unittest.cc",
       "base/wrapping_active_ice_controller_unittest.cc",
       "client/basic_port_allocator_unittest.cc",
+      "dtls/dtls_stun_piggyback_controller_unittest.cc",
       "dtls/dtls_transport_unittest.cc",
       "dtls/dtls_utils_unittest.cc",
     ]
@@ -1153,6 +1177,7 @@ if (rtc_include_tests) {
       ":basic_packet_socket_factory",
       ":basic_port_allocator",
       ":connection",
+      ":dtls_stun_piggyback_controller",
       ":dtls_transport",
       ":dtls_transport_internal",
       ":dtls_utils",

p2p/DEPS

@@ -2,4 +2,5 @@ include_rules = [
   "+logging",
   "+net",
   "+system_wrappers",
+  "+absl/functional/any_invocable.h",
 ]

p2p/dtls/dtls_stun_piggyback_controller.cc

@@ -0,0 +1,168 @@
+/*
+ *  Copyright 2024 The WebRTC Project Authors. All rights reserved.
+ *
+ *  Use of this source code is governed by a BSD-style license
+ *  that can be found in the LICENSE file in the root of the source
+ *  tree. An additional intellectual property rights grant can be found
+ *  in the file PATENTS.  All contributing project authors may
+ *  be found in the AUTHORS file in the root of the source tree.
+ */
+
+#include "p2p/dtls/dtls_stun_piggyback_controller.h"
+
+#include <cstdint>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include "absl/functional/any_invocable.h"
+#include "absl/strings/string_view.h"
+#include "api/array_view.h"
+#include "api/sequence_checker.h"
+#include "api/transport/stun.h"
+#include "p2p/dtls/dtls_utils.h"
+#include "rtc_base/checks.h"
+#include "rtc_base/logging.h"
+#include "rtc_base/string_encode.h"
+
+namespace cricket {
+
+DtlsStunPiggybackController::DtlsStunPiggybackController(
+    absl::AnyInvocable<void(rtc::ArrayView<const uint8_t>)> dtls_data_callback)
+    : dtls_data_callback_(std::move(dtls_data_callback)) {}
+
+DtlsStunPiggybackController::~DtlsStunPiggybackController() {}
+
+void DtlsStunPiggybackController::SetDtlsHandshakeComplete(
+    bool is_dtls_client) {
+  RTC_DCHECK_RUN_ON(&sequence_checker_);
+  // Peer does not support this so fallback to a normal DTLS handshake
+  // happened.
+  if (state_ == State::OFF) {
+    return;
+  }
+  // As DTLS server we need to keep the last flight around until
+  // we receive the post-handshake acknowledgment.
+  // As DTLS client we have nothing more to send at this point
+  // but will continue to send ACK attributes until receiving
+  // the last flight from the server.
+  state_ = State::PENDING;
+  if (is_dtls_client) {
+    pending_packet_.Clear();
+  }
+}
+
+void DtlsStunPiggybackController::SetDataToPiggyback(
+    rtc::ArrayView<const uint8_t> data) {
+  RTC_DCHECK_RUN_ON(&sequence_checker_);
+  if (state_ == State::OFF) {
+    return;
+  }
+  // Note: this overwrites the existing packets which is an issue
+  // if this gets called with fragmented DTLS flights.
+  pending_packet_.SetData(data);
+}
+
+std::optional<absl::string_view>
+DtlsStunPiggybackController::GetDataToPiggyback(
+    StunMessageType stun_message_type) {
+  RTC_DCHECK_RUN_ON(&sequence_checker_);
+  RTC_DCHECK(stun_message_type == STUN_BINDING_REQUEST ||
+             stun_message_type == STUN_BINDING_RESPONSE);
+  if (state_ == State::OFF || state_ == State::COMPLETE) {
+    return std::nullopt;
+  }
+  if (pending_packet_.size() == 0) {
+    return std::nullopt;
+  }
+  return absl::string_view(pending_packet_);
+}
+
+std::optional<absl::string_view> DtlsStunPiggybackController::GetAckToPiggyback(
+    StunMessageType stun_message_type) {
+  RTC_DCHECK_RUN_ON(&sequence_checker_);
+  if (state_ == State::OFF || state_ == State::COMPLETE) {
+    return std::nullopt;
+  }
+  return handshake_ack_writer_.DataAsStringView();
+}
+
+void DtlsStunPiggybackController::ReportDataPiggybacked(
+    const StunByteStringAttribute* data,
+    const StunByteStringAttribute* ack) {
+  RTC_DCHECK_RUN_ON(&sequence_checker_);
+
+  // Drop silently when receiving acked data when the peer previously did not
+  // support or we already moved to the complete state.
+  if (state_ == State::OFF || state_ == State::COMPLETE) {
+    return;
+  }
+
+  // We sent dtls piggybacked but got nothing in return or
+  // we received a stun request with neither attribute set
+  // => peer does not support.
+  if (state_ == State::TENTATIVE && data == nullptr && ack == nullptr) {
+    state_ = State::OFF;
+    pending_packet_.Clear();
+    RTC_LOG(LS_INFO) << "DTLS-STUN piggybacking not supported by peer.";
+    return;
+  }
+
+  // In PENDING state the peer may have stopped sending the ack
+  // when it moved to the COMPLETE state. Move to the same state.
+  if (state_ == State::PENDING && data == nullptr && ack == nullptr) {
+    RTC_LOG(LS_INFO) << "DTLS-STUN piggybacking complete.";
+    state_ = State::COMPLETE;
+    pending_packet_.Clear();
+    handshake_ack_writer_.Clear();
+    handshake_messages_received_.clear();
+    return;
+  }
+
+  // We sent dtls piggybacked and got something in return => peer does support.
+  if (state_ == State::TENTATIVE) {
+    state_ = State::CONFIRMED;
+  }
+
+  if (ack != nullptr) {
+    RTC_LOG(LS_VERBOSE) << "DTLS-STUN piggybacking ACK: "
+                        << rtc::hex_encode(ack->string_view());
+  }
+  // The response to the final flight of the handshake will not contain
+  // the DTLS data but will contain an ack.
+  // Must not happen on the initial server to client packet which
+  // has no DTLS data yet.
+  if (data == nullptr && ack != nullptr && state_ == State::PENDING) {
+    RTC_LOG(LS_INFO) << "DTLS-STUN piggybacking complete.";
+    state_ = State::COMPLETE;
+    pending_packet_.Clear();
+    handshake_ack_writer_.Clear();
+    handshake_messages_received_.clear();
+    return;
+  }
+  if (!data || data->length() == 0) {
+    return;
+  }
+
+  // Extract the received message sequence numbers of the handshake
+  // from the packet and prepare the ack to be sent.
+  std::optional<std::vector<uint16_t>> new_message_sequences =
+      GetDtlsHandshakeAcks(data->array_view());
+  if (!new_message_sequences) {
+    RTC_LOG(LS_ERROR) << "DTLS-STUN piggybacking failed to parse DTLS packet.";
+    return;
+  }
+  if (!new_message_sequences->empty()) {
+    for (const auto& message_seq : *new_message_sequences) {
+      handshake_messages_received_.insert(message_seq);
+    }
+    handshake_ack_writer_.Clear();
+    for (const auto& message_seq : handshake_messages_received_) {
+      handshake_ack_writer_.WriteUInt16(message_seq);
+    }
+  }
+
+  dtls_data_callback_(data->array_view());
+}
+
+}  // namespace cricket

p2p/dtls/dtls_stun_piggyback_controller.h

@@ -0,0 +1,94 @@
+/*
+ *  Copyright 2024 The WebRTC Project Authors. All rights reserved.
+ *
+ *  Use of this source code is governed by a BSD-style license
+ *  that can be found in the LICENSE file in the root of the source
+ *  tree. An additional intellectual property rights grant can be found
+ *  in the file PATENTS.  All contributing project authors may
+ *  be found in the AUTHORS file in the root of the source tree.
+ */
+
+#ifndef P2P_DTLS_DTLS_STUN_PIGGYBACK_CONTROLLER_H_
+#define P2P_DTLS_DTLS_STUN_PIGGYBACK_CONTROLLER_H_
+
+#include <cstdint>
+#include <optional>
+#include <set>
+
+#include "absl/functional/any_invocable.h"
+#include "absl/strings/string_view.h"
+#include "api/array_view.h"
+#include "api/sequence_checker.h"
+#include "api/transport/stun.h"
+#include "rtc_base/buffer.h"
+#include "rtc_base/byte_buffer.h"
+#include "rtc_base/system/no_unique_address.h"
+#include "rtc_base/thread_annotations.h"
+
+namespace cricket {
+
+// This class is not thread safe; all methods must be called on the same thread
+// as the constructor.
+class DtlsStunPiggybackController {
+ public:
+  // dtls_data_callback will be called with any DTLS packets received
+  // piggybacked.
+  DtlsStunPiggybackController(
+      absl::AnyInvocable<void(rtc::ArrayView<const uint8_t>)>
+          dtls_data_callback);
+  ~DtlsStunPiggybackController();
+
+  enum class State {
+    // We don't know if peer support DTLS piggybacked in STUN.
+    // We will piggyback DTLS until we get a piggybacked response
+    // or a STUN response with piggyback support.
+    TENTATIVE = 0,
+    // The peer supports DTLS in STUN and we continue the handshake.
+    CONFIRMED = 1,
+    // We are waiting for the final ack. Semantic differs depending
+    // on DTLS role.
+    PENDING = 2,
+    // We successfully completed the DTLS handshake in STUN.
+    COMPLETE = 3,
+    // The peer does not support piggybacking DTLS in STUN.
+    OFF = 4,
+  };
+
+  State state() const {
+    RTC_DCHECK_RUN_ON(&sequence_checker_);
+    return state_;
+  }
+
+  // Called by DtlsTransport when handshake is complete.
+  void SetDtlsHandshakeComplete(bool is_dtls_client);
+
+  // Called by DtlsTransport transport when there is data to piggyback.
+  void SetDataToPiggyback(rtc::ArrayView<const uint8_t> data);
+
+  // Called by Connection, when sending a STUN BINDING { REQUEST / RESPONSE }
+  // to obtain optional DTLS data or ACKs.
+  std::optional<absl::string_view> GetDataToPiggyback(
+      StunMessageType stun_message_type);
+  std::optional<absl::string_view> GetAckToPiggyback(
+      StunMessageType stun_message_type);
+
+  // Called by Connection when receiving a STUN BINDING { REQUEST / RESPONSE }.
+  void ReportDataPiggybacked(const StunByteStringAttribute* data,
+                             const StunByteStringAttribute* ack);
+
+ private:
+  State state_ RTC_GUARDED_BY(sequence_checker_) = State::TENTATIVE;
+  rtc::Buffer pending_packet_ RTC_GUARDED_BY(sequence_checker_);
+  absl::AnyInvocable<void(rtc::ArrayView<const uint8_t>)> dtls_data_callback_;
+
+  std::set<uint16_t> handshake_messages_received_
+      RTC_GUARDED_BY(sequence_checker_);
+  rtc::ByteBufferWriter handshake_ack_writer_ RTC_GUARDED_BY(sequence_checker_);
+
+  // In practice this will be the network thread.
+  RTC_NO_UNIQUE_ADDRESS webrtc::SequenceChecker sequence_checker_;
+};
+
+}  // namespace cricket
+
+#endif  // P2P_DTLS_DTLS_STUN_PIGGYBACK_CONTROLLER_H_

p2p/dtls/dtls_stun_piggyback_controller_unittest.cc

@@ -0,0 +1,288 @@
+/*
+ *  Copyright 2024 The WebRTC Project Authors. All rights reserved.
+ *
+ *  Use of this source code is governed by a BSD-style license
+ *  that can be found in the LICENSE file in the root of the source
+ *  tree. An additional intellectual property rights grant can be found
+ *  in the file PATENTS.  All contributing project authors may
+ *  be found in the AUTHORS file in the root of the source tree.
+ */
+
+#include "p2p/dtls/dtls_stun_piggyback_controller.h"
+
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "api/array_view.h"
+#include "api/transport/stun.h"
+#include "test/gtest.h"
+
+namespace {
+// Extracted from a stock DTLS call using Wireshark.
+// Each packet (apart from the last) is truncated to
+// the first fragment to keep things short.
+
+// Based on a "server hello done" but with different msg_seq.
+const std::vector<uint8_t> dtls_flight1 = {
+    0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  //
+    0x00, 0x01,                                            // seq=1
+    0x00, 0x0c, 0x0e, 0x00, 0x00, 0x00, 0x12, 0x34, 0x00,  // msg_seq=0x1234
+    0x00, 0x00, 0x00, 0x00, 0x00};
+
+const std::vector<uint8_t> dtls_flight2 = {
+    0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  //
+    0x00, 0x02,                                            // seq=2
+    0x00, 0x0c, 0x0e, 0x00, 0x00, 0x00, 0x43, 0x21, 0x00,  // msg_seq=0x4321
+    0x00, 0x00, 0x00, 0x00, 0x00};
+
+const std::vector<uint8_t> dtls_flight3 = {
+    0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  //
+    0x00, 0x03,                                            // seq=3
+    0x00, 0x0c, 0x0e, 0x00, 0x00, 0x00, 0x44, 0x44, 0x00,  // msg_seq=0x4444
+    0x00, 0x00, 0x00, 0x00, 0x00};
+
+const std::vector<uint8_t> dtls_flight4 = {
+    0x16, 0xfe, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,  //
+    0x00, 0x04,                                            // seq=4
+    0x00, 0x0c, 0x0e, 0x00, 0x00, 0x00, 0x54, 0x86, 0x00,  // msg_seq=0x5486
+    0x00, 0x00, 0x00, 0x00, 0x00};
+
+const std::vector<uint8_t> empty = {};
+}  // namespace
+
+namespace cricket {
+
+using State = DtlsStunPiggybackController::State;
+
+class DtlsStunPiggybackControllerTest : public ::testing::Test {
+ protected:
+  DtlsStunPiggybackControllerTest()
+      : client_([](rtc::ArrayView<const uint8_t> data) {}),
+        server_([](rtc::ArrayView<const uint8_t> data) {}) {}
+
+  void SendClientToServer(const std::vector<uint8_t> data,
+                          StunMessageType type) {
+    client_.SetDataToPiggyback(data);
+    std::unique_ptr<StunByteStringAttribute> attr_data;
+    if (client_.GetDataToPiggyback(type)) {
+      attr_data = std::make_unique<StunByteStringAttribute>(
+          STUN_ATTR_META_DTLS_IN_STUN, *client_.GetDataToPiggyback(type));
+    }
+    std::unique_ptr<StunByteStringAttribute> attr_ack;
+    if (client_.GetAckToPiggyback(type)) {
+      attr_ack = std::make_unique<StunByteStringAttribute>(
+          STUN_ATTR_META_DTLS_IN_STUN_ACK, *client_.GetAckToPiggyback(type));
+    }
+    server_.ReportDataPiggybacked(attr_data.get(), attr_ack.get());
+    if (data == dtls_flight3) {
+      // When receiving flight 3, server handshake is complete.
+      server_.SetDtlsHandshakeComplete(/*is_client=*/false);
+    }
+  }
+  void SendServerToClient(const std::vector<uint8_t> data,
+                          StunMessageType type) {
+    server_.SetDataToPiggyback(data);
+    std::unique_ptr<StunByteStringAttribute> attr_data;
+    if (server_.GetDataToPiggyback(type)) {
+      attr_data = std::make_unique<StunByteStringAttribute>(
+          STUN_ATTR_META_DTLS_IN_STUN, *server_.GetDataToPiggyback(type));
+    }
+    std::unique_ptr<StunByteStringAttribute> attr_ack;
+    if (server_.GetAckToPiggyback(type)) {
+      attr_ack = std::make_unique<StunByteStringAttribute>(
+          STUN_ATTR_META_DTLS_IN_STUN_ACK, *server_.GetAckToPiggyback(type));
+    }
+    client_.ReportDataPiggybacked(attr_data.get(), attr_ack.get());
+    if (data == dtls_flight4) {
+      // When receiving flight 4, client handshake is complete.
+      client_.SetDtlsHandshakeComplete(/*is_client=*/true);
+    }
+  }
+
+  void DisableSupport(DtlsStunPiggybackController& client_or_server) {
+    ASSERT_EQ(client_or_server.state(), State::TENTATIVE);
+    client_or_server.ReportDataPiggybacked(nullptr, nullptr);
+    ASSERT_EQ(client_or_server.state(), State::OFF);
+  }
+
+  DtlsStunPiggybackController client_;
+  DtlsStunPiggybackController server_;
+};
+
+TEST_F(DtlsStunPiggybackControllerTest, BasicHandshake) {
+  // Flight 1+2
+  SendClientToServer(dtls_flight1, STUN_BINDING_REQUEST);
+  EXPECT_EQ(server_.state(), State::CONFIRMED);
+  SendServerToClient(dtls_flight2, STUN_BINDING_RESPONSE);
+  EXPECT_EQ(client_.state(), State::CONFIRMED);
+
+  // Flight 3+4
+  SendClientToServer(dtls_flight3, STUN_BINDING_REQUEST);
+  SendServerToClient(dtls_flight4, STUN_BINDING_RESPONSE);
+  EXPECT_EQ(server_.state(), State::PENDING);
+  EXPECT_EQ(client_.state(), State::PENDING);
+
+  // Post-handshake ACK
+  SendServerToClient(empty, STUN_BINDING_REQUEST);
+  SendClientToServer(empty, STUN_BINDING_RESPONSE);
+  EXPECT_EQ(server_.state(), State::COMPLETE);
+  EXPECT_EQ(client_.state(), State::COMPLETE);
+}
+
+TEST_F(DtlsStunPiggybackControllerTest, FirstClientPacketLost) {
+  // Client to server got lost (or arrives late)
+  // Flight 1
+  SendServerToClient(empty, STUN_BINDING_REQUEST);
+  SendClientToServer(dtls_flight1, STUN_BINDING_RESPONSE);
+  EXPECT_EQ(server_.state(), State::CONFIRMED);
+  EXPECT_EQ(client_.state(), State::CONFIRMED);
+
+  // Flight 2+3
+  SendServerToClient(dtls_flight2, STUN_BINDING_REQUEST);
+  SendClientToServer(dtls_flight3, STUN_BINDING_RESPONSE);
+  EXPECT_EQ(server_.state(), State::PENDING);
+  EXPECT_EQ(client_.state(), State::CONFIRMED);
+
+  // Flight 4
+  SendServerToClient(dtls_flight4, STUN_BINDING_REQUEST);
+  SendClientToServer(empty, STUN_BINDING_RESPONSE);
+  EXPECT_EQ(server_.state(), State::COMPLETE);
+  EXPECT_EQ(client_.state(), State::PENDING);
+
+  // Post-handshake ACK
+  SendServerToClient(empty, STUN_BINDING_REQUEST);
+  EXPECT_EQ(client_.state(), State::COMPLETE);
+}
+
+TEST_F(DtlsStunPiggybackControllerTest, NotSupportedByServer) {
+  DisableSupport(server_);
+
+  // Flight 1
+  SendClientToServer(dtls_flight1, STUN_BINDING_REQUEST);
+  SendServerToClient(empty, STUN_BINDING_RESPONSE);
+  EXPECT_EQ(client_.state(), State::OFF);
+}
+
+TEST_F(DtlsStunPiggybackControllerTest, NotSupportedByServerClientReceives) {
+  DisableSupport(server_);
+
+  // Client to server got lost (or arrives late)
+  SendServerToClient(empty, STUN_BINDING_REQUEST);
+  EXPECT_EQ(client_.state(), State::OFF);
+}
+
+TEST_F(DtlsStunPiggybackControllerTest, NotSupportedByClient) {
+  DisableSupport(client_);
+
+  SendServerToClient(empty, STUN_BINDING_REQUEST);
+  SendClientToServer(empty, STUN_BINDING_RESPONSE);
+  EXPECT_EQ(server_.state(), State::OFF);
+}
+
+TEST_F(DtlsStunPiggybackControllerTest, SomeRequestsDoNotGoThrough) {
+  // Client to server got lost (or arrives late)
+  // Flight 1
+  SendServerToClient(empty, STUN_BINDING_REQUEST);
+  SendClientToServer(dtls_flight1, STUN_BINDING_RESPONSE);
+  EXPECT_EQ(server_.state(), State::CONFIRMED);
+  EXPECT_EQ(client_.state(), State::CONFIRMED);
+
+  // Flight 1+2, server sent request got lost.
+  SendClientToServer(dtls_flight1, STUN_BINDING_REQUEST);
+  SendServerToClient(dtls_flight2, STUN_BINDING_RESPONSE);
+  EXPECT_EQ(server_.state(), State::CONFIRMED);
+  EXPECT_EQ(client_.state(), State::CONFIRMED);
+
+  // Flight 3+4
+  SendClientToServer(dtls_flight3, STUN_BINDING_REQUEST);
+  SendServerToClient(dtls_flight4, STUN_BINDING_RESPONSE);
+  EXPECT_EQ(server_.state(), State::PENDING);
+  EXPECT_EQ(client_.state(), State::PENDING);
+
+  // Post-handshake ACK
+  SendClientToServer(empty, STUN_BINDING_REQUEST);
+  SendServerToClient(empty, STUN_BINDING_RESPONSE);
+  EXPECT_EQ(server_.state(), State::COMPLETE);
+  EXPECT_EQ(client_.state(), State::COMPLETE);
+}
+
+TEST_F(DtlsStunPiggybackControllerTest, LossOnPostHandshakeAck) {
+  // Flight 1+2
+  SendClientToServer(dtls_flight1, STUN_BINDING_REQUEST);
+  EXPECT_EQ(server_.state(), State::CONFIRMED);
+  SendServerToClient(dtls_flight2, STUN_BINDING_RESPONSE);
+  EXPECT_EQ(client_.state(), State::CONFIRMED);
+
+  // Flight 3+4
+  SendClientToServer(dtls_flight3, STUN_BINDING_REQUEST);
+  SendServerToClient(dtls_flight4, STUN_BINDING_RESPONSE);
+  EXPECT_EQ(server_.state(), State::PENDING);
+  EXPECT_EQ(client_.state(), State::PENDING);
+
+  // Post-handshake ACK. Client to server gets lost
+  SendServerToClient(empty, STUN_BINDING_REQUEST);
+  SendClientToServer(empty, STUN_BINDING_RESPONSE);
+  EXPECT_EQ(server_.state(), State::COMPLETE);
+  EXPECT_EQ(client_.state(), State::COMPLETE);
+}
+
+TEST_F(DtlsStunPiggybackControllerTest,
+       UnsupportedStateAfterFallbackHandshakeRemainsOff) {
+  DisableSupport(client_);
+  DisableSupport(server_);
+
+  // Set DTLS complete after normal handshake.
+  client_.SetDtlsHandshakeComplete(true);
+  EXPECT_EQ(client_.state(), State::OFF);
+  server_.SetDtlsHandshakeComplete(true);
+  EXPECT_EQ(server_.state(), State::OFF);
+}
+
+TEST_F(DtlsStunPiggybackControllerTest, BasicHandshakeAckData) {
+  EXPECT_EQ(server_.GetAckToPiggyback(STUN_BINDING_RESPONSE), "");
+  EXPECT_EQ(client_.GetAckToPiggyback(STUN_BINDING_REQUEST), "");
+
+  // Flight 1+2
+  SendClientToServer(dtls_flight1, STUN_BINDING_REQUEST);
+  SendServerToClient(dtls_flight2, STUN_BINDING_RESPONSE);
+  EXPECT_EQ(server_.GetAckToPiggyback(STUN_BINDING_REQUEST),
+            std::string("\x12\x34", 2));
+  EXPECT_EQ(client_.GetAckToPiggyback(STUN_BINDING_RESPONSE),
+            std::string("\x43\x21", 2));
+
+  // Flight 3+4
+  SendClientToServer(dtls_flight3, STUN_BINDING_REQUEST);
+  SendServerToClient(dtls_flight4, STUN_BINDING_RESPONSE);
+  EXPECT_EQ(server_.GetAckToPiggyback(STUN_BINDING_RESPONSE),
+            std::string("\x12\x34\x44\x44", 4));
+  EXPECT_EQ(client_.GetAckToPiggyback(STUN_BINDING_REQUEST),
+            std::string("\x43\x21\x54\x86", 4));
+
+  // Post-handshake ACK
+  SendServerToClient(empty, STUN_BINDING_REQUEST);
+  SendClientToServer(empty, STUN_BINDING_RESPONSE);
+  EXPECT_EQ(server_.state(), State::COMPLETE);
+  EXPECT_EQ(client_.state(), State::COMPLETE);
+  EXPECT_EQ(server_.GetAckToPiggyback(STUN_BINDING_RESPONSE), std::nullopt);
+  EXPECT_EQ(client_.GetAckToPiggyback(STUN_BINDING_REQUEST), std::nullopt);
+}
+
+TEST_F(DtlsStunPiggybackControllerTest, AckDataNoDuplicates) {
+  // Flight 1+2
+  SendClientToServer(dtls_flight1, STUN_BINDING_REQUEST);
+  EXPECT_EQ(server_.GetAckToPiggyback(STUN_BINDING_REQUEST),
+            std::string("\x12\x34", 2));
+  SendClientToServer(dtls_flight3, STUN_BINDING_REQUEST);
+  EXPECT_EQ(server_.GetAckToPiggyback(STUN_BINDING_REQUEST),
+            std::string("\x12\x34\x44\x44", 4));
+
+  // Receive Flight 1 again, no change expected.
+  SendClientToServer(dtls_flight1, STUN_BINDING_REQUEST);
+  EXPECT_EQ(server_.GetAckToPiggyback(STUN_BINDING_REQUEST),
+            std::string("\x12\x34\x44\x44", 4));
+}
+
+}  // namespace cricket