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

#include "chrome/browser/ui/views/drive_picker_host/drive_picker_sanitizer.h"

#include "chrome/browser/ui/views/drive_picker_host/drive_picker_result_handler.mojom.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "url/gurl.h"

namespace {

drive_picker_host::mojom::DriveFilePtr CreateValidFile() {
  auto file = drive_picker_host::mojom::DriveFile::New();
  file->id = "valid-id_123";
  file->mime_type = "application/pdf";
  file->name = "test.pdf";
  file->type = "document";
  file->size_bytes = 1024;
  return file;
}

}  // namespace

class DrivePickerSanitizerTest : public testing::Test {};

TEST_F(DrivePickerSanitizerTest, SanitizeValidDocument) {
  auto file = CreateValidFile();
  file->resource_key = "valid-key";

  auto sanitized = DrivePickerSanitizer::Sanitize(file);

  ASSERT_TRUE(sanitized.has_value());
  EXPECT_EQ(sanitized->drive_id, "valid-id_123");
  EXPECT_EQ(sanitized->mime_type, "application/pdf");
  EXPECT_EQ(sanitized->file_name, "test.pdf");
  EXPECT_EQ(sanitized->size_bytes, 1024u);
  EXPECT_EQ(sanitized->resource_key, "valid-key");
  EXPECT_FALSE(sanitized->thumbnail_url.has_value());
}

TEST_F(DrivePickerSanitizerTest, SanitizeValidFile) {
  auto file = CreateValidFile();
  file->type = "file";
  file->mime_type = "application/zip";
  file->name = "test.zip";
  file->resource_key = "valid-key";

  auto sanitized = DrivePickerSanitizer::Sanitize(file);

  ASSERT_TRUE(sanitized.has_value());
  EXPECT_EQ(sanitized->drive_id, "valid-id_123");
  EXPECT_EQ(sanitized->mime_type, "application/zip");
  EXPECT_EQ(sanitized->file_name, "test.zip");
  EXPECT_EQ(sanitized->size_bytes, 1024u);
  EXPECT_EQ(sanitized->resource_key, "valid-key");
  EXPECT_FALSE(sanitized->thumbnail_url.has_value());
}

TEST_F(DrivePickerSanitizerTest, SanitizeValidPhotoWithThumbnail) {
  auto file = CreateValidFile();
  file->type = "photo";
  GURL valid_url(
      "https://lh3.googleusercontent.com/drive-storage/"
      "AJQWtBOI_xcNPWAVolLfcNDZGWus_ELL8GP-ZYLdkxOcPlUxsVOMr5jn-"
      "i261zxtYVcpA0kZeP"
      "sVWT2Ghrxdg03aoJYX7t9vlc4ojecNyW7QNrP6muY9-"
      "wvmO77SsiXHrwnU2GbRCt2V9NDv62u2"
      "R96rpvI=s220");
  file->thumbnail_url = valid_url;

  auto sanitized = DrivePickerSanitizer::Sanitize(file);

  ASSERT_TRUE(sanitized.has_value());
  ASSERT_TRUE(sanitized->thumbnail_url.has_value());
  EXPECT_EQ(sanitized->thumbnail_url.value(), valid_url);
}

TEST_F(DrivePickerSanitizerTest, SanitizeValidPhotoWithDThumbnail) {
  auto file = CreateValidFile();
  file->type = "photo";
  GURL valid_url(
      "https://lh3.googleusercontent.com/d/"
      "1aBcD2eFgH3iJkL4mN-oPqR5sTuvWxYz0=s220");
  file->thumbnail_url = valid_url;

  auto sanitized = DrivePickerSanitizer::Sanitize(file);

  ASSERT_TRUE(sanitized.has_value());
  ASSERT_TRUE(sanitized->thumbnail_url.has_value());
  EXPECT_EQ(sanitized->thumbnail_url.value(), valid_url);
}

