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

#include <optional>
#include <string_view>

#include "base/base64.h"
#include "base/files/file_path.h"
#include "base/files/file_util.h"
#include "base/files/scoped_temp_dir.h"
#include "base/functional/bind.h"
#include "base/run_loop.h"
#include "base/strings/utf_string_conversions.h"
#include "base/unguessable_token.h"
#include "build/build_config.h"
#include "chrome/browser/ui/browser.h"
#include "chrome/test/base/in_process_browser_test.h"
#include "chrome/test/base/ui_test_utils.h"
#include "components/paint_preview/browser/paint_preview_client.h"
#include "components/paint_preview/common/file_stream.h"
#include "components/paint_preview/common/mojom/paint_preview_recorder.mojom-shared.h"
#include "components/paint_preview/common/proto/paint_preview.pb.h"
#include "components/paint_preview/common/recording_map.h"
#include "components/paint_preview/common/serialized_recording.h"
#include "components/paint_preview/common/test_utils.h"
#include "components/ukm/test_ukm_recorder.h"
#include "content/public/browser/render_process_host.h"
#include "content/public/common/content_switches.h"
#include "content/public/test/browser_test.h"
#include "content/public/test/browser_test_utils.h"
#include "content/public/test/fenced_frame_test_util.h"
#include "net/dns/mock_host_resolver.h"
#include "net/test/embedded_test_server/embedded_test_server.h"
#include "services/metrics/public/cpp/ukm_builders.h"
#include "third_party/blink/public/common/associated_interfaces/associated_interface_provider.h"
#include "third_party/skia/include/core/SkPicture.h"
#include "third_party/skia/include/core/SkStream.h"
#include "url/gurl.h"

namespace paint_preview {

class NoOpPaintPreviewRecorder : public mojom::PaintPreviewRecorder {
 public:
  NoOpPaintPreviewRecorder() = default;
  ~NoOpPaintPreviewRecorder() override = default;

  NoOpPaintPreviewRecorder(const NoOpPaintPreviewRecorder&) = delete;
  NoOpPaintPreviewRecorder& operator=(const NoOpPaintPreviewRecorder&) = delete;

  void SetRequestedClosure(base::OnceClosure requested) {
    requested_ = std::move(requested);
  }

  void CapturePaintPreview(
      mojom::PaintPreviewCaptureParamsPtr params,
      mojom::PaintPreviewRecorder::CapturePaintPreviewCallback callback)
      override {
    callback_ = std::move(callback);
    std::move(requested_).Run();
  }

  void BindRequest(mojo::ScopedInterfaceEndpointHandle handle) {
    binding_.Bind(mojo::PendingAssociatedReceiver<mojom::PaintPreviewRecorder>(
        std::move(handle)));
  }

