// 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 "content/public/browser/web_ui_url_loader_factory.h"

#include <optional>

#include "base/memory/raw_ptr.h"
#include "base/memory/ref_counted_memory.h"
#include "base/notreached.h"
#include "base/strings/strcat.h"
#include "base/strings/stringprintf.h"
#include "build/build_config.h"
#include "content/browser/webui/url_data_manager.h"
#include "content/public/browser/url_data_source.h"
#include "content/public/common/content_client.h"
#include "content/public/common/url_constants.h"
#include "content/public/test/test_browser_context.h"
#include "content/public/test/test_content_browser_client.h"
#include "content/public/test/test_content_client.h"
#include "content/public/test/test_renderer_host.h"
#include "mojo/public/cpp/bindings/remote.h"
#include "mojo/public/cpp/system/data_pipe_utils.h"
#include "net/traffic_annotation/network_traffic_annotation_test_helper.h"
#include "services/network/public/cpp/resource_request.h"
#include "services/network/public/mojom/url_loader.mojom.h"
#include "services/network/public/mojom/url_loader_factory.mojom.h"
#include "services/network/test/test_url_loader_client.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "url/url_util.h"

namespace content {

namespace {

const char* kTestWebUIScheme = kChromeUIScheme;
const char kNonChromeDummyScheme[] = "non-chrome";
constexpr char kTestWebUIHost[] = "testhost";
constexpr size_t kMaxTestResourceSize = 10;

// A URLDataSource that always returns the byte sequence:
// {0, 1, 2, ..., resource_size - 1}.
class TestWebUIDataSource final : public URLDataSource {
 public:
  static std::vector<unsigned char> GetResource(size_t size) {
    std::vector<unsigned char> resource(size);
    for (size_t i = 0; i < resource.size(); ++i)
      resource[i] = i;
    return resource;
  }

  explicit TestWebUIDataSource(size_t resource_size)
      : resource_size_(resource_size) {}

  // URLDataSource implementation:
  std::string GetSource() override { return kTestWebUIHost; }

  void StartDataRequest(
      const GURL& /*url*/,
      const content::WebContents::Getter& /*wc_getter*/,
      content::URLDataSource::GotDataCallback callback) override {
    std::move(callback).Run(base::MakeRefCounted<base::RefCountedBytes>(
        GetResource(resource_size_)));
  }

  std::string GetMimeType(const GURL& url) override { return "video/webm"; }

 private:
  const size_t resource_size_;
};

class OversizedWebUIDataSource final : public URLDataSource {
 public:
  explicit OversizedWebUIDataSource(size_t resource_size)
      : resource_size_(resource_size) {}

  // URLDataSource implementation:
  std::string GetSource() override { return kTestWebUIHost; }

  void StartDataRequest(
      const GURL& /*url*/,
      const content::WebContents::Getter& /*wc_getter*/,
      content::URLDataSource::GotDataCallback callback) override {
    auto resource =
        base::MakeRefCounted<OversizedRefCountedMemory>(resource_size_);
    std::move(callback).Run(std::move(resource));
  }

  std::string GetMimeType(const GURL& url) override { return "video/webm"; }

 private:
  // An implementation of RefCountedMemory representing a very large "virtual"
  // buffer that is never actually allocated.
  class OversizedRefCountedMemory final : public base::RefCountedMemory {
   public:
    explicit OversizedRefCountedMemory(size_t size) : size_(size) {}

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

    // base::RefCountedMemory implementation:
    base::span<const uint8_t> AsSpan() const override {
      // This uses `reinterpret_cast` of `1` to avoid nullness `CHECK` in the
      // constructor of `span`.
      //
      // SAFETY: This is unsound, but any use of the pointer will crash as the
      // first page is not mapped. The test does not actually use the pointer.
      return UNSAFE_BUFFERS(base::span<const uint8_t>(
          reinterpret_cast<const uint8_t*>(1), size_));
    }

   private:
    ~OversizedRefCountedMemory() override = default;

    const size_t size_;
  };

