/*
 *  POC: Latent heap-use-after-free in
 *  webrtc::JsepTransportCollection::RollbackTransports().
 *
 *  RollbackTransports() chains map_change_callback_ invocations with
 *  short-circuiting `&&`, but UNCONDITIONALLY rewrites mid_to_transport_
 *  and calls DestroyUnusedTransports() afterwards. If any callback in the
 *  "restore old mappings" loop returns false, every later mid is silently
 *  skipped — yet the offer-created JsepTransport (and its owned
 *  RtpTransportInternal) those mids' channels still point at is freed.
 *
 *  In production the callback eventually reaches
 *  BaseChannel::SetRtpTransport(), which stores a *raw*
 *  RtpTransportInternal* (pc/channel.h:326, no MiraclePtr in
 *  third_party/webrtc). A skipped channel keeps that raw pointer aimed at
 *  freed memory; the next SendPacket / DisconnectFromRtpTransport_n /
 *  transport_name() dereferences it.
 *
 *  This test simulates BaseChannel's pointer storage and demonstrates the
 *  resulting heap-use-after-free under ASAN. It exercises the public
 *  JsepTransportCollection contract directly because, under current code,
 *  the only production source of a `false` return
 *  (RtpDemuxer::CriteriaWouldConflict) cannot be induced on the stable
 *  transport during rollback without first putting the PeerConnection into
 *  a session-error state that blocks rollback.
 */

#include <map>
#include <memory>
#include <set>
#include <string>
#include <utility>

#include "absl/strings/string_view.h"
#include "api/field_trials.h"
#include "api/ice_transport_interface.h"
#include "api/make_ref_counted.h"
#include "api/scoped_refptr.h"
#include "p2p/base/p2p_constants.h"
#include "p2p/dtls/dtls_transport_internal.h"
#include "p2p/dtls/fake_dtls_transport.h"
#include "p2p/test/fake_ice_transport.h"
#include "pc/dtls_srtp_transport.h"
#include "pc/dtls_transport.h"
#include "pc/jsep_transport.h"
#include "pc/jsep_transport_collection.h"
#include "pc/rtp_transport.h"
#include "pc/rtp_transport_internal.h"
#include "rtc_base/logging.h"
#include "rtc_base/thread.h"
#include "test/create_test_field_trials.h"
#include "test/gtest.h"