TEST_F(DrivePickerSanitizerTest, BlocksInvalidThumbnailUrl) {
  {
    auto file = CreateValidFile();
    file->type = "photo";
    file->thumbnail_url = GURL("https://malicious.com/not-drive-storage/AJQWt");
    auto sanitized = DrivePickerSanitizer::Sanitize(file);
    ASSERT_TRUE(sanitized.has_value());
    EXPECT_FALSE(sanitized->thumbnail_url.has_value());
  }
  // Blocks http scheme even if host is valid.
  {
    auto file = CreateValidFile();
    file->type = "photo";
    file->thumbnail_url =
        GURL("http://lh3.googleusercontent.com/drive-storage/AJQWt");
    auto sanitized = DrivePickerSanitizer::Sanitize(file);
    ASSERT_TRUE(sanitized.has_value());
    EXPECT_FALSE(sanitized->thumbnail_url.has_value());
  }
  // Blocks chrome-extension scheme even if host is valid.
  {
    auto file = CreateValidFile();
    file->type = "photo";
    file->thumbnail_url = GURL(
        "chrome-extension://lh3.googleusercontent.com/drive-storage/AJQWt");
    auto sanitized = DrivePickerSanitizer::Sanitize(file);
    ASSERT_TRUE(sanitized.has_value());
    EXPECT_FALSE(sanitized->thumbnail_url.has_value());
  }
}

TEST_F(DrivePickerSanitizerTest, BlocksThumbnailForNonPhotoOrVideo) {
  auto file = CreateValidFile();
  file->type = "document";
  file->thumbnail_url = GURL(
      "https://lh3.googleusercontent.com/drive-storage/"
      "AJQWtBOI_xcNPWAVolLfcNDZGWus_ELL8GP-ZYLdkxOcPlUxsVOMr5jn-"
      "i261zxtYVcpA0kZeP"
      "sVWT2Ghrxdg03aoJYX7t9vlc4ojecNyW7QNrP6muY9-"
      "wvmO77SsiXHrwnU2GbRCt2V9NDv62u2"
      "R96rpvI=s220");

  auto sanitized = DrivePickerSanitizer::Sanitize(file);

  ASSERT_TRUE(sanitized.has_value());
  EXPECT_FALSE(sanitized->thumbnail_url.has_value());
}

TEST_F(DrivePickerSanitizerTest, SanitizeValidVideoWithThumbnail) {
  auto file = CreateValidFile();
  file->type = "video";
  GURL valid_url(
      "https://lh3.googleusercontent.com/drive-storage/"
      "AJQWtBOI_xcNPWAVolLfcNDZGWus_ELL8GP-ZYLdkxOcPlUxsVOMr5jn-"
      "i261zxtYVcpA0kZeP"
      "sVWT2Ghrxdg03aoJYX7t9vlc4ojecNyW7QNrP6muY9-"
      "wvmO77SsiXHrwnU2GbRCt2V9NDv62u2"
      "R96rpvI=s220");
  file->thumbnail_url = valid_url;

  auto sanitized = DrivePickerSanitizer::Sanitize(file);

  ASSERT_TRUE(sanitized.has_value());
  ASSERT_TRUE(sanitized->thumbnail_url.has_value());
  EXPECT_EQ(sanitized->thumbnail_url.value(), valid_url);
}

TEST_F(DrivePickerSanitizerTest, BlocksInvalidThumbnailUrlForVideo) {
  auto file = CreateValidFile();
  file->type = "video";
  file->thumbnail_url = GURL("https://malicious.com/not-drive-storage/AJQWt");
  auto sanitized = DrivePickerSanitizer::Sanitize(file);
  ASSERT_TRUE(sanitized.has_value());
  EXPECT_FALSE(sanitized->thumbnail_url.has_value());
}

TEST_F(DrivePickerSanitizerTest, AllowsCorrectDriveGoogleComUrls) {
  auto valid_cases = {"https://drive.google.com/thumbnail",
                      "https://drive.google.com/thumbnail?id=123"};

  for (const auto* url : valid_cases) {
    auto file = CreateValidFile();
    file->type = "photo";
    file->thumbnail_url = GURL(url);
    auto sanitized = DrivePickerSanitizer::Sanitize(file);
    ASSERT_TRUE(sanitized.has_value());
    EXPECT_TRUE(sanitized->thumbnail_url.has_value()) << "Failed for: " << url;
  }
}

TEST_F(DrivePickerSanitizerTest, BlocksInvalidDriveGoogleComUrls) {
  auto invalid_cases = {
      "https://drive.google.com/thumbnail_abc",
      "https://drive.google.com/file/d/ABC-123_xyz=/edit",
      "https://drive.google.com/file/d//view",
      "https://drive.google.com/file/d/abc^def/view",
      "https://drive.google.com/file/d/ABC-123_xyz=/view",
      "https://drive.google.com/file/d/ABC-123_xyz=/preview",
      "https://drive.google.com/file/d/ABC-123_xyz=/view?usp=sharing"};

  for (const auto* url : invalid_cases) {
    auto file = CreateValidFile();
    file->type = "photo";
    file->thumbnail_url = GURL(url);
    auto sanitized = DrivePickerSanitizer::Sanitize(file);
    ASSERT_TRUE(sanitized.has_value());
    EXPECT_FALSE(sanitized->thumbnail_url.has_value())
        << "Should have failed for: " << url;
  }
}

