// 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 "claim/local_signer.h"

#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/functional/any_invocable.h"
#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/notification.h"
#include "async/async_context.h"
#include "claim/signer.h"
#include "cose/sig_structure.h"
#include "cose/verifier.h"
#include "crypto/algorithms.h"
#include "crypto/crypto_read_handler.h"
#include "crypto/crypto_write_handler.h"
#include "crypto/default/default_crypto_read_handler.h"
#include "crypto/default/hasher.h"
#include "crypto/timestamp_write_handler.h"
#include "gmock/gmock.h"
#include "google/protobuf/util/message_differencer.h"
#include "gtest/gtest.h"
#include "proto/cose_verification_result.pb.h"
#include "proto/signature_info.pb.h"
#include "testing/create_test_crypto_write_handler.h"
#include "testing/local_key_crypto_write_handler.h"
#include "testing/proto_test_utils.h"

namespace credentio {
namespace {

using ::absl_testing::StatusIs;
using ::credentio_testing::ParseTextProtoOrDie;
using ::testing::ElementsAre;
using ::testing::HasSubstr;
using ::testing::IsFalse;
using ::testing::IsTrue;

MATCHER_P(EqualsProtoIgnoringTimestampAndOrdering, expected, "") {
  google::protobuf::util::MessageDifferencer differencer;
  differencer.set_repeated_field_comparison(
      google::protobuf::util::MessageDifferencer::AS_SET);
  differencer.IgnoreField(
      SignatureInfo::descriptor()->FindFieldByName("timestamping_time"));
  std::string diff;
  differencer.ReportDifferencesToString(&diff);
  bool match = differencer.Compare(expected, arg);
  if (!match) {
    *result_listener << "\nDifference:\n" << diff;
  }
  return match;
}

absl::StatusOr<std::unique_ptr<CryptoReadHandler>>
CreateTrustAnyoneCryptoReadHandlerForTest() {
  return CreateDefaultCryptoReadHandler({
      .trust_environment = CryptoReadHandler::TrustEnvironment::kAny,
  });
}

absl::StatusOr<std::string> Sign(LocalClaimSigner& signer,
                                 absl::string_view claim,
                                 ClaimSigner::Options options = {}) {
  absl::Notification done;
  absl::StatusOr<std::string> signature =
      absl::InternalError("Callback not called");
  {
    std::string claim_copy(claim);
    signer.Sign(claim_copy, std::move(options),
                [&](absl::StatusOr<ClaimSigner::SignResult> cb_result) {
                  if (cb_result.ok()) {
                    signature = std::move(cb_result->signature);
                  } else {
                    signature = cb_result.status();
                  }
                  done.Notify();
                });
  }
  done.WaitForNotification();
  return signature;
}

TEST(ClaimSigner, Sign) {
  absl::StatusOr<std::unique_ptr<CryptoWriteHandler>> crypto_writer =
      CreateTestCryptoWriteHandler({});
  ABSL_ASSERT_OK(crypto_writer);
  Sha256HasherFactory hasher_factory;
  LocalClaimSigner claim_signer(crypto_writer->get(), &hasher_factory);
  absl::StatusOr<std::string> signature = Sign(claim_signer, "abc");
  ABSL_ASSERT_OK(signature);

  absl::StatusOr<std::unique_ptr<CryptoReadHandler>>
      crypto_read_handler_for_verifier =
          CreateTrustAnyoneCryptoReadHandlerForTest();
  ABSL_ASSERT_OK(crypto_read_handler_for_verifier);
  std::unique_ptr<CoseVerifier> cose_verifier = CreateCoseVerifier(
      {.crypto_read_handler = std::move(*crypto_read_handler_for_verifier)});
  CoseVerificationResult result = cose_verifier->Verify(*signature, "abc");
  CoseVerificationResult expected_result = ParseTextProtoOrDie(R"pb(
    signature_info {
      cert_serial_number: "02"
      issuer: "OU=Test Unit,O=Test Org,CN=Test Claim Signer"
      certificate_issuer: "OU=Test Unit,O=Test Org,CN=Test Root CA"
    }
    validation_status {
      successes { code: "timeStamp.validated" url: "c2pa.signature" }
      successes { code: "timeStamp.trusted" url: "c2pa.signature" }
      successes { code: "signingCredential.trusted" url: "c2pa.signature" }
      successes { code: "claimSignature.insideValidity" url: "c2pa.signature" }
      successes { code: "claimSignature.validated" url: "c2pa.signature" }
      informationals {
        code: "signingCredential.ocsp.skipped"
        url: "c2pa.signature"
      }
    }
  )pb");
  EXPECT_THAT(result, EqualsProtoIgnoringTimestampAndOrdering(expected_result));
}

// This signer produces signatures that are larger than its advertised
// max_signature_byte_size.
class InvalidFixedKeySigner : public FixedKeySigner {
 public:
  SigningAlgorithm algorithm() const override {
    return SigningAlgorithm::kPs256;
  }
  const std::vector<std::string>& certificate_chain() const override {
    return certificate_chain_;
  }
  const std::vector<std::string>& ocsp_responses() const override {
    return ocsp_responses_;
  }
  uint64_t max_signature_byte_size() const override { return 100; }
  absl::StatusOr<SignResponse> Sign(
      absl::string_view to_be_signed) const override {
    return SignResponse{.signature = std::string(200, '\0')};
  }