namespace webrtc {
namespace {

FieldTrials g_field_trials = CreateTestFieldTrials();

std::unique_ptr<JsepTransport> MakeJsepTransport(const std::string& name) {
  auto ice = make_ref_counted<FakeIceTransport>(
      std::make_unique<FakeIceTransportInternal>(name,
                                                 ICE_CANDIDATE_COMPONENT_RTP));
  auto dtls = std::make_unique<FakeDtlsTransport>(ice);
  DtlsTransportInternal* dtls_ptr = dtls.get();

  auto rtp = std::make_unique<DtlsSrtpTransport>(/*rtcp_mux_enabled=*/true,
                                                 g_field_trials);
  rtp->SetDtlsTransportsOwned(std::move(dtls), nullptr);

  auto dtls_wrapper = make_ref_counted<DtlsTransport>(dtls_ptr);

  return std::make_unique<JsepTransport>(
      /*local_certificate=*/nullptr, std::move(rtp), std::move(dtls_wrapper),
      /*sctp_transport=*/nullptr,
      /*rtcp_mux_active_callback=*/[]() {});
}

TEST(JsepTransportCollectionRollbackUaf,
     ShortCircuitLeavesDanglingRtpTransport) {
  AutoThread main_thread;

  // Simulated per-mid BaseChannel::rtp_transport_ raw pointers
  // (pc/channel.h:326 — plain `RtpTransportInternal*`).
  std::map<std::string, RtpTransportInternal*> channel_rtp_transport;
  std::set<std::string> notified_during_rollback;

  bool rollback_phase = false;
  std::string fail_on_mid;

  auto map_change_callback = [&](absl::string_view mid,
                                 JsepTransport* t) -> bool {
    std::string m(mid);
    if (rollback_phase) {
      notified_during_rollback.insert(m);
      if (!fail_on_mid.empty() && m == fail_on_mid && t != nullptr) {
        // Simulate BaseChannel::SetRtpTransport when ConnectToRtpTransport_n
        // (RegisterRtpDemuxerSink) fails: rtp_transport_ ends up nullptr,
        // and `false` is propagated up the OnTransportChanged chain.
        channel_rtp_transport[m] = nullptr;
        return false;
      }
    }
    // Simulate BaseChannel storing the RtpTransportInternal* it was handed.
    channel_rtp_transport[m] = t ? t->rtp_transport() : nullptr;
    return true;
  };
  auto state_change_callback = []() {};

  JsepTransportCollection collection(map_change_callback,
                                     state_change_callback);

  // --- Reach stable state: bundle {A, B} on transport TA (named "A"). ---
  auto ta_owned = MakeJsepTransport("A");
  JsepTransport* ta = ta_owned.get();
  collection.RegisterTransport(std::move(ta_owned));   // A -> TA
  EXPECT_TRUE(collection.SetTransportForMid("B", ta));  // B -> TA
  collection.CommitTransports();  // stable_mid_to_transport_ = {A->TA, B->TA}

  // --- Apply remote offer: bundle group reordered to {C, A, B}. ---
  // A new mid "C" is BUNDLE-tagged; ApplyDescription_n would create a fresh
  // JsepTransport TC and move A/B onto it via SetTransportForMid.
  auto tc_owned = MakeJsepTransport("C");
  JsepTransport* tc = tc_owned.get();
  RtpTransportInternal* tc_rtp = tc->rtp_transport();
  collection.RegisterTransport(std::move(tc_owned));    // C -> TC
  EXPECT_TRUE(collection.SetTransportForMid("A", tc));  // A -> TC
  EXPECT_TRUE(collection.SetTransportForMid("B", tc));  // B -> TC

  ASSERT_EQ(channel_rtp_transport["A"], tc_rtp);
  ASSERT_EQ(channel_rtp_transport["B"], tc_rtp);

  // --- Rollback. ---
  // Arm the callback so that restoring mid "A" to its stable transport fails
  // (simulating RegisterRtpDemuxerSink -> CriteriaWouldConflict -> false).
  // Because of `ret = ret && callback(...)`, mid "B" is NEVER notified — yet
  // TC is destroyed unconditionally by DestroyUnusedTransports().
  rollback_phase = true;
  fail_on_mid = "A";

  bool ok = collection.RollbackTransports();
  EXPECT_FALSE(ok);

  // Prove the short-circuit happened.
  EXPECT_TRUE(notified_during_rollback.count("A"));
  EXPECT_TRUE(notified_during_rollback.count("C"));   // loop 1 (nullptr)
  EXPECT_FALSE(notified_during_rollback.count("B"))   // SKIPPED!
      << "mid B was notified — short-circuit did not occur";

  // The collection no longer references TC.
  EXPECT_EQ(collection.GetTransportForMid("C"), nullptr);
  EXPECT_EQ(collection.GetTransportByName("C"), nullptr);

  // Channel B's raw pointer was never updated; it still points at TC's
  // RtpTransportInternal, which was just freed.
  ASSERT_EQ(channel_rtp_transport["B"], tc_rtp);

  // *** Heap-use-after-free *** — this is exactly what
  // BaseChannel::DisconnectFromRtpTransport_n / SendPacket / transport_name()
  // would do on the next access. ASAN reports heap-use-after-free here.
  RTC_LOG(LS_ERROR) << "Dereferencing dangling rtp_transport_ for mid B: "
                    << channel_rtp_transport["B"]->transport_name();
}

}  // namespace
}  // namespace webrtc

// Minimal main() so this target builds inside a Chromium checkout without
// depending on //third_party/webrtc/test:test_main (which is gated on
// rtc_include_tests && !build_with_chromium).
int main(int argc, char** argv) {
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}