// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#include "utils/cbor_padding.h"

#include <cstddef>
#include <cstdint>
#include <string>

#include "absl/log/absl_check.h"
#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "cose/sig_structure.h"
#include "crypto/algorithms.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "proto/boxes_hash_assertion.pb.h"
#include "proto/data_hash_assertion.cbor.h"
#include "proto/data_hash_assertion.pb.h"
#include "testing/proto_test_utils.h"

namespace credentio {
namespace {

using ::absl_testing::StatusIs;
using ::credentio_testing::EqualsProto;
using ::credentio_testing::ParseTextProtoOrDie;
using ::testing::HasSubstr;
using ::testing::Le;
using ::testing::SizeIs;
using ::testing::TestWithParam;
using ::testing::Values;
using ::testing::ValuesIn;

constexpr size_t kTstContainerHeaderByteSize = 26;

struct Range {
  int64_t start;
  int64_t length;
};

size_t EncodeByteSize(const DataHashAssertion& data_hash) {
  auto res = cbor::FromProto(data_hash, {});
  ABSL_CHECK_OK(res);
  return (*res)->toString().size();
}

TEST(CborPaddingTest, PadDataHashAssertion) {
  DataHashAssertion data_hash;
  data_hash.set_alg("sha256");
  data_hash.set_hash(std::string(HashSizeBytes(HashAlgorithm::kSha256), '\0'));
  auto exclusion = data_hash.add_exclusions();
  exclusion->set_start(0);
  exclusion->set_length(0);
  data_hash.set_pad("to be replaced");
  data_hash.set_pad2("to be cleared");
  PadDataHashAssertion(data_hash);
  DataHashAssertion expected = ParseTextProtoOrDie<DataHashAssertion>(R"pb(
    exclusions { start: 0 length: 0 }
    alg: "sha256"
    hash: "\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000"
    pad: "\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000"
  )pb");
  EXPECT_THAT(data_hash, EqualsProto(expected));
}

TEST(CborPaddingTest, PadDataHashAssertionEmptyPadding) {
  DataHashAssertion data_hash;
  data_hash.set_alg("sha256");
  data_hash.set_hash(std::string(HashSizeBytes(HashAlgorithm::kSha256), '\0'));
  data_hash.set_pad("to be cleared");
  data_hash.set_pad2("to be cleared");
  PadDataHashAssertion(data_hash);
  DataHashAssertion expected = ParseTextProtoOrDie<DataHashAssertion>(R"pb(
    alg: "sha256"
    hash: "\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000"
    pad: ""
  )pb");
  EXPECT_THAT(data_hash, EqualsProto(expected));
}

TEST(CborPaddingTest, PadDataHashAssertionWithPad2) {
  DataHashAssertion data_hash;
  data_hash.set_alg("sha256");
  data_hash.set_hash(std::string(HashSizeBytes(HashAlgorithm::kSha256), '\0'));
  for (int i = 0; i < 2; i++) {
    auto exclusion = data_hash.add_exclusions();
    exclusion->set_start(0);
    exclusion->set_length(0);
  }
  PadDataHashAssertion(data_hash);
  DataHashAssertion expected = ParseTextProtoOrDie<DataHashAssertion>(R"pb(
    exclusions { start: 0 length: 0 }
    exclusions { start: 0 length: 0 }
    alg: "sha256"
    hash: "\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000"
    pad: "\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000"
    pad2: ""
  )pb");
  EXPECT_THAT(data_hash, EqualsProto(expected));
}

TEST(CborPaddingTest, PadDataHashAssertionWithPad2Filled) {
  DataHashAssertion data_hash;
  data_hash.set_alg("sha256");
  data_hash.set_hash(std::string(HashSizeBytes(HashAlgorithm::kSha256), '\0'));
  for (int i = 0; i < 2; i++) {
    auto exclusion = data_hash.add_exclusions();
    exclusion->set_start(0);
    exclusion->set_length(4294967296);
  }
  PadDataHashAssertion(data_hash);
  DataHashAssertion expected = ParseTextProtoOrDie<DataHashAssertion>(R"pb(
    exclusions { start: 0 length: 4294967296 }
    exclusions { start: 0 length: 4294967296 }
    alg: "sha256"
    hash: "\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000"
    pad: "\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000"
    pad2: "\000"
  )pb");
  EXPECT_THAT(data_hash, EqualsProto(expected));
}

class PadDataHashTest : public TestWithParam<Range> {};

TEST_P(PadDataHashTest, PadDataHashAssertionTestWithParam) {
  credentio::DataHashAssertion data_hash;
  data_hash.set_alg("sha256");
  data_hash.set_hash(std::string(HashSizeBytes(HashAlgorithm::kSha256), '\0'));
  auto exclusion = data_hash.add_exclusions();
  exclusion->set_start(GetParam().start);
  exclusion->set_length(GetParam().length);
  PadDataHashAssertion(data_hash);
  auto size_1 = EncodeByteSize(data_hash);
  EXPECT_THAT(data_hash.pad().size() + data_hash.pad2().size(), Le(16));

  data_hash.mutable_exclusions(0)->set_start(0);
  data_hash.mutable_exclusions(0)->set_length(0);
  PadDataHashAssertion(data_hash);
  auto size_2 = EncodeByteSize(data_hash);

  EXPECT_EQ(size_1, size_2);
}

constexpr Range kRanges[] = {
    {.start = 0, .length = 0},
    {.start = 0, .length = 23},
    {.start = 0, .length = 24},
    {.start = 23, .length = 24},
    {.start = 24, .length = 255},
    {.start = 255, .length = 255},
    {.start = 255, .length = 256},
    {.start = 256, .length = 256},
    {.start = 256, .length = 65535},
    {.start = 256, .length = 65536},
    {.start = 65535, .length = 65536},
    {.start = 65536, .length = 65536},
    {.start = 65536, .length = 4294967295},
    {.start = 65536, .length = 4294967296},
    {.start = 4294967295, .length = 4294967296},
    {.start = 4294967296, .length = 4294967296},
    {.start = 9223372036854775807, .length = 9223372036854775807},
};

INSTANTIATE_TEST_SUITE_P(PadDataHashAssertionTestWithParam, PadDataHashTest,
                         ValuesIn(kRanges));

TEST(CborPaddingTest, PadBoxesHashAssertion) {
  BoxesHashAssertion boxes_hash = ParseTextProtoOrDie<BoxesHashAssertion>(R"pb(
    boxes { names: "name1" alg: "sha256" hash: "\000" }
    boxes { names: "name2" alg: "sha256" hash: "\000" }
    alg: "sha256"
  )pb");
  PadBoxesHashAssertion(boxes_hash);
  BoxesHashAssertion expected = ParseTextProtoOrDie<BoxesHashAssertion>(R"pb(
    boxes { names: "name1" alg: "sha256" hash: "\000" pad: "" pad2: "" }
    boxes { names: "name2" alg: "sha256" hash: "\000" pad: "" pad2: "" }
    alg: "sha256"
  )pb");
  EXPECT_THAT(boxes_hash, EqualsProto(expected));
}

TEST(CborPaddingTest, PadCoseSign1TaggedStructureNoTimestamp) {
  int64_t max_signature_size = 10;
  int64_t max_timestamp_size = 10;
  CoseSign1TaggedStructure cose_sign1;
  cose_sign1.signature = std::string(max_signature_size, '\0');

  CoseSign1TaggedStructure expected = cose_sign1;
  expected.unprotected_header.pad =
      std::string(max_timestamp_size + kTstContainerHeaderByteSize, '\0');
  expected.unprotected_header.pad2 = "";

  ABSL_ASSERT_OK(PadCoseSign1TaggedStructure(cose_sign1, max_signature_size,
                                             max_timestamp_size));
  EXPECT_EQ(cose_sign1, expected);
}

TEST(CborPaddingTest, PadCoseSign1TaggedStructureEmptySignature) {
  int64_t max_signature_size = 10;
  int64_t max_timestamp_size = 10;
  CoseSign1TaggedStructure cose_sign1;
  cose_sign1.signature = "";
  cose_sign1.unprotected_header.sig_tst2 = TstContainer{
      .tst_tokens = {{.val = std::string(max_timestamp_size, '\0')}}};

  CoseSign1TaggedStructure expected = cose_sign1;
  expected.unprotected_header.pad = std::string(max_signature_size, '\0');
  expected.unprotected_header.pad2 = std::string(1, '\0');

  ABSL_ASSERT_OK(PadCoseSign1TaggedStructure(cose_sign1, max_signature_size,
                                             max_timestamp_size));
  EXPECT_EQ(cose_sign1, expected);
}

TEST(CborPaddingTest, PadCoseSign1TaggedStructureEmptyTimestamp) {
  int64_t max_signature_size = 10;
  int64_t max_timestamp_size = 10;
  CoseSign1TaggedStructure cose_sign1;
  cose_sign1.signature = std::string(max_signature_size, '\0');
  cose_sign1.unprotected_header.sig_tst2 =
      TstContainer{.tst_tokens = {{.val = ""}}};

  CoseSign1TaggedStructure expected = cose_sign1;
  expected.unprotected_header.pad = std::string(max_timestamp_size, '\0');
  expected.unprotected_header.pad2 = std::string(1, '\0');

  ABSL_ASSERT_OK(PadCoseSign1TaggedStructure(cose_sign1, max_signature_size,
                                             max_timestamp_size));
  EXPECT_EQ(cose_sign1, expected);
}

TEST(CborPaddingTest, PadCoseSign1TaggedStructureEmptySignatureAndTimestamp) {
  int64_t max_signature_size = 10;
  int64_t max_timestamp_size = 10;
  CoseSign1TaggedStructure cose_sign1;
  cose_sign1.signature = "";
  cose_sign1.unprotected_header.sig_tst2 =
      TstContainer{.tst_tokens = {{.val = ""}}};

  CoseSign1TaggedStructure expected = cose_sign1;
  expected.unprotected_header.pad =
      std::string(max_signature_size + max_timestamp_size, '\0');
  expected.unprotected_header.pad2 = std::string(1, '\0');

  ABSL_ASSERT_OK(PadCoseSign1TaggedStructure(cose_sign1, max_signature_size,
                                             max_timestamp_size));
  EXPECT_EQ(cose_sign1, expected);
}

TEST(CborPaddingTest, PadCoseSign1TaggedStructureFullSignatureAndTimestamp) {
  int64_t max_signature_size = 10;
  int64_t max_timestamp_size = 10;
  CoseSign1TaggedStructure cose_sign1;
  cose_sign1.signature = std::string(max_signature_size, '\0');
  cose_sign1.unprotected_header.sig_tst2 = TstContainer{
      .tst_tokens = {{.val = std::string(max_timestamp_size, '\0')}}};

  CoseSign1TaggedStructure expected = cose_sign1;
  expected.unprotected_header.pad = "";
  expected.unprotected_header.pad2 = std::string(1, '\0');

  ABSL_ASSERT_OK(PadCoseSign1TaggedStructure(cose_sign1, max_signature_size,
                                             max_timestamp_size));
  EXPECT_EQ(cose_sign1, expected);
}

TEST(CborPaddingTest,
     PadCoseSign1TaggedStructureWithSignatureAndTimestampDeltas) {
  int64_t max_signature_size = 20;
  int64_t max_timestamp_size = 200;

  CoseSign1TaggedStructure cose_sign1;
  cose_sign1.signature = std::string(10, '\0');
  cose_sign1.unprotected_header.sig_tst2 =
      TstContainer{.tst_tokens = {{.val = std::string(10, '\0')}}};

  CoseSign1TaggedStructure expected = cose_sign1;
  expected.unprotected_header.pad =
      std::string(max_signature_size - 10 + max_timestamp_size - 10, '\0');
  expected.unprotected_header.pad2 = std::string(1, '\0');

  ABSL_ASSERT_OK(PadCoseSign1TaggedStructure(cose_sign1, max_signature_size,
                                             max_timestamp_size));
  EXPECT_EQ(cose_sign1, expected);
}

TEST(CborPaddingTest, PadCoseSign1TaggedStructureSignatureTooBig) {
  int64_t max_signature_size = 10;
  int64_t max_timestamp_size = 10;
  CoseSign1TaggedStructure cose_sign1;
  cose_sign1.signature = std::string(max_signature_size + 1, '\0');
  cose_sign1.unprotected_header.sig_tst2 = TstContainer{
      .tst_tokens = {{.val = std::string(max_timestamp_size, '\0')}}};
  EXPECT_THAT(PadCoseSign1TaggedStructure(cose_sign1, max_signature_size,
                                          max_timestamp_size),
              StatusIs(absl::StatusCode::kInternal, HasSubstr("signature")));
}

TEST(CborPaddingTest, PadCoseSign1TaggedStructureTimestampTooBig) {
  int64_t max_signature_size = 10;
  int64_t max_timestamp_size = 10;
  CoseSign1TaggedStructure cose_sign1;
  cose_sign1.signature = std::string(max_signature_size, '\0');
  cose_sign1.unprotected_header.sig_tst2 = TstContainer{
      .tst_tokens = {{.val = std::string(max_timestamp_size + 1, '\0')}}};
  EXPECT_THAT(
      PadCoseSign1TaggedStructure(cose_sign1, max_signature_size,
                                  max_timestamp_size),
      StatusIs(absl::StatusCode::kInternal, HasSubstr("timestamp size")));
}

TEST(CborPaddingTest, PadCoseSign1TaggedStructureMultipleTimestampTokens) {
  int64_t max_signature_size = 10;
  int64_t max_timestamp_size = 10;
  CoseSign1TaggedStructure cose_sign1;
  cose_sign1.signature = std::string(max_signature_size, '\0');
  cose_sign1.unprotected_header.sig_tst2 = TstContainer{
      .tst_tokens = {{.val = std::string(max_timestamp_size, '\0')},
                     {.val = std::string(max_timestamp_size, '\0')}}};
  EXPECT_THAT(
      PadCoseSign1TaggedStructure(cose_sign1, max_signature_size,
                                  max_timestamp_size),
      StatusIs(absl::StatusCode::kInternal, HasSubstr("multiple timestamp")));
}

struct CoseSign1TaggedStructureInput {
  int64_t signature_size;
  int64_t timestamp_size;
  int64_t signature_max_size;
  int64_t timestamp_max_size;
};

class PadCoseSign1TaggedStructureTest
    : public TestWithParam<CoseSign1TaggedStructureInput> {};

TEST_P(PadCoseSign1TaggedStructureTest, PadCoseSign1TaggedStructureWithParams) {
  CoseSign1TaggedStructure cose_sign1;
  cose_sign1.signature = std::string(GetParam().signature_size, '\0');
  if (GetParam().timestamp_size > 0) {
    cose_sign1.unprotected_header.sig_tst2 = TstContainer{
        .tst_tokens = {{.val = std::string(GetParam().timestamp_size, '\0')}}};
  }
  int64_t max_signature_size = GetParam().signature_max_size;
  int64_t max_timestamp_size = GetParam().timestamp_max_size;
  ABSL_ASSERT_OK(PadCoseSign1TaggedStructure(cose_sign1, max_signature_size,
                                             max_timestamp_size));

  CoseSign1TaggedStructure cose_sign2;
  cose_sign2.signature = std::string(max_signature_size, '\0');
  ABSL_ASSERT_OK(PadCoseSign1TaggedStructure(cose_sign2, max_signature_size,
                                             max_timestamp_size));

  EXPECT_EQ(EncodeCoseSign1TaggedStructure(cose_sign1).size(),
            EncodeCoseSign1TaggedStructure(cose_sign2).size());
}
constexpr CoseSign1TaggedStructureInput kCoseSign1TaggedInputs[] = {
    {
        .signature_size = 10,
        .timestamp_size = 0,
        .signature_max_size = 10,
        .timestamp_max_size = 10,
    },
    {
        .signature_size = 10,
        .timestamp_size = 0,
        .signature_max_size = 100,
        .timestamp_max_size = 100,
    },
    {
        .signature_size = 10,
        .timestamp_size = 0,
        .signature_max_size = 1000,
        .timestamp_max_size = 1000,
    },
    {
        .signature_size = 10,
        .timestamp_size = 10,
        .signature_max_size = 10,
        .timestamp_max_size = 10,
    },
    {
        .signature_size = 10,
        .timestamp_size = 10,
        .signature_max_size = 100,
        .timestamp_max_size = 100,
    },
    {
        .signature_size = 10,
        .timestamp_size = 10,
        .signature_max_size = 1000,
        .timestamp_max_size = 1000,
    },
    {
        .signature_size = 10,
        .timestamp_size = 10,
        .signature_max_size = 10000,
        .timestamp_max_size = 10000,
    },
    {
        .signature_size = 100,
        .timestamp_size = 100,
        .signature_max_size = 100,
        .timestamp_max_size = 100,
    },
    {
        .signature_size = 1000,
        .timestamp_size = 1000,
        .signature_max_size = 1000,
        .timestamp_max_size = 1000,
    },
    {
        .signature_size = 10000,
        .timestamp_size = 10000,
        .signature_max_size = 10000,
        .timestamp_max_size = 10000,
    },
};
INSTANTIATE_TEST_SUITE_P(
    PadCoseSign1TaggedStructureWithParams, PadCoseSign1TaggedStructureTest,
    ValuesIn(kCoseSign1TaggedInputs),
    [](const testing::TestParamInfo<CoseSign1TaggedStructureInput>& info) {
      return absl::StrCat("signature_", info.param.signature_size,
                          "_timestamp_", info.param.timestamp_size,
                          "_signature_max_", info.param.signature_max_size,
                          "_timestamp_max_", info.param.timestamp_max_size);
    });

TEST(CborPaddingTest,
     PadCoseSign1TaggedStructureSmallExpectedSizeReturnsError) {
  CoseSign1TaggedStructure cose_sign1;
  cose_sign1.unprotected_header.pad = "";
  cose_sign1.unprotected_header.pad2 = std::string(8, '\0');
  int initial_size = EncodeCoseSign1TaggedStructure(cose_sign1).size();
  EXPECT_THAT(PadCoseSign1TaggedStructure(cose_sign1, initial_size - 1),
              StatusIs(absl::StatusCode::kInternal,
                       HasSubstr(absl::StrCat(
                           "Expected size ", initial_size - 1,
                           " is smaller than base size ", initial_size))));
}

// Returns the size of structure that PadCoseSign1TaggedStructure expects as
// initial size to pad from.
int GetInitialSizeForPadCoseSign1TaggedStructure() {
  CoseSign1TaggedStructure cose_sign1;
  cose_sign1.unprotected_header.pad = "";
  cose_sign1.unprotected_header.pad2 = std::string(8, '\0');
  return EncodeCoseSign1TaggedStructure(cose_sign1).size();
}

class PadOverallCoseSign1TaggedStructureTest : public TestWithParam<size_t> {};

TEST_P(PadOverallCoseSign1TaggedStructureTest, PadCoseSign1TaggedStructure) {
  CoseSign1TaggedStructure cose_sign1;
  size_t target_size = GetParam();

  ABSL_ASSERT_OK(PadCoseSign1TaggedStructure(cose_sign1, target_size));

  EXPECT_THAT(EncodeCoseSign1TaggedStructure(cose_sign1), SizeIs(target_size));
}

INSTANTIATE_TEST_SUITE_P(
    PadOverallCoseSign1TaggedStructureTest,
    PadOverallCoseSign1TaggedStructureTest,
    Values(GetInitialSizeForPadCoseSign1TaggedStructure(),
           GetInitialSizeForPadCoseSign1TaggedStructure() + 24,
           GetInitialSizeForPadCoseSign1TaggedStructure() + 256,
           GetInitialSizeForPadCoseSign1TaggedStructure() + 65536,
           GetInitialSizeForPadCoseSign1TaggedStructure() + 1000000));
}  // namespace
}  // namespace credentio