 private:
  base::OnceClosure requested_;
  mojom::PaintPreviewRecorder::CapturePaintPreviewCallback callback_;
  mojo::AssociatedReceiver<mojom::PaintPreviewRecorder> binding_{this};
};

// Test harness for a integration test of paint previews. In this test:
// - Each RenderFrame has an instance of PaintPreviewRecorder attached.
// - Each WebContents has an instance of PaintPreviewClient attached.
// This permits end-to-end testing of the flow of paint previews.
class PaintPreviewBrowserTest
    : public InProcessBrowserTest,
      public testing::WithParamInterface<RecordingPersistence> {
 public:
  PaintPreviewBrowserTest(const PaintPreviewBrowserTest&) = delete;
  PaintPreviewBrowserTest& operator=(const PaintPreviewBrowserTest&) = delete;

 protected:
  PaintPreviewBrowserTest() = default;
  ~PaintPreviewBrowserTest() override = default;

  void SetUp() override {
    ASSERT_TRUE(temp_dir_.CreateUniqueTempDir());

    InProcessBrowserTest::SetUp();
  }

  void SetUpOnMainThread() override {
    host_resolver()->AddRule("*", "127.0.0.1");

    http_server_.ServeFilesFromSourceDirectory("content/test/data");

    ASSERT_TRUE(http_server_.Start());
  }

  void SetUpCommandLine(base::CommandLine* command_line) override {
    content::IsolateAllSitesForTesting(command_line);
  }

  void CreateClient() {
    PaintPreviewClient::CreateForWebContents(GetWebContents());
  }

  content::WebContents* GetWebContents() {
    return browser()->tab_strip_model()->GetActiveWebContents();
  }

  void LoadPage(const GURL& url) const {
    ASSERT_TRUE(ui_test_utils::NavigateToURL(browser(), url));
  }

  void LoadHtml(std::string_view html) const {
    std::string base64_html = base::Base64Encode(html);
    GURL url(std::string("data:text/html;base64,") + base64_html);
    ASSERT_TRUE(ui_test_utils::NavigateToURL(browser(), url));
  }

  PaintPreviewClient::PaintPreviewParams MakeParams() const {
    PaintPreviewClient::PaintPreviewParams params(GetParam());
    params.inner.is_main_frame = true;
    params.root_dir = temp_dir_.GetPath();
    params.inner.capture_links = true;
    params.inner.max_capture_size = 0;
    return params;
  }

  void OverrideInterface(NoOpPaintPreviewRecorder* service,
                         content::RenderFrameHost* rfh) {
    blink::AssociatedInterfaceProvider* remote_interfaces =
        rfh->GetRemoteAssociatedInterfaces();
    remote_interfaces->OverrideBinderForTesting(
        mojom::PaintPreviewRecorder::Name_,
        base::BindRepeating(&NoOpPaintPreviewRecorder::BindRequest,
                            base::Unretained(service)));
  }

  void WaitForLoadStopWithoutSuccessCheck() {
    // In many cases, the load may have finished before we get here.  Only wait
    // if the tab still has a pending navigation.
    auto* web_contents = GetWebContents();
    if (web_contents->IsLoading()) {
      content::LoadStopObserver load_stop_observer(web_contents);
      load_stop_observer.Wait();
    }
  }

  // Check that |recording_map| contains the frame |frame_proto| and is a valid
  // SkPicture. Don't bother checking the contents as this is non-trivial and
  // could change. Instead check that the SkPicture can be read correctly and
  // has a cull rect of at least |size|.
  //
  // Consumes the recording from |recording_map|.
  static void EnsureSkPictureIsValid(RecordingMap* recording_map,
                                     const PaintPreviewFrameProto& frame_proto,
                                     size_t expected_subframe_count,
                                     const gfx::Size& size = gfx::Size(1, 1)) {
    base::ScopedAllowBlockingForTesting scoped_blocking;

    auto it = recording_map->find(
        base::UnguessableToken::Deserialize(frame_proto.embedding_token_high(),
                                            frame_proto.embedding_token_low())
            .value());
    ASSERT_NE(it, recording_map->end());

    std::optional<SkpResult> result = std::move(it->second).Deserialize();
    ASSERT_TRUE(result.has_value());
    EXPECT_NE(result->skp, nullptr);
    EXPECT_GE(result->skp->cullRect().width(), 0);
    EXPECT_GE(result->skp->cullRect().height(), 0);
    EXPECT_EQ(result->ctx.size(), expected_subframe_count);

    recording_map->erase(it);
  }

  base::ScopedTempDir temp_dir_;
  net::EmbeddedTestServer http_server_;
  net::EmbeddedTestServer http_server_different_origin_;
};

IN_PROC_BROWSER_TEST_P(PaintPreviewBrowserTest, CaptureFrame) {
  LoadPage(http_server_.GetURL("a.com", "/cross_site_iframe_factory.html?a"));
  ukm::TestAutoSetUkmRecorder ukm_recorder;
  auto params = MakeParams();

  base::RunLoop loop;

  CreateClient();
  auto* client = PaintPreviewClient::FromWebContents(GetWebContents());
  WaitForLoadStopWithoutSuccessCheck();
  client->CapturePaintPreview(
      params, GetWebContents()->GetPrimaryMainFrame(),
      base::BindOnce(
          [](base::RepeatingClosure quit,
             const PaintPreviewClient::PaintPreviewParams& params,
             base::UnguessableToken guid, mojom::PaintPreviewStatus status,
             std::unique_ptr<CaptureResult> result) {
            EXPECT_EQ(guid, params.inner.document_guid);
            EXPECT_EQ(status, mojom::PaintPreviewStatus::kOk);
            EXPECT_TRUE(result->proto.has_root_frame());
            EXPECT_EQ(result->proto.subframes_size(), 0);
            EXPECT_EQ(result->proto.root_frame()
                          .content_id_to_embedding_tokens_size(),
                      0);
            EXPECT_TRUE(result->proto.root_frame().is_main_frame());
            {
              base::ScopedAllowBlockingForTesting scoped_blocking;
              auto pair = RecordingMapFromCaptureResult(std::move(*result));
              EnsureSkPictureIsValid(&pair.first, pair.second.root_frame(), 0);
            }
            quit.Run();
          },
          loop.QuitClosure(), params));
  loop.Run();

  auto entries = ukm_recorder.GetEntriesByName(
      ukm::builders::PaintPreviewCapture::kEntryName);
  EXPECT_EQ(1u, entries.size());
}

IN_PROC_BROWSER_TEST_P(PaintPreviewBrowserTest,
                       CaptureMainFrameWithCrossProcessSubframe) {
  LoadPage(
      http_server_.GetURL("a.com", "/cross_site_iframe_factory.html?a(b)"));
  auto params = MakeParams();

  base::RunLoop loop;

  CreateClient();
  auto* client = PaintPreviewClient::FromWebContents(GetWebContents());
  WaitForLoadStopWithoutSuccessCheck();
  client->CapturePaintPreview(
      params, GetWebContents()->GetPrimaryMainFrame(),
      base::BindOnce(
          [](base::RepeatingClosure quit,
             const PaintPreviewClient::PaintPreviewParams& params,
             base::UnguessableToken guid, mojom::PaintPreviewStatus status,
             std::unique_ptr<CaptureResult> result) {
            EXPECT_EQ(guid, params.inner.document_guid);
            EXPECT_EQ(status, mojom::PaintPreviewStatus::kOk);
            EXPECT_TRUE(result->proto.has_root_frame());
            EXPECT_EQ(result->proto.subframes_size(), 1);
            EXPECT_EQ(result->proto.root_frame()
                          .content_id_to_embedding_tokens_size(),
                      1);
            EXPECT_TRUE(result->proto.root_frame().is_main_frame());
            EXPECT_EQ(result->proto.subframes(0)
                          .content_id_to_embedding_tokens_size(),
                      0);
            EXPECT_FALSE(result->proto.subframes(0).is_main_frame());
            {
              base::ScopedAllowBlockingForTesting scoped_blocking;
              auto pair = RecordingMapFromCaptureResult(std::move(*result));
              EnsureSkPictureIsValid(&pair.first, pair.second.root_frame(), 1);
              EnsureSkPictureIsValid(&pair.first, pair.second.subframes(0), 0);
            }
            quit.Run();
          },
          loop.QuitClosure(), params));
  loop.Run();
}

class PaintPreviewFencedFrameBrowserTest : public PaintPreviewBrowserTest {
 public:
  PaintPreviewFencedFrameBrowserTest() = default;
  ~PaintPreviewFencedFrameBrowserTest() override = default;