  const size_t resource_size_;
};

}  // namespace

const struct RangeRequestTestData {
  size_t resource_size = 0;
  std::optional<int> first_byte_position;
  std::optional<int> last_byte_position;
  int expected_error_code = net::OK;
  uint32_t expected_size = 0;
} kRangeRequestTestData[] = {
    // No range.
    {kMaxTestResourceSize, std::nullopt, std::nullopt, net::OK,
     kMaxTestResourceSize},

    // No range, 0-size resource.
    {0, std::nullopt, std::nullopt, net::OK, 0},

    {kMaxTestResourceSize, 3, std::nullopt, net::OK, kMaxTestResourceSize - 3},

    {kMaxTestResourceSize, 1, 1, net::OK, 1},

    // Range too large by 1, truncated to resource size.
    {kMaxTestResourceSize, 0, kMaxTestResourceSize, net::OK,
     kMaxTestResourceSize},

    // Range starts after the last resource byte.
    {kMaxTestResourceSize, kMaxTestResourceSize, kMaxTestResourceSize + 5,
     net::ERR_REQUEST_RANGE_NOT_SATISFIABLE, 0},

    // Invalid range.
    {kMaxTestResourceSize, 2, 1, net::ERR_REQUEST_RANGE_NOT_SATISFIABLE, 0},

#if defined(ARCH_CPU_64_BITS)
    // Resource too large.
    {static_cast<size_t>(std::numeric_limits<uint32_t>::max()) + 1,
     std::nullopt, std::nullopt, net::ERR_INSUFFICIENT_RESOURCES, 0},
#endif  // defined(ARCH_CPU_64_BITS)
};

class WebUIURLLoaderFactoryTest
    : public RenderViewHostTestHarness,
      public testing::WithParamInterface<RangeRequestTestData> {
 public:
  void SetUp() override {
    RenderViewHostTestHarness::SetUp();

    const auto resource_size = GetParam().resource_size;
    std::unique_ptr<URLDataSource> data_source;
    if (resource_size <= kMaxTestResourceSize)
      data_source = std::make_unique<TestWebUIDataSource>(resource_size);
    else
      data_source = std::make_unique<OversizedWebUIDataSource>(resource_size);
    URLDataManager::AddDataSource(browser_context(), std::move(data_source));
  }
};

INSTANTIATE_TEST_SUITE_P(,
                         WebUIURLLoaderFactoryTest,
                         testing::ValuesIn(kRangeRequestTestData));

TEST_P(WebUIURLLoaderFactoryTest, RangeRequest) {
  mojo::Remote<network::mojom::URLLoaderFactory> loader_factory(
      CreateWebUIURLLoaderFactory(main_rfh(), kTestWebUIScheme,
                                  /*allowed_hosts=*/{}));

  network::ResourceRequest request;
  request.url = GURL(base::StrCat({kTestWebUIScheme, "://", kTestWebUIHost}));

  if (GetParam().first_byte_position) {
    const std::string range_value =
        GetParam().last_byte_position
            ? base::StringPrintf("bytes=%d-%d", *GetParam().first_byte_position,
                                 *GetParam().last_byte_position)
            : base::StringPrintf("bytes=%d-", *GetParam().first_byte_position);
    request.headers.SetHeader(net::HttpRequestHeaders::kRange, range_value);
  }

  mojo::PendingRemote<network::mojom::URLLoader> loader;
  network::TestURLLoaderClient loader_client;
  loader_factory->CreateLoaderAndStart(
      loader.InitWithNewPipeAndPassReceiver(), /*request_id=*/0,
      /*options=*/0, request, loader_client.CreateRemote(),
      net::MutableNetworkTrafficAnnotationTag(TRAFFIC_ANNOTATION_FOR_TESTS));
  loader_client.RunUntilComplete();

  EXPECT_EQ(loader_client.completion_status().error_code,
            GetParam().expected_error_code);

  if (loader_client.completion_status().error_code == net::OK) {
    ASSERT_TRUE(loader_client.response_body().is_valid());
    std::string response;
    ASSERT_TRUE(mojo::BlockingCopyToString(
        loader_client.response_body_release(), &response));
    ASSERT_EQ(response.size(), GetParam().expected_size);

    std::vector<unsigned char> expected_resource =
        TestWebUIDataSource::GetResource(GetParam().resource_size);
    expected_resource.erase(
        expected_resource.begin(),
        expected_resource.begin() + GetParam().first_byte_position.value_or(0));
    expected_resource.resize(GetParam().expected_size);
    EXPECT_EQ(std::vector<uint8_t>(response.begin(), response.end()),
              expected_resource);
  }
}

class WebUIURLLoaderFactoryInvalidUrlTest
    : public RenderViewHostTestHarness,
      public testing::WithParamInterface<std::string> {
 public:
  void SetUp() override {
    RenderViewHostTestHarness::SetUp();
    browser_client_ = std::make_unique<InvalidUrlTestBrowserClient>();
    old_browser_client_ = SetBrowserClientForTesting(browser_client_.get());
  }

  void TearDown() override {
    SetBrowserClientForTesting(old_browser_client_);
    RenderViewHostTestHarness::TearDown();
  }

 private:
  class InvalidUrlTestBrowserClient : public TestContentBrowserClient {
   public:
    void GetAdditionalWebUISchemes(
        std::vector<std::string>* additional_schemes) override {
      additional_schemes->push_back(kNonChromeDummyScheme);
    }
  };

  std::unique_ptr<InvalidUrlTestBrowserClient> browser_client_;
  raw_ptr<ContentBrowserClient> old_browser_client_;
};

TEST_P(WebUIURLLoaderFactoryInvalidUrlTest, InvalidUrl) {
  mojo::Remote<network::mojom::URLLoaderFactory> loader_factory(
      CreateWebUIURLLoaderFactory(main_rfh(), kNonChromeDummyScheme,
                                  /*allowed_hosts=*/{}));

  network::ResourceRequest request;
  request.url = GURL(base::StrCat({kNonChromeDummyScheme, "://", GetParam()}));

  mojo::PendingRemote<network::mojom::URLLoader> loader;
  network::TestURLLoaderClient loader_client;
  loader_factory->CreateLoaderAndStart(
      loader.InitWithNewPipeAndPassReceiver(), /*request_id=*/0,
      /*options=*/0, request, loader_client.CreateRemote(),
      net::MutableNetworkTrafficAnnotationTag(TRAFFIC_ANNOTATION_FOR_TESTS));
  loader_client.RunUntilComplete();

  EXPECT_EQ(loader_client.completion_status().error_code, net::ERR_INVALID_URL);
}

// Test that non-chrome://blob-internals, non-chrome://dino and
// non-chrome://network-error/<xyz> are not reachable.
INSTANTIATE_TEST_SUITE_P(
    All,
    WebUIURLLoaderFactoryInvalidUrlTest,
    testing::Values(kChromeUIBlobInternalsHost,
                    kChromeUIDinoHost,
                    base::StrCat({kChromeUINetworkErrorHost, "/-147"})),
    [](const testing::TestParamInfo<std::string>& info) {
      std::string name = base::StrCat({kNonChromeDummyScheme, "_", info.param});
      std::replace_if(
          name.begin(), name.end(), [](char c) { return !std::isalnum(c); },
          '_');
      return name;
    });

}  // namespace content