 private:
  std::vector<std::string> certificate_chain_ = {"cert1", "cert2"};
  std::vector<std::string> ocsp_responses_;
};

class InvalidCryptoWriteHandler : public CryptoWriteHandler {
 public:
  const FixedKeySigner& GetSigner() override { return signer_; }
  void CreateTimestamp(HashAlgorithm hash_algorithm, std::string hashed_message,
                       TimestampWriteHandler::Options options,
                       absl::AnyInvocable<void(absl::StatusOr<std::string>) &&>
                           cb) const override {
    AsyncResolver<std::string>(options.async_context, std::move(cb))
        .Return(std::string(100, '\0'));
  }
  uint64_t max_timestamp_byte_size(
      HashAlgorithm hash_algorithm) const override {
    return 100;
  }

 private:
  InvalidFixedKeySigner signer_;
};

TEST(ClaimSigner, SignatureTooLarge) {
  InvalidCryptoWriteHandler crypto_writer;
  Sha256HasherFactory hasher_factory;
  LocalClaimSigner claim_signer(&crypto_writer, &hasher_factory);
  EXPECT_THAT(Sign(claim_signer, "abc"), StatusIs(absl::StatusCode::kInternal,
                                                  HasSubstr("Signature size")));
}

class DelegatingCryptoWriteHandler : public CryptoWriteHandler {
 public:
  explicit DelegatingCryptoWriteHandler(CryptoWriteHandler* delegate)
      : delegate_(delegate) {}

  const FixedKeySigner& GetSigner() override { return delegate_->GetSigner(); }

  void CreateTimestamp(HashAlgorithm hash_algorithm, std::string hashed_message,
                       TimestampWriteHandler::Options options,
                       absl::AnyInvocable<void(absl::StatusOr<std::string>) &&>
                           cb) const override {
    delegate_->CreateTimestamp(hash_algorithm, std::move(hashed_message),
                               options, std::move(cb));
  }

  uint64_t max_timestamp_byte_size(
      HashAlgorithm hash_algorithm) const override {
    return delegate_->max_timestamp_byte_size(hash_algorithm);
  }

 protected:
  CryptoWriteHandler* delegate_;
};

class InvalidTimestampCryptoWriteHandler : public DelegatingCryptoWriteHandler {
 public:
  using DelegatingCryptoWriteHandler::DelegatingCryptoWriteHandler;