  content::test::FencedFrameTestHelper& fenced_frame_test_helper() {
    return fenced_frame_test_helper_;
  }

 private:
  content::test::FencedFrameTestHelper fenced_frame_test_helper_;
};

IN_PROC_BROWSER_TEST_P(PaintPreviewFencedFrameBrowserTest,
                       CaptureMainFrameWithCrossProcessFencedFrames) {
  LoadPage(http_server_.GetURL("a.com", "/title1.html"));
  content::RenderFrameHost* primary_main_rfh =
      GetWebContents()->GetPrimaryMainFrame();

  // Create two fenced frames.
  fenced_frame_test_helper().CreateFencedFrame(
      primary_main_rfh,
      http_server_.GetURL("b.com", "/fenced_frames/title1.html"));
  fenced_frame_test_helper().CreateFencedFrame(
      primary_main_rfh,
      http_server_.GetURL("c.com", "/fenced_frames/title1.html"));

  base::RunLoop finished_loop;
  CreateClient();
  auto* client = PaintPreviewClient::FromWebContents(GetWebContents());
  auto params = MakeParams();

  client->CapturePaintPreview(
      params, primary_main_rfh,
      base::BindOnce(
          [](const PaintPreviewClient::PaintPreviewParams& params,
             base::UnguessableToken guid, mojom::PaintPreviewStatus status,
             std::unique_ptr<CaptureResult> result) {
            // This callback should have a success result without any DCHECK
            // error.
            EXPECT_EQ(guid, params.inner.document_guid);
            EXPECT_EQ(status, mojom::PaintPreviewStatus::kOk);
            EXPECT_TRUE(result->proto.has_root_frame());
            EXPECT_EQ(result->proto.subframes_size(), 2);
            EXPECT_EQ(result->proto.root_frame()
                          .content_id_to_embedding_tokens_size(),
                      2);
            EXPECT_TRUE(result->proto.root_frame().is_main_frame());
            EXPECT_FALSE(result->proto.subframes(0).is_main_frame());
            EXPECT_FALSE(result->proto.subframes(1).is_main_frame());
            {
              base::ScopedAllowBlockingForTesting scoped_blocking;
              auto pair = RecordingMapFromCaptureResult(std::move(*result));
              EnsureSkPictureIsValid(&pair.first, pair.second.root_frame(), 2);
              EnsureSkPictureIsValid(&pair.first, pair.second.subframes(0), 0);
              EnsureSkPictureIsValid(&pair.first, pair.second.subframes(1), 0);
            }
          },
          params)
          .Then(finished_loop.QuitClosure()));
  finished_loop.Run();
}

IN_PROC_BROWSER_TEST_P(PaintPreviewFencedFrameBrowserTest,
                       DoNotAffectAnotherFrameWhenRemovingFencedFrame) {
  base::ScopedAllowBlockingForTesting scope;

  LoadPage(http_server_.GetURL("a.com", "/title1.html"));
  content::RenderFrameHost* primary_main_rfh =
      GetWebContents()->GetPrimaryMainFrame();

  // Create two fenced frames.
  content::RenderFrameHostWrapper fenced_rfh_wrapper(
      fenced_frame_test_helper().CreateFencedFrame(
          primary_main_rfh,
          http_server_.GetURL("b.com", "/fenced_frames/title1.html")));
  fenced_frame_test_helper().CreateFencedFrame(
      primary_main_rfh,
      http_server_.GetURL("c.com", "/fenced_frames/title1.html"));

  // Override remote interfaces of the fenced frame with a no-op.
  base::RunLoop started_loop;
  NoOpPaintPreviewRecorder noop_recorder;
  noop_recorder.SetRequestedClosure(started_loop.QuitClosure());

  OverrideInterface(&noop_recorder, fenced_rfh_wrapper.get());

  base::RunLoop finished_loop;
  CreateClient();
  auto* client = PaintPreviewClient::FromWebContents(GetWebContents());
  auto params = MakeParams();

  client->CapturePaintPreview(
      params, primary_main_rfh,
      base::BindOnce(
          [](const PaintPreviewClient::PaintPreviewParams& params,
             base::UnguessableToken guid, mojom::PaintPreviewStatus status,
             std::unique_ptr<CaptureResult> result) {
            // This callback should have a partial success result since the
            // fenced frame has been removed during running the capture.
            EXPECT_EQ(guid, params.inner.document_guid);
            EXPECT_EQ(status, mojom::PaintPreviewStatus::kPartialSuccess);
            EXPECT_TRUE(result->proto.has_root_frame());
            EXPECT_EQ(result->proto.subframes_size(), 1);
            EXPECT_EQ(result->proto.root_frame()
                          .content_id_to_embedding_tokens_size(),
                      2);
            EXPECT_TRUE(result->proto.root_frame().is_main_frame());
            EXPECT_EQ(result->proto.subframes(0)
                          .content_id_to_embedding_tokens_size(),
                      0);
            EXPECT_FALSE(result->proto.subframes(0).is_main_frame());
            {
              base::ScopedAllowBlockingForTesting scoped_blocking;
              auto pair = RecordingMapFromCaptureResult(std::move(*result));
              EnsureSkPictureIsValid(&pair.first, pair.second.root_frame(), 2);
              EnsureSkPictureIsValid(&pair.first, pair.second.subframes(0), 0);
            }
          },
          params)
          .Then(finished_loop.QuitClosure()));

  // Wait for the request to execute before removing the fenced frame.
  started_loop.Run();

  // Remove the fenced frame.
  EXPECT_TRUE(
      ExecJs(primary_main_rfh,
             "const ff = document.querySelector('fencedframe'); ff.remove();"));
  ASSERT_TRUE(fenced_rfh_wrapper.WaitUntilRenderFrameDeleted());

  finished_loop.Run();
}

IN_PROC_BROWSER_TEST_P(PaintPreviewBrowserTest,
                       CaptureMainFrameWithScrollableSameProcessSubframe) {
  std::string html = R"(<html>
                          <iframe
                            srcdoc="<div
                                      style='width: 300px;
                                             height: 300px;
                                             background-color: #ff0000'>
                                      &nbsp;
                                    </div>"
                            title="subframe"
                            width="100px"
                            height="100px">
                          </iframe>
                        </html>)";
  LoadHtml(html);
  auto params = MakeParams();

