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

#include "base/test/bind.h"
#include "content/public/browser/navigation_handle_user_data.h"

#include "content/browser/renderer_host/navigation_request.h"
#include "content/browser/renderer_host/render_frame_host_impl.h"
#include "content/browser/web_contents/web_contents_impl.h"
#include "content/public/test/browser_test.h"
#include "content/public/test/browser_test_utils.h"
#include "content/public/test/content_browser_test.h"
#include "content/public/test/content_browser_test_utils.h"
#include "content/public/test/web_contents_observer_test_utils.h"
#include "content/shell/browser/shell.h"
#include "net/dns/mock_host_resolver.h"
#include "net/test/embedded_test_server/embedded_test_server.h"

namespace content {

namespace {
class Data : public NavigationHandleUserData<Data> {
 public:
  ~Data() override = default;

  base::WeakPtr<Data> GetWeakPtr() { return weak_ptr_factory_.GetWeakPtr(); }

  const std::string& value() const { return value_; }

 private:
  Data(NavigationHandle& navigation, std::string value) : value_(value) {}
  friend class NavigationHandleUserData<Data>;

  std::string value_;

  base::WeakPtrFactory<Data> weak_ptr_factory_{this};

  NAVIGATION_HANDLE_USER_DATA_KEY_DECL();
};

NAVIGATION_HANDLE_USER_DATA_KEY_IMPL(Data);

}  // namespace

class NavigationHandleUserDataBrowserTest : public ContentBrowserTest {
 public:
  ~NavigationHandleUserDataBrowserTest() override = default;

 protected:
  void SetUpOnMainThread() override {
    host_resolver()->AddRule("*", "127.0.0.1");
    ContentBrowserTest::SetUpOnMainThread();
  }

  WebContentsImpl* web_contents() const {
    return static_cast<WebContentsImpl*>(shell()->web_contents());
  }

  RenderFrameHostImpl* top_frame_host() {
    return web_contents()->GetPrimaryFrameTree().root()->current_frame_host();
  }
};

IN_PROC_BROWSER_TEST_F(NavigationHandleUserDataBrowserTest,
                       SuccessfulNavigation) {
  ASSERT_TRUE(embedded_test_server()->Start());
  GURL url(embedded_test_server()->GetURL("a.com", "/title1.html"));

  TestNavigationManager navigation_manager(web_contents(), url);

  // Start a navigation.
  shell()->LoadURL(url);
  EXPECT_TRUE(navigation_manager.WaitForRequestStart());

  // Attach some data to the navigation.
  Data::CreateForNavigationHandle(*navigation_manager.GetNavigationHandle(),
                                  "data");

  // Check that we can retrieve the data.
  base::WeakPtr<Data> data =
      Data::GetForNavigationHandle(*navigation_manager.GetNavigationHandle())
          ->GetWeakPtr();
  ASSERT_TRUE(data);
  EXPECT_EQ(data->value(), "data");

  // Wait for the navigation to finish.
  bool did_finish_navigation = false;
  NavigationFinishObserver dfn(
      web_contents(),
      base::BindLambdaForTesting([&](NavigationHandle* navigation) {
        EXPECT_TRUE(navigation->HasCommitted());
        ASSERT_TRUE(Data::GetForNavigationHandle(*navigation));
        ASSERT_TRUE(data);
        EXPECT_EQ(data->value(), "data");
        did_finish_navigation = true;
      }));
  ASSERT_TRUE(navigation_manager.WaitForNavigationFinished());
  EXPECT_TRUE(did_finish_navigation);

  // Ensure that the data is deleted after navigation finished.
  EXPECT_FALSE(data);
}

IN_PROC_BROWSER_TEST_F(NavigationHandleUserDataBrowserTest, FailedNavigation) {
  ASSERT_TRUE(embedded_test_server()->Start());
  GURL url(embedded_test_server()->GetURL("a.com", "/title1.html"));

  TestNavigationManager navigation_manager(web_contents(), url);

  // Start a navigation.
  shell()->LoadURL(url);
  EXPECT_TRUE(navigation_manager.WaitForRequestStart());

  // Attach some data to the navigation.
  Data::CreateForNavigationHandle(*navigation_manager.GetNavigationHandle(),
                                  "data");

  // Check that we can retrieve the data.
  base::WeakPtr<Data> data =
      Data::GetForNavigationHandle(*navigation_manager.GetNavigationHandle())
          ->GetWeakPtr();
  ASSERT_TRUE(data);
  EXPECT_EQ(data->value(), "data");

  // Wait for DidFinishNavigation event.
  bool did_finish_navigation = false;
  NavigationFinishObserver dfn(
      web_contents(),
      base::BindLambdaForTesting([&](NavigationHandle* navigation) {
        EXPECT_TRUE(!navigation->HasCommitted());
        ASSERT_TRUE(Data::GetForNavigationHandle(*navigation));
        ASSERT_TRUE(data);
        EXPECT_EQ(data->value(), "data");
        did_finish_navigation = true;
      }));
  shell()->Stop();
  EXPECT_TRUE(did_finish_navigation);

  // Ensure that the data is deleted after navigation finished.
  EXPECT_FALSE(data);
}

IN_PROC_BROWSER_TEST_F(NavigationHandleUserDataBrowserTest, VeryEarlyAttach) {
  ASSERT_TRUE(embedded_test_server()->Start());
  GURL url(embedded_test_server()->GetURL("a.com", "/title1.html"));

  TestNavigationManager navigation_manager(web_contents(), url);

  int64_t navigation_id;
  base::WeakPtr<Data> data;

  // Expect that the data will be available inside DidStartNavigation and
  // DidFinishNavigation.
  bool did_start_navigation = false;
  NavigationStartObserver dsn(
      web_contents(),
      base::BindLambdaForTesting([&](NavigationHandle* navigation) {
        // Attach some data to the navigation.
        Data::CreateForNavigationHandle(*navigation, "data");
        ASSERT_TRUE(Data::GetForNavigationHandle(*navigation));
        EXPECT_EQ(Data::GetForNavigationHandle(*navigation)->value(), "data");
        did_start_navigation = true;
      }));

  bool did_finish_navigation = false;
  NavigationFinishObserver dfn(
      web_contents(),
      base::BindLambdaForTesting([&](NavigationHandle* navigation) {
        EXPECT_EQ(navigation_id, navigation->GetNavigationId());
        ASSERT_TRUE(Data::GetForNavigationHandle(*navigation));
        ASSERT_TRUE(data);
        EXPECT_EQ(data->value(), "data");
        did_finish_navigation = true;
      }));

  // Start a navigation.
  auto* web_contents_impl =
      static_cast<WebContentsImpl*>(shell()->web_contents());
  web_contents_impl->GetController().LoadURLWithParams(
      NavigationController::LoadURLParams(url));
  auto* navigation =
      web_contents_impl->GetPrimaryFrameTree().root()->navigation_request();
  ASSERT_TRUE(navigation);
  navigation_id = navigation->GetNavigationId();

  // Check that we can retrieve the data.
  data = Data::GetForNavigationHandle(*navigation_manager.GetNavigationHandle())
             ->GetWeakPtr();
  ASSERT_TRUE(data);
  EXPECT_EQ(data->value(), "data");

  // Wait for the navigation to finish.
  ASSERT_TRUE(navigation_manager.WaitForNavigationFinished());
  EXPECT_TRUE(did_start_navigation);
  EXPECT_TRUE(did_finish_navigation);

  // Ensure that the data is deleted after navigation finished.
  EXPECT_FALSE(data);
}

}  // namespace content