  uint64_t max_timestamp_byte_size(
      HashAlgorithm hash_algorithm) const override {
    return 10;
  }
};

TEST(ClaimSigner, TimestampTooLarge) {
  absl::StatusOr<std::unique_ptr<CryptoWriteHandler>> crypto_writer =
      CreateTestCryptoWriteHandler({});
  ABSL_ASSERT_OK(crypto_writer);
  InvalidTimestampCryptoWriteHandler invalid_crypto_writer(
      crypto_writer->get());
  Sha256HasherFactory hasher_factory;
  LocalClaimSigner claim_signer(&invalid_crypto_writer, &hasher_factory);
  absl::StatusOr<std::string> signature = Sign(claim_signer, "abc");
  ABSL_ASSERT_OK(signature);
  absl::StatusOr<CoseSign1TaggedStructure> cose_sign1 =
      DecodeCoseSign1TaggedStructure(*signature);
  ABSL_ASSERT_OK(cose_sign1);
  EXPECT_THAT(cose_sign1->unprotected_header.sig_tst2.has_value(), IsFalse());
}

TEST(ClaimSigner, TimestampTooLargeTimestampingRequired) {
  absl::StatusOr<std::unique_ptr<CryptoWriteHandler>> crypto_writer =
      CreateTestCryptoWriteHandler({});
  ABSL_ASSERT_OK(crypto_writer);
  InvalidTimestampCryptoWriteHandler invalid_crypto_writer(
      crypto_writer->get());
  Sha256HasherFactory hasher_factory;
  LocalClaimSigner claim_signer(&invalid_crypto_writer, &hasher_factory);
  EXPECT_THAT(
      Sign(claim_signer, "abc",
           {.time_stamping_mode = TimeStampingMode::kRequired}),
      StatusIs(absl::StatusCode::kInternal, HasSubstr("Timestamp size")));
}

class FixedMaxTimestampCryptoWriteHandler
    : public DelegatingCryptoWriteHandler {
 public:
  FixedMaxTimestampCryptoWriteHandler(CryptoWriteHandler* delegate,
                                      uint64_t max_timestamp_size)
      : DelegatingCryptoWriteHandler(delegate),
        max_timestamp_size_(max_timestamp_size) {}

  uint64_t max_timestamp_byte_size(
      HashAlgorithm hash_algorithm) const override {
    return max_timestamp_size_;
  }

 private:
  uint64_t max_timestamp_size_;
};

class UnavailableTimestampCryptoWriteHandler
    : public DelegatingCryptoWriteHandler {
 public:
  UnavailableTimestampCryptoWriteHandler(CryptoWriteHandler* delegate,
                                         uint64_t max_timestamp_size)
      : DelegatingCryptoWriteHandler(delegate),
        max_timestamp_size_(max_timestamp_size) {}

  void CreateTimestamp(HashAlgorithm hash_algorithm, std::string hashed_message,
                       TimestampWriteHandler::Options options,
                       absl::AnyInvocable<void(absl::StatusOr<std::string>) &&>
                           cb) const override {
    AsyncResolver<std::string>(options.async_context, std::move(cb))
        .Return(absl::UnavailableError("Timestamping is unavailable."));
  }

  uint64_t max_timestamp_byte_size(
      HashAlgorithm hash_algorithm) const override {
    return max_timestamp_size_;
  }

 private:
  uint64_t max_timestamp_size_;
};

TEST(ClaimSignerTest, TimeStampingRequired) {
  absl::StatusOr<std::unique_ptr<CryptoWriteHandler>> delegate =
      CreateTestCryptoWriteHandler({});
  ABSL_ASSERT_OK(delegate);
  UnavailableTimestampCryptoWriteHandler crypto_writer(delegate->get(), 100);
  Sha256HasherFactory hasher_factory;
  LocalClaimSigner claim_signer(&crypto_writer, &hasher_factory);
  EXPECT_THAT(Sign(claim_signer, "abc",
                   {.time_stamping_mode = TimeStampingMode::kRequired}),
              StatusIs(absl::StatusCode::kUnavailable,
                       HasSubstr("Timestamping is unavailable.")));
}

TEST(ClaimSigner, TimeStampingOmittedWithPadding) {
  absl::StatusOr<std::unique_ptr<CryptoWriteHandler>> delegate =
      CreateTestCryptoWriteHandler({});
  ABSL_ASSERT_OK(delegate);
  FixedMaxTimestampCryptoWriteHandler crypto_writer(delegate->get(), 2000);
  Sha256HasherFactory hasher_factory;
  LocalClaimSigner claim_signer(&crypto_writer, &hasher_factory);
  absl::StatusOr<std::string> signature =
      Sign(claim_signer, "abc",
           {.time_stamping_mode = TimeStampingMode::kOmittedWithPadding});
  ABSL_ASSERT_OK(signature);
  absl::StatusOr<CoseSign1TaggedStructure> cose_sign1 =
      DecodeCoseSign1TaggedStructure(*signature);
  ABSL_ASSERT_OK(cose_sign1);
  EXPECT_FALSE(cose_sign1->signature.empty());
  EXPECT_THAT(cose_sign1->unprotected_header.sig_tst2.has_value(), IsFalse());

  // Check if the padding is of correct size.
  absl::StatusOr<std::string> signature2 = Sign(
      claim_signer, "abc", {.time_stamping_mode = TimeStampingMode::kRequired});
  ABSL_ASSERT_OK(signature2);
  EXPECT_EQ(signature->size(), signature2->size());
}

class ClaimSignerTest : public testing::TestWithParam<int64_t> {};

TEST_P(ClaimSignerTest, SignatureSizeUnaffectedByTimestamp) {
  absl::StatusOr<std::unique_ptr<CryptoWriteHandler>> delegate =
      CreateTestCryptoWriteHandler({});
  ABSL_ASSERT_OK(delegate);
  int64_t timestamp_size = GetParam();
  UnavailableTimestampCryptoWriteHandler unavailable_crypto_writer(
      delegate->get(), timestamp_size);
  Sha256HasherFactory hasher_factory;
  LocalClaimSigner claim_signer(&unavailable_crypto_writer, &hasher_factory);
  absl::StatusOr<std::string> signature1 = Sign(claim_signer, "abc");
  ABSL_ASSERT_OK(signature1);

  FixedMaxTimestampCryptoWriteHandler normal_crypto_writer(delegate->get(),
                                                           timestamp_size);
  LocalClaimSigner claim_signer2(&normal_crypto_writer, &hasher_factory);
  absl::StatusOr<std::string> signature2 = Sign(claim_signer2, "abc");
  ABSL_ASSERT_OK(signature2);
  EXPECT_EQ(signature1->size(), signature2->size());
}
INSTANTIATE_TEST_SUITE_P(ClaimSignerTestWithParams, ClaimSignerTest,
                         testing::Values(100, 200, 500, 1000, 10000));

class FixedKeySignerWithOcsp : public FixedKeySigner {
 public:
  FixedKeySignerWithOcsp(const FixedKeySigner& delegate,
                         std::vector<std::string> ocsp_responses)
      : delegate_(delegate), ocsp_responses_(std::move(ocsp_responses)) {}
  SigningAlgorithm algorithm() const override { return delegate_.algorithm(); }
  const std::vector<std::string>& certificate_chain() const override {
    return delegate_.certificate_chain();
  }
  const std::vector<std::string>& ocsp_responses() const override {
    return ocsp_responses_;
  }
  uint64_t max_signature_byte_size() const override {
    return delegate_.max_signature_byte_size();
  }
  absl::StatusOr<SignResponse> Sign(
      absl::string_view to_be_signed) const override {
    return delegate_.Sign(to_be_signed);
  }