  base::RunLoop loop;
  CreateClient();
  auto* client = PaintPreviewClient::FromWebContents(GetWebContents());
  WaitForLoadStopWithoutSuccessCheck();
  client->CapturePaintPreview(
      params, GetWebContents()->GetPrimaryMainFrame(),
      base::BindOnce(
          [](base::RepeatingClosure quit,
             const PaintPreviewClient::PaintPreviewParams& params,
             base::UnguessableToken guid, mojom::PaintPreviewStatus status,
             std::unique_ptr<CaptureResult> result) {
            EXPECT_EQ(guid, params.inner.document_guid);
            EXPECT_EQ(status, mojom::PaintPreviewStatus::kOk);
            EXPECT_TRUE(result->proto.has_root_frame());
            EXPECT_EQ(result->proto.subframes_size(), 1);
            EXPECT_EQ(result->proto.root_frame()
                          .content_id_to_embedding_tokens_size(),
                      1);
            EXPECT_TRUE(result->proto.root_frame().is_main_frame());
            EXPECT_EQ(result->proto.subframes(0)
                          .content_id_to_embedding_tokens_size(),
                      0);
            EXPECT_FALSE(result->proto.subframes(0).is_main_frame());
            {
              base::ScopedAllowBlockingForTesting scoped_blocking;
              auto pair = RecordingMapFromCaptureResult(std::move(*result));
              EnsureSkPictureIsValid(&pair.first, pair.second.root_frame(), 1);
              EnsureSkPictureIsValid(&pair.first, pair.second.subframes(0), 0,
                                     gfx::Size(300, 300));
            }
            quit.Run();
          },
          loop.QuitClosure(), params));
  loop.Run();
}

