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

#include "ui/ozone/platform/headless/headless_window.h"

#include <memory>

#include "testing/gtest/include/gtest/gtest.h"
#include "ui/display/test/test_screen.h"
#include "ui/gfx/geometry/rect.h"
#include "ui/ozone/platform/headless/headless_window_manager.h"
#include "ui/platform_window/platform_window_delegate.h"

namespace ui {
namespace {

// A PlatformWindowDelegate that simulates the production behaviour of
// DesktopWindowTreeHostPlatform: when the platform window calls back into the
// delegate (OnBoundsChanged / OnWindowStateChanged), an observer may
// synchronously call Widget::CloseNow(), which ends up calling
// HeadlessWindow::Close() -> delegate_->OnClosed() ->
// DWTHP::OnClosed() -> SetPlatformWindow(nullptr) -> ~HeadlessWindow().
//
// This delegate models that by resetting its owned unique_ptr<PlatformWindow>
// from inside the chosen callback.
class DestroyingDelegate : public PlatformWindowDelegate {
 public:
  enum class DestroyOn {
    kBoundsChanged,
    kWindowStateChanged,
  };

  explicit DestroyingDelegate(DestroyOn destroy_on)
      : destroy_on_(destroy_on) {}
  ~DestroyingDelegate() override = default;

  void SetWindow(std::unique_ptr<PlatformWindow> window) {
    window_ = std::move(window);
  }

  // PlatformWindowDelegate:
  void OnBoundsChanged(const BoundsChange& change) override {
    if (destroy_on_ == DestroyOn::kBoundsChanged) {
      // Simulates: WindowTreeHostPlatform::OnBoundsChanged ->
      //   OnHostResizedInPixels -> WidgetObserver::OnWidgetBoundsChanged ->
      //   Widget::CloseNow() -> DWTHP::CloseNow() -> platform_window()->Close()
      //   -> delegate_->OnClosed() -> DWTHP::OnClosed() ->
      //   SetPlatformWindow(nullptr) -> ~HeadlessWindow().
      window_.reset();
    }
  }
  void OnDamageRect(const gfx::Rect& damaged_region) override {}
  void DispatchEvent(Event* event) override {}
  void OnCloseRequest() override {}
  void OnClosed() override { window_.reset(); }
  void OnWindowStateChanged(PlatformWindowState old_state,
                            PlatformWindowState new_state) override {
    if (destroy_on_ == DestroyOn::kWindowStateChanged) {
      // Simulates: DWTHP::OnWindowStateChanged ->
      //   Widget::OnNativeWidgetWindowShowStateChanged ->
      //   WidgetObserver::OnWidgetShowStateChanged -> Widget::CloseNow() ->
      //   ... -> SetPlatformWindow(nullptr) -> ~HeadlessWindow().
      window_.reset();
    }
  }
  void OnLostCapture() override {}
  void OnAcceleratedWidgetAvailable(gfx::AcceleratedWidget widget) override {}
  void OnWillDestroyAcceleratedWidget() override {}
  void OnAcceleratedWidgetDestroyed() override {}
  void OnActivationChanged(bool active) override {}
  void OnCursorUpdate() override {}
  bool CanMaximize() const override { return true; }
  bool CanFullscreen() const override { return true; }
  gfx::Rect ConvertRectToPixels(const gfx::Rect& rect_in_dip) const override {
    return rect_in_dip;
  }
  gfx::Rect ConvertRectToDIP(const gfx::Rect& rect_in_pixels) const override {
    return rect_in_pixels;
  }

 private:
  const DestroyOn destroy_on_;
  std::unique_ptr<PlatformWindow> window_;
};

// HeadlessWindow::SetFullscreen() calls ZoomWindowBounds() (which dispatches
// delegate_->OnBoundsChanged()) followed by UpdateWindowState() with no
// liveness check between them. If the OnBoundsChanged() dispatch synchronously
// destroys |this| (as DesktopWindowTreeHostPlatform::OnClosed() does via
// SetPlatformWindow(nullptr)), UpdateWindowState() reads and writes
// |window_state_| on the freed object and dispatches a virtual call through the
// freed |delegate_| slot.
//
// Under ASAN this test produces a heap-use-after-free at
// HeadlessWindow::UpdateWindowState().
TEST(HeadlessWindowUAFTest, SetFullscreenAfterDestroyInOnBoundsChanged) {
  display::test::TestScreen test_screen(/*create_display=*/true,
                                        /*register_screen=*/true);
  HeadlessWindowManager manager;
  DestroyingDelegate delegate(DestroyingDelegate::DestroyOn::kBoundsChanged);

  auto window = std::make_unique<HeadlessWindow>(&delegate, &manager,
                                                 gfx::Rect(0, 0, 100, 100));
  HeadlessWindow* raw_window = window.get();
  delegate.SetWindow(std::move(window));

  // Web-reachable entry point: element.requestFullscreen() ->
  // Widget::SetFullscreen() -> DWTHP::SetFullscreen() ->
  // HeadlessWindow::SetFullscreen().
  raw_window->SetFullscreen(/*fullscreen=*/true,
                            /*target_display_id=*/-1);
}

// HeadlessWindow::Minimize() calls UpdateWindowState() (which dispatches
// delegate_->OnWindowStateChanged()) followed by Deactivate() with no liveness
// check between them. If the OnWindowStateChanged() dispatch synchronously
// destroys |this|, Deactivate() reads and writes |activation_state_| on the
// freed object and dispatches delegate_->OnActivationChanged() through the
// freed |delegate_| slot.
//
// Under ASAN this test produces a heap-use-after-free at
// HeadlessWindow::Deactivate().
TEST(HeadlessWindowUAFTest, MinimizeAfterDestroyInOnWindowStateChanged) {
  HeadlessWindowManager manager;
  DestroyingDelegate delegate(
      DestroyingDelegate::DestroyOn::kWindowStateChanged);

  auto window = std::make_unique<HeadlessWindow>(&delegate, &manager,
                                                 gfx::Rect(0, 0, 100, 100));
  HeadlessWindow* raw_window = window.get();
  delegate.SetWindow(std::move(window));

  raw_window->Minimize();
}

}  // namespace
}  // namespace ui