components/enterprise/obfuscation/core/utils_unittest.cc - chromium/src - Git at Google

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

 #include "components/enterprise/obfuscation/core/utils.h"

 #include "base/containers/span_reader.h"
 #include "base/files/file_enumerator.h"
 #include "base/files/file_path.h"
 #include "base/files/file_util.h"
 #include "base/files/scoped_temp_dir.h"
 #include "base/rand_util.h"
 #include "base/strings/string_number_conversions.h"
 #include "base/test/metrics/histogram_tester.h"
 #include "base/test/scoped_feature_list.h"
 #include "testing/gtest/include/gtest/gtest.h"

 namespace enterprise_obfuscation {
 namespace {

 // Helper function to divide data in chunks of random sizes.
 void ObfuscateTestDataInChunks(base::span<const uint8_t> test_data,
                                std::vector<uint8_t>& obfuscated_content) {
   std::array<uint8_t, kKeySize> derived_key;
   std::vector<uint8_t> nonce_prefix;
   auto header = CreateHeader(&derived_key, &nonce_prefix);
   ASSERT_TRUE(header.has_value());
   obfuscated_content.insert(obfuscated_content.end(), header.value().begin(),
                             header.value().end());

   base::SpanReader reader(test_data);
   uint32_t counter = 0;

   while (reader.remaining() > 0) {
     // Generate a random chunk size between 1 and remaining data size, capped at
     // `kMaxChunkSize`.
     size_t chunk_size =
         base::RandInt(1, std::min(reader.remaining(), kMaxChunkSize));

     // Read in the next chunk.
     auto current_chunk = reader.Read(chunk_size);
     ASSERT_TRUE(current_chunk.has_value());

     auto obfuscated_result =
         ObfuscateDataChunk(*current_chunk, derived_key, nonce_prefix, counter++,
                            reader.remaining() == 0);
     ASSERT_TRUE(obfuscated_result.has_value());

     obfuscated_content.insert(obfuscated_content.end(),
                               obfuscated_result.value().begin(),
                               obfuscated_result.value().end());
   }
 }

 // Helper function to count the number of files in a directory.
 int CountFilesInDirectory(const base::FilePath& dir_path) {
   base::FileEnumerator enum_files(dir_path, false, base::FileEnumerator::FILES);
   int count = 0;
   while (!enum_files.Next().empty()) {
     count++;
   }
   return count;
 }

 }  // namespace

 class ObfuscationUtilsTest
     : public ::testing::TestWithParam<std::tuple<bool, size_t>> {
  protected:
   ObfuscationUtilsTest() {
     feature_list_.InitWithFeatureState(kEnterpriseFileObfuscation,
                                        file_obfuscation_feature_enabled());
   }

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

   base::FilePath test_file_path() const {
     return temp_dir_.GetPath().AppendASCII("test_file.txt");
   }

   size_t test_data_size() const { return std::get<1>(GetParam()); }

   bool file_obfuscation_feature_enabled() const {
     return std::get<0>(GetParam());
   }

  private:
   base::test::ScopedFeatureList feature_list_;
   base::ScopedTempDir temp_dir_;
 };

 TEST_P(ObfuscationUtilsTest, ObfuscateAndDeobfuscateSingleDataChunk) {
   // Obfuscate the data chunk.
   std::vector<uint8_t> test_data = base::RandBytesAsVector(test_data_size());

   std::array<uint8_t, kKeySize> derived_key;
   std::vector<uint8_t> nonce_prefix;
   auto header = CreateHeader(&derived_key, &nonce_prefix);
   constexpr uint32_t kInitialChunkCounter = 0;

   auto obfuscated_chunk =
       ObfuscateDataChunk(base::as_byte_span(test_data), derived_key,
                          nonce_prefix, kInitialChunkCounter, true);

   if (!file_obfuscation_feature_enabled()) {
     ASSERT_EQ(obfuscated_chunk.error(), Error::kDisabled);
     ASSERT_EQ(header.error(), Error::kDisabled);
     return;
   }
   ASSERT_TRUE(header.has_value());
   ASSERT_TRUE(obfuscated_chunk.has_value());
   ASSERT_NE(obfuscated_chunk.value(), test_data);

   // Deobfuscate the data chunk.
   auto header_data = GetHeaderData(header.value());
   ASSERT_TRUE(header_data.has_value());

   auto chunk_size = GetObfuscatedChunkSize(obfuscated_chunk.value());

   if (test_data.size() > kMaxChunkSize) {
     ASSERT_EQ(chunk_size.error(), Error::kDeobfuscationFailed);
     return;
   }

   ASSERT_TRUE(chunk_size.has_value());
   EXPECT_EQ(chunk_size.value(), test_data.size() + kAuthTagSize);

   auto deobfuscated_chunk = DeobfuscateDataChunk(
       base::span(obfuscated_chunk.value())
           .subspan(kChunkSizePrefixSize, chunk_size.value()),
       header_data.value().derived_key, header_data.value().nonce_prefix,
       kInitialChunkCounter, true);
   ASSERT_TRUE(deobfuscated_chunk.has_value());
   EXPECT_EQ(deobfuscated_chunk.value(), test_data);

   // Deobfuscation should fail when we modify the ciphertext.
   obfuscated_chunk.value()[kChunkSizePrefixSize] ^= 1;
   deobfuscated_chunk = DeobfuscateDataChunk(
       base::span(obfuscated_chunk.value())
           .subspan(kChunkSizePrefixSize, chunk_size.value()),
       header_data.value().derived_key, header_data.value().nonce_prefix,
       kInitialChunkCounter, true);
   ASSERT_EQ(deobfuscated_chunk.error(), Error::kDeobfuscationFailed);
 }

 TEST_P(ObfuscationUtilsTest, DeobfuscateFileInPlace) {
   base::HistogramTester histogram_tester;

   std::vector<uint8_t> test_data = base::RandBytesAsVector(test_data_size());
   ASSERT_TRUE(base::WriteFile(test_file_path(), test_data));

   int64_t original_size = base::GetFileSize(test_file_path()).value_or(0);

   auto result = DeobfuscateFileInPlace(test_file_path());

   if (!file_obfuscation_feature_enabled()) {
     ASSERT_EQ(result.error(), Error::kDisabled);
     histogram_tester.ExpectUniqueSample(kObfuscationResultHistogram,
                                         Error::kDisabled, 1);
     return;
   }

   // Deobfuscating an unobfuscated file should fail.
   Error unobfuscated_error = original_size == 0 ? Error::kFileOperationError
                                                 : Error::kDeobfuscationFailed;
   ASSERT_EQ(result.error(), unobfuscated_error);
   histogram_tester.ExpectUniqueSample(kObfuscationResultHistogram,
                                       unobfuscated_error, 1);

   // Only the original test file should remain.
   EXPECT_EQ(CountFilesInDirectory(test_file_path().DirName()), 1);

   std::vector<uint8_t> obfuscated_content;

   ObfuscateTestDataInChunks(test_data, obfuscated_content);

   ASSERT_TRUE(base::WriteFile(test_file_path(), obfuscated_content));
   ASSERT_TRUE(DeobfuscateFileInPlace(test_file_path()).has_value());

   histogram_tester.ExpectBucketCount(kObfuscationResultHistogram,
                                      Error::kSuccess, 1);

   auto deobfuscated_content = base::ReadFileToBytes(test_file_path());
   ASSERT_TRUE(deobfuscated_content.has_value());
   EXPECT_EQ(deobfuscated_content.value(), test_data);

   // Get deobfuscated file size which should match original.
   std::optional<int64_t> deobfuscated_size =
       base::GetFileSize(test_file_path());
   ASSERT_TRUE(deobfuscated_size.has_value());
   EXPECT_EQ(deobfuscated_size.value(), original_size);

   // Deobfuscating to an invalid path should fail.
   base::FilePath invalid_path(
       test_file_path().InsertBeforeExtensionASCII("_invalid"));
   ASSERT_EQ(DeobfuscateFileInPlace(invalid_path).error(),
             Error::kFileOperationError);

   // Only the original test file should remain.
   EXPECT_EQ(CountFilesInDirectory(test_file_path().DirName()), 1);

   int expected_file_op_count = (original_size == 0) ? 2 : 1;
   histogram_tester.ExpectBucketCount(kObfuscationResultHistogram,
                                      Error::kFileOperationError,
                                      expected_file_op_count);
   histogram_tester.ExpectTotalCount(kObfuscationResultHistogram, 3);
 }

 TEST_P(ObfuscationUtilsTest, ObfuscateAndDeobfuscateVariableChunks) {
   if (!file_obfuscation_feature_enabled()) {
     GTEST_SKIP() << "File obfuscation feature is disabled.";
   }

   // Create test data.
   std::vector<uint8_t> test_data = base::RandBytesAsVector(test_data_size());

   // Obfuscate data in chunks of random sizes.
   std::vector<uint8_t> obfuscated_content;
   ObfuscateTestDataInChunks(test_data, obfuscated_content);

   // Deobfuscate chunk by chunk.
   std::vector<uint8_t> deobfuscated_content;
   size_t offset = kHeaderSize;
   uint32_t counter = 0;

   std::vector<uint8_t> header(obfuscated_content.begin(),
                               obfuscated_content.begin() + kHeaderSize);

   auto header_data = GetHeaderData(header);
   ASSERT_TRUE(header_data.has_value());

   while (offset < obfuscated_content.size()) {
     // Read chunk size
     auto chunk_size = GetObfuscatedChunkSize(
         base::span(obfuscated_content).subspan(offset, kChunkSizePrefixSize));
     ASSERT_TRUE(chunk_size.has_value());
     offset += kChunkSizePrefixSize;

     // Deobfuscate chunk
     auto deobfuscated_chunk = DeobfuscateDataChunk(
         base::span(obfuscated_content).subspan(offset, chunk_size.value()),
         header_data.value().derived_key, header_data.value().nonce_prefix,
         counter++, (offset + chunk_size.value() >= obfuscated_content.size()));
     ASSERT_TRUE(deobfuscated_chunk.has_value());

     std::move(deobfuscated_chunk.value().begin(),
               deobfuscated_chunk.value().end(),
               std::back_inserter(deobfuscated_content));

     offset += chunk_size.value();
   }

   // Compare deobfuscated content with original test data
   EXPECT_EQ(deobfuscated_content, test_data);
 }

 INSTANTIATE_TEST_SUITE_P(
     ObfuscationUtilsFeatureTest,
     ObfuscationUtilsTest,
     ::testing::Combine(
         ::testing::Bool(),  // File obfuscator feature enabled/disabled
         ::testing::Values(0,
                           10,
                           kMaxChunkSize + 1024,
                           kMaxChunkSize * 2 + 1024)));

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

	#include "components/enterprise/obfuscation/core/utils.h"

	#include "base/containers/span_reader.h"
	#include "base/files/file_enumerator.h"
	#include "base/files/file_path.h"
	#include "base/files/file_util.h"
	#include "base/files/scoped_temp_dir.h"
	#include "base/rand_util.h"
	#include "base/strings/string_number_conversions.h"
	#include "base/test/metrics/histogram_tester.h"
	#include "base/test/scoped_feature_list.h"
	#include "testing/gtest/include/gtest/gtest.h"

	namespace enterprise_obfuscation {
	namespace {

	// Helper function to divide data in chunks of random sizes.
	void ObfuscateTestDataInChunks(base::span<const uint8_t> test_data,
	std::vector<uint8_t>& obfuscated_content) {
	std::array<uint8_t, kKeySize> derived_key;
	std::vector<uint8_t> nonce_prefix;
	auto header = CreateHeader(&derived_key, &nonce_prefix);
	ASSERT_TRUE(header.has_value());
	obfuscated_content.insert(obfuscated_content.end(), header.value().begin(),
	header.value().end());

	base::SpanReader reader(test_data);
	uint32_t counter = 0;

	while (reader.remaining() > 0) {
	// Generate a random chunk size between 1 and remaining data size, capped at
	// `kMaxChunkSize`.
	size_t chunk_size =
	base::RandInt(1, std::min(reader.remaining(), kMaxChunkSize));

	// Read in the next chunk.
	auto current_chunk = reader.Read(chunk_size);
	ASSERT_TRUE(current_chunk.has_value());

	auto obfuscated_result =
	ObfuscateDataChunk(*current_chunk, derived_key, nonce_prefix, counter++,
	reader.remaining() == 0);
	ASSERT_TRUE(obfuscated_result.has_value());

	obfuscated_content.insert(obfuscated_content.end(),
	obfuscated_result.value().begin(),
	obfuscated_result.value().end());
	}
	}

	// Helper function to count the number of files in a directory.
	int CountFilesInDirectory(const base::FilePath& dir_path) {
	base::FileEnumerator enum_files(dir_path, false, base::FileEnumerator::FILES);
	int count = 0;
	while (!enum_files.Next().empty()) {
	count++;
	}
	return count;
	}

	} // namespace

	class ObfuscationUtilsTest
	: public ::testing::TestWithParam<std::tuple<bool, size_t>> {
	protected:
	ObfuscationUtilsTest() {
	feature_list_.InitWithFeatureState(kEnterpriseFileObfuscation,
	file_obfuscation_feature_enabled());
	}

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

	base::FilePath test_file_path() const {
	return temp_dir_.GetPath().AppendASCII("test_file.txt");
	}

	size_t test_data_size() const { return std::get<1>(GetParam()); }

	bool file_obfuscation_feature_enabled() const {
	return std::get<0>(GetParam());
	}

	private:
	base::test::ScopedFeatureList feature_list_;
	base::ScopedTempDir temp_dir_;
	};

	TEST_P(ObfuscationUtilsTest, ObfuscateAndDeobfuscateSingleDataChunk) {
	// Obfuscate the data chunk.
	std::vector<uint8_t> test_data = base::RandBytesAsVector(test_data_size());

	std::array<uint8_t, kKeySize> derived_key;
	std::vector<uint8_t> nonce_prefix;
	auto header = CreateHeader(&derived_key, &nonce_prefix);
	constexpr uint32_t kInitialChunkCounter = 0;

	auto obfuscated_chunk =
	ObfuscateDataChunk(base::as_byte_span(test_data), derived_key,
	nonce_prefix, kInitialChunkCounter, true);

	if (!file_obfuscation_feature_enabled()) {
	ASSERT_EQ(obfuscated_chunk.error(), Error::kDisabled);
	ASSERT_EQ(header.error(), Error::kDisabled);
	return;
	}
	ASSERT_TRUE(header.has_value());
	ASSERT_TRUE(obfuscated_chunk.has_value());
	ASSERT_NE(obfuscated_chunk.value(), test_data);

	// Deobfuscate the data chunk.
	auto header_data = GetHeaderData(header.value());
	ASSERT_TRUE(header_data.has_value());

	auto chunk_size = GetObfuscatedChunkSize(obfuscated_chunk.value());

	if (test_data.size() > kMaxChunkSize) {
	ASSERT_EQ(chunk_size.error(), Error::kDeobfuscationFailed);
	return;
	}

	ASSERT_TRUE(chunk_size.has_value());
	EXPECT_EQ(chunk_size.value(), test_data.size() + kAuthTagSize);

	auto deobfuscated_chunk = DeobfuscateDataChunk(
	base::span(obfuscated_chunk.value())
	.subspan(kChunkSizePrefixSize, chunk_size.value()),
	header_data.value().derived_key, header_data.value().nonce_prefix,
	kInitialChunkCounter, true);
	ASSERT_TRUE(deobfuscated_chunk.has_value());
	EXPECT_EQ(deobfuscated_chunk.value(), test_data);

	// Deobfuscation should fail when we modify the ciphertext.
	obfuscated_chunk.value()[kChunkSizePrefixSize] ^= 1;
	deobfuscated_chunk = DeobfuscateDataChunk(
	base::span(obfuscated_chunk.value())
	.subspan(kChunkSizePrefixSize, chunk_size.value()),
	header_data.value().derived_key, header_data.value().nonce_prefix,
	kInitialChunkCounter, true);
	ASSERT_EQ(deobfuscated_chunk.error(), Error::kDeobfuscationFailed);
	}

	TEST_P(ObfuscationUtilsTest, DeobfuscateFileInPlace) {
	base::HistogramTester histogram_tester;

	std::vector<uint8_t> test_data = base::RandBytesAsVector(test_data_size());
	ASSERT_TRUE(base::WriteFile(test_file_path(), test_data));

	int64_t original_size = base::GetFileSize(test_file_path()).value_or(0);

	auto result = DeobfuscateFileInPlace(test_file_path());

	if (!file_obfuscation_feature_enabled()) {
	ASSERT_EQ(result.error(), Error::kDisabled);
	histogram_tester.ExpectUniqueSample(kObfuscationResultHistogram,
	Error::kDisabled, 1);
	return;
	}

	// Deobfuscating an unobfuscated file should fail.
	Error unobfuscated_error = original_size == 0 ? Error::kFileOperationError
	: Error::kDeobfuscationFailed;
	ASSERT_EQ(result.error(), unobfuscated_error);
	histogram_tester.ExpectUniqueSample(kObfuscationResultHistogram,
	unobfuscated_error, 1);

	// Only the original test file should remain.
	EXPECT_EQ(CountFilesInDirectory(test_file_path().DirName()), 1);

	std::vector<uint8_t> obfuscated_content;

	ObfuscateTestDataInChunks(test_data, obfuscated_content);

	ASSERT_TRUE(base::WriteFile(test_file_path(), obfuscated_content));
	ASSERT_TRUE(DeobfuscateFileInPlace(test_file_path()).has_value());

	histogram_tester.ExpectBucketCount(kObfuscationResultHistogram,
	Error::kSuccess, 1);

	auto deobfuscated_content = base::ReadFileToBytes(test_file_path());
	ASSERT_TRUE(deobfuscated_content.has_value());
	EXPECT_EQ(deobfuscated_content.value(), test_data);

	// Get deobfuscated file size which should match original.
	std::optional<int64_t> deobfuscated_size =
	base::GetFileSize(test_file_path());
	ASSERT_TRUE(deobfuscated_size.has_value());
	EXPECT_EQ(deobfuscated_size.value(), original_size);

	// Deobfuscating to an invalid path should fail.
	base::FilePath invalid_path(
	test_file_path().InsertBeforeExtensionASCII("_invalid"));
	ASSERT_EQ(DeobfuscateFileInPlace(invalid_path).error(),
	Error::kFileOperationError);

	// Only the original test file should remain.
	EXPECT_EQ(CountFilesInDirectory(test_file_path().DirName()), 1);

	int expected_file_op_count = (original_size == 0) ? 2 : 1;
	histogram_tester.ExpectBucketCount(kObfuscationResultHistogram,
	Error::kFileOperationError,
	expected_file_op_count);
	histogram_tester.ExpectTotalCount(kObfuscationResultHistogram, 3);
	}

	TEST_P(ObfuscationUtilsTest, ObfuscateAndDeobfuscateVariableChunks) {
	if (!file_obfuscation_feature_enabled()) {
	GTEST_SKIP() << "File obfuscation feature is disabled.";
	}

	// Create test data.
	std::vector<uint8_t> test_data = base::RandBytesAsVector(test_data_size());

	// Obfuscate data in chunks of random sizes.
	std::vector<uint8_t> obfuscated_content;
	ObfuscateTestDataInChunks(test_data, obfuscated_content);

	// Deobfuscate chunk by chunk.
	std::vector<uint8_t> deobfuscated_content;
	size_t offset = kHeaderSize;
	uint32_t counter = 0;

	std::vector<uint8_t> header(obfuscated_content.begin(),
	obfuscated_content.begin() + kHeaderSize);

	auto header_data = GetHeaderData(header);
	ASSERT_TRUE(header_data.has_value());

	while (offset < obfuscated_content.size()) {
	// Read chunk size
	auto chunk_size = GetObfuscatedChunkSize(
	base::span(obfuscated_content).subspan(offset, kChunkSizePrefixSize));
	ASSERT_TRUE(chunk_size.has_value());
	offset += kChunkSizePrefixSize;

	// Deobfuscate chunk
	auto deobfuscated_chunk = DeobfuscateDataChunk(
	base::span(obfuscated_content).subspan(offset, chunk_size.value()),
	header_data.value().derived_key, header_data.value().nonce_prefix,
	counter++, (offset + chunk_size.value() >= obfuscated_content.size()));
	ASSERT_TRUE(deobfuscated_chunk.has_value());

	std::move(deobfuscated_chunk.value().begin(),
	deobfuscated_chunk.value().end(),
	std::back_inserter(deobfuscated_content));

	offset += chunk_size.value();
	}

	// Compare deobfuscated content with original test data
	EXPECT_EQ(deobfuscated_content, test_data);
	}

	INSTANTIATE_TEST_SUITE_P(
	ObfuscationUtilsFeatureTest,
	ObfuscationUtilsTest,
	::testing::Combine(
	::testing::Bool(), // File obfuscator feature enabled/disabled
	::testing::Values(0,
	10,
	kMaxChunkSize + 1024,
	kMaxChunkSize * 2 + 1024)));

	} // namespace enterprise_obfuscation