IN_PROC_BROWSER_TEST_P(PaintPreviewBrowserTest,
                       CaptureMainFrameWithNonScrollableSameProcessSubframe) {
  std::string html = R"(<html>
                          <iframe
                            srcdoc="<div
                                      style='width: 50px;
                                             height: 50px;
                                             background-color: #ff0000'>
                                      &nbsp;
                                    </div>"
                            title="subframe"
                            width="100px"
                            height="100px">
                          </iframe>
                        </html>)";
  LoadHtml(html);
  auto params = MakeParams();

  base::RunLoop loop;
  CreateClient();
  auto* client = PaintPreviewClient::FromWebContents(GetWebContents());
  WaitForLoadStopWithoutSuccessCheck();
  client->CapturePaintPreview(
      params, GetWebContents()->GetPrimaryMainFrame(),
      base::BindOnce(
          [](base::RepeatingClosure quit,
             const PaintPreviewClient::PaintPreviewParams& params,
             base::UnguessableToken guid, mojom::PaintPreviewStatus status,
             std::unique_ptr<CaptureResult> result) {
            EXPECT_EQ(guid, params.inner.document_guid);
            EXPECT_EQ(status, mojom::PaintPreviewStatus::kOk);
            EXPECT_TRUE(result->proto.has_root_frame());
            EXPECT_EQ(result->proto.subframes_size(), 0);
            EXPECT_EQ(result->proto.root_frame()
                          .content_id_to_embedding_tokens_size(),
                      0);
            EXPECT_TRUE(result->proto.root_frame().is_main_frame());
            {
              base::ScopedAllowBlockingForTesting scoped_blocking;
              auto pair = RecordingMapFromCaptureResult(std::move(*result));
              EnsureSkPictureIsValid(&pair.first, pair.second.root_frame(), 0);
            }
            quit.Run();
          },
          loop.QuitClosure(), params));
  loop.Run();
}