 private:
  const FixedKeySigner& delegate_;
  std::vector<std::string> ocsp_responses_;
};

class CryptoWriteHandlerWithOcsp : public DelegatingCryptoWriteHandler {
 public:
  CryptoWriteHandlerWithOcsp(CryptoWriteHandler* delegate,
                             std::vector<std::string> ocsp_responses)
      : DelegatingCryptoWriteHandler(delegate),
        signer_(delegate->GetSigner(), std::move(ocsp_responses)) {}
  const FixedKeySigner& GetSigner() override { return signer_; }

 private:
  FixedKeySignerWithOcsp signer_;
};

TEST(ClaimSigner, OcspResponsesWithTimestamp) {
  absl::StatusOr<std::unique_ptr<CryptoWriteHandler>> delegate =
      CreateTestCryptoWriteHandler({});
  ABSL_ASSERT_OK(delegate);
  CryptoWriteHandlerWithOcsp crypto_writer(delegate->get(), {"ocsp1", "ocsp2"});
  Sha256HasherFactory hasher_factory;
  LocalClaimSigner claim_signer(&crypto_writer, &hasher_factory);
  absl::StatusOr<std::string> signature = Sign(claim_signer, "abc");
  ABSL_ASSERT_OK(signature);

  absl::StatusOr<CoseSign1TaggedStructure> cose_sign1 =
      DecodeCoseSign1TaggedStructure(*signature);
  ABSL_ASSERT_OK(cose_sign1);
  EXPECT_THAT(cose_sign1->unprotected_header.ocsp_responses,
              ElementsAre("ocsp1", "ocsp2"));
  EXPECT_THAT(cose_sign1->unprotected_header.sig_tst2.has_value(), IsTrue());
}

TEST(ClaimSigner, OcspResponsesWithoutTimestamp) {
  absl::StatusOr<std::unique_ptr<CryptoWriteHandler>> delegate =
      CreateTestCryptoWriteHandler({.disable_timestamping = true});
  ABSL_ASSERT_OK(delegate);
  CryptoWriteHandlerWithOcsp crypto_writer(delegate->get(), {"ocsp1", "ocsp2"});
  Sha256HasherFactory hasher_factory;
  LocalClaimSigner claim_signer(&crypto_writer, &hasher_factory);
  absl::StatusOr<std::string> signature = Sign(claim_signer, "abc");
  ABSL_ASSERT_OK(signature);

  absl::StatusOr<CoseSign1TaggedStructure> cose_sign1 =
      DecodeCoseSign1TaggedStructure(*signature);
  ABSL_ASSERT_OK(cose_sign1);
  EXPECT_THAT(cose_sign1->unprotected_header.ocsp_responses,
              ElementsAre("ocsp1", "ocsp2"));
  EXPECT_THAT(cose_sign1->unprotected_header.sig_tst2.has_value(), IsFalse());
}

TEST(ClaimSigner, PadToSize) {
  absl::StatusOr<std::unique_ptr<CryptoWriteHandler>> crypto_writer =
      CreateTestCryptoWriteHandler({});
  ABSL_ASSERT_OK(crypto_writer);
  Sha256HasherFactory hasher_factory;
  LocalClaimSigner claim_signer(crypto_writer->get(), &hasher_factory,
                                /*pad_to_size=*/99999);
  absl::StatusOr<std::string> signature = Sign(claim_signer, "abc");
  ABSL_ASSERT_OK(signature);
  EXPECT_EQ(signature->size(), 99999);
}

}  // namespace
}  // namespace credentio