// Copyright 2024 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "components/openscreen_platform/udp_socket.h"

#include <memory>
#include <optional>
#include <utility>

#include "base/test/task_environment.h"
#include "mojo/public/cpp/bindings/pending_receiver.h"
#include "mojo/public/cpp/bindings/pending_remote.h"
#include "mojo/public/cpp/bindings/remote.h"
#include "net/base/net_errors.h"
#include "services/network/public/mojom/udp_socket.mojom.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "third_party/openscreen/src/platform/api/udp_socket.h"
#include "third_party/openscreen/src/platform/base/error.h"
#include "third_party/openscreen/src/platform/base/ip_address.h"
#include "third_party/openscreen/src/platform/base/udp_packet.h"

namespace openscreen_platform {

namespace {

// A UdpSocket::Client that owns the socket and destroys it synchronously
// from within OnRead(). This models the Cast-receiver browser-process
// teardown chain where a recv error propagates through
//   Environment::OnRead -> Environment::OnError ->
//   ReceiverSession::OnSocketInvalid -> ... ->
//   StreamingRuntimeApplication::StopApplication() ->
//   receiver_session_client_.reset()
// which synchronously runs ~openscreen_platform::UdpSocket while
// UdpSocket::OnReceived is still on the stack.
class SelfDestructingClient : public openscreen::UdpSocket::Client {
 public:
  SelfDestructingClient() = default;
  ~SelfDestructingClient() override = default;

  void TakeOwnership(std::unique_ptr<openscreen::UdpSocket> socket) {
    socket_ = std::move(socket);
  }

  // openscreen::UdpSocket::Client:
  void OnError(openscreen::UdpSocket* socket,
               const openscreen::Error& error) override {}
  void OnSendError(openscreen::UdpSocket* socket,
                   const openscreen::Error& error) override {}
  void OnRead(openscreen::UdpSocket* socket,
              openscreen::ErrorOr<openscreen::UdpPacket> packet) override {
    // Synchronously free the socket while UdpSocket::OnReceived() is still on
    // the stack. After this returns, OnReceived() unconditionally executes
    //   udp_socket_->ReceiveMore(1);
    // which dereferences a destructed mojo::Remote<> stored inside the freed
    // |this| allocation.
    socket_.reset();
  }

 private:
  std::unique_ptr<openscreen::UdpSocket> socket_;
};

}  // namespace

// POC: heap-use-after-free in openscreen_platform::UdpSocket::OnReceived.
//
// UdpSocket::OnReceived (the network::mojom::UDPSocketListener handler) calls
// client_->OnRead(...) and then unconditionally touches |this->udp_socket_|.
// The Mojo listener dispatch is not weak-guarded (unlike Bind/JoinGroup/Send
// callbacks which use weak_ptr_factory_), so a Client that synchronously
// destroys the UdpSocket from within OnRead() leaves OnReceived() holding a
// dangling |this|. In production (Cast receiver products) the destruction is
// reached via StreamingRuntimeApplication::StopApplication; here we model that
// directly.
TEST(OpenscreenPlatformUdpSocketTest, OnReceivedAfterClientDeletesSocket) {
  base::test::TaskEnvironment task_environment;

  // Bound Remote so that, post-free, ReceiveMore(1) attempts a virtual call
  // through the destructed proxy (matching the production state).
  mojo::Remote<network::mojom::UDPSocket> udp_remote;
  std::ignore = udp_remote.BindNewPipeAndPassReceiver();

  mojo::PendingRemote<network::mojom::UDPSocketListener> listener_remote;
  mojo::PendingReceiver<network::mojom::UDPSocketListener> pending_listener =
      listener_remote.InitWithNewPipeAndPassReceiver();

  SelfDestructingClient client;
  auto socket = std::make_unique<UdpSocket>(
      &client,
      openscreen::IPEndpoint{openscreen::IPAddress{127, 0, 0, 1}, 0},
      std::move(udp_remote), std::move(pending_listener));
  UdpSocket* raw_socket = socket.get();
  client.TakeOwnership(std::move(socket));

  // Network process delivers OnReceived(result < 0, nullopt, nullopt) over the
  // UDPSocketListener pipe (services/network/udp_socket.cc forwards recv
  // errors unchanged). The listener_ receiver is not weak-guarded, so this is
  // a direct call into the implementation method:
  raw_socket->OnReceived(net::ERR_FAILED, std::nullopt, std::nullopt);
  // ^ heap-use-after-free under ASAN at udp_socket.cc:166:
  //     udp_socket_->ReceiveMore(1);
}

}  // namespace openscreen_platform