// https://crbug.com/1146573 reproduction. If a renderer crashes,
// WebContentsObserver::RenderFrameDeleted. Paint preview implements this in an
// observer which in turn releases the capture handle which can cause the
// WebContents to be reloaded on Android where we have auto-reload. This reload
// occurs *during* crash handling, leaving the frame in an invalid state and
// leading to a crash when it subsequently unloaded.
// This is fixed by deferring it to a PostTask.
IN_PROC_BROWSER_TEST_P(PaintPreviewBrowserTest, DontReloadInRenderProcessExit) {
  // In the FileSystem variant of this test, blocking needs to be permitted to
  // allow cleanup to work during the crash.
  base::ScopedAllowBlockingForTesting scope;
  LoadPage(http_server_.GetURL("a.com", "/title1.html"));

  content::WebContents* web_contents = GetWebContents();

  // Override remote interfaces with a no-op.
  base::RunLoop started_loop;
  NoOpPaintPreviewRecorder noop_recorder;
  noop_recorder.SetRequestedClosure(started_loop.QuitClosure());
  OverrideInterface(&noop_recorder, GetWebContents()->GetPrimaryMainFrame());

  CreateClient();
  auto* client = PaintPreviewClient::FromWebContents(web_contents);
  // Do this twice to simulate conditions for crash.
  auto handle1 = web_contents->IncrementCapturerCount(
      gfx::Size(), /*stay_hidden=*/true,
      /*stay_awake=*/true, /*is_activity=*/true);
  auto handle2 = web_contents->IncrementCapturerCount(
      gfx::Size(), /*stay_hidden=*/true,
      /*stay_awake=*/true, /*is_activity=*/true);

  // A callback that causes the frame to reload and end up in an invalid state
  // if it is allowed to run during crash handling.
  base::RunLoop finished_loop;
  auto params = MakeParams();
  bool did_run = false;
  client->CapturePaintPreview(
      params, web_contents->GetPrimaryMainFrame(),
      // This callback is now posted so it shouldn't cause a crash.
      base::BindOnce(
          [](content::WebContents* web_contents, bool* did_run_ptr,
             base::ScopedClosureRunner handle1,
             base::ScopedClosureRunner handle2, base::UnguessableToken guid,
             mojom::PaintPreviewStatus status,
             std::unique_ptr<CaptureResult> result) {
            EXPECT_EQ(status, mojom::PaintPreviewStatus::kFailed);
            EXPECT_EQ(result, nullptr);
            // On Android crashed frames are marked as needing reload.
            web_contents->GetController().SetNeedsReload();
            handle1.RunAndReset();
            handle2.RunAndReset();
            *did_run_ptr = true;
          },
          web_contents, &did_run, std::move(handle1), std::move(handle2))
          .Then(finished_loop.QuitClosure()));
  // Wait for the request to execute before crashing the renderer. Otherwise in
  // the FileSystem variant it is possible there will be a race during creation
  // of the file with the renderer crash. If this happens the callback for
  // `finished_loop` will not be run as no request to capture succeeded leading
  // to a timeout.
  started_loop.Run();

  // Crash the renderer.
  content::RenderProcessHost* process =
      GetWebContents()->GetPrimaryMainFrame()->GetProcess();
  content::RenderProcessHostWatcher crash_observer(
      process, content::RenderProcessHostWatcher::WATCH_FOR_PROCESS_EXIT);
  process->Shutdown(0);
  crash_observer.Wait();

  // The browser would have crashed before the loop exited if the callback was
  // not posted.
  if (!did_run)
    finished_loop.Run();

  // Now navigate away and ensure that the frame unloads successfully.
  LoadPage(http_server_.GetURL("a.com", "/title2.html"));
}

INSTANTIATE_TEST_SUITE_P(All,
                         PaintPreviewBrowserTest,
                         testing::Values(RecordingPersistence::kFileSystem,
                                         RecordingPersistence::kMemoryBuffer),
                         PersistenceParamToString);
INSTANTIATE_TEST_SUITE_P(All,
                         PaintPreviewFencedFrameBrowserTest,
                         testing::Values(RecordingPersistence::kFileSystem,
                                         RecordingPersistence::kMemoryBuffer),
                         PersistenceParamToString);
}  // namespace paint_preview