TEST_F(DrivePickerSanitizerTest, ThrowsOnInvalidId) {
  auto file = CreateValidFile();
  file->id = "invalid id with spaces";

  auto sanitized = DrivePickerSanitizer::Sanitize(file);

  EXPECT_FALSE(sanitized.has_value());
}

TEST_F(DrivePickerSanitizerTest, ThrowsOnUnsupportedType) {
  auto file = CreateValidFile();
  file->type = "unsupported";

  auto sanitized = DrivePickerSanitizer::Sanitize(file);

  EXPECT_FALSE(sanitized.has_value());
}

TEST_F(DrivePickerSanitizerTest, SanitizesInvalidResourceKey) {
  auto file = CreateValidFile();
  file->resource_key = "invalid key";

  auto sanitized = DrivePickerSanitizer::Sanitize(file);

  ASSERT_TRUE(sanitized.has_value());
  EXPECT_FALSE(sanitized->resource_key.has_value());
}

TEST_F(DrivePickerSanitizerTest, SanitizeDriveIconUrl) {
  // Valid hosts and paths.
  EXPECT_EQ(DrivePickerSanitizer::SanitizeDriveIconUrl(
                "https://drive-thirdparty.googleusercontent.com/16/type/pdf"),
            GURL("https://drive-thirdparty.googleusercontent.com/16/type/pdf"));
  EXPECT_EQ(DrivePickerSanitizer::SanitizeDriveIconUrl(
                "https://lh3.googleusercontent.com/16/hype/image/png"),
            GURL("https://lh3.googleusercontent.com/16/hype/image/png"));
  EXPECT_EQ(DrivePickerSanitizer::SanitizeDriveIconUrl(
                "https://lh6.googleusercontent.com/16/type/video/mp4"),
            GURL("https://lh6.googleusercontent.com/16/type/video/mp4"));

  // Invalid hosts
  EXPECT_TRUE(DrivePickerSanitizer::SanitizeDriveIconUrl(
                  "https://malicious.com/16/type/pdf")
                  .is_empty());
  EXPECT_TRUE(DrivePickerSanitizer::SanitizeDriveIconUrl(
                  "https://lh7.googleusercontent.com/16/type/pdf")
                  .is_empty());

  // Invalid URL string
  EXPECT_TRUE(
      DrivePickerSanitizer::SanitizeDriveIconUrl("not-a-url").is_empty());
}

TEST_F(DrivePickerSanitizerTest, SanitizeCopiesSanitizedIconUrl) {
  auto file = CreateValidFile();
  GURL valid_icon_url(
      "https://drive-thirdparty.googleusercontent.com/16/type/pdf");
  file->icon_url = valid_icon_url;

  auto sanitized = DrivePickerSanitizer::Sanitize(file);

  ASSERT_TRUE(sanitized.has_value());
  ASSERT_TRUE(sanitized->icon_url.has_value());
  EXPECT_EQ(sanitized->icon_url.value(), valid_icon_url);
}

TEST_F(DrivePickerSanitizerTest, ThrowsOnInvalidMetadata) {
  {
    auto file = CreateValidFile();
    file->name = "";
    auto sanitized = DrivePickerSanitizer::Sanitize(file);
    EXPECT_FALSE(sanitized.has_value());
  }
  {
    auto file = CreateValidFile();
    file->mime_type = "";
    auto sanitized = DrivePickerSanitizer::Sanitize(file);
    EXPECT_FALSE(sanitized.has_value());
  }
  {
    auto file = CreateValidFile();
    file->mime_type = "invalid-mime-no-slash";
    auto sanitized = DrivePickerSanitizer::Sanitize(file);
    EXPECT_FALSE(sanitized.has_value());
  }
  {
    auto file = CreateValidFile();
    file->mime_type = "text / html";  // Spaces are generally not allowed in the
                                      // type/subtype tokens
    auto sanitized = DrivePickerSanitizer::Sanitize(file);
    EXPECT_FALSE(sanitized.has_value());
  }
}