// 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 "content/browser/webid/delegation/jwt_signer.h"

#include "base/base64.h"
#include "base/base64url.h"
#include "base/compiler_specific.h"
#include "base/containers/span.h"
#include "base/containers/to_vector.h"
#include "base/functional/callback.h"
#include "base/json/json_reader.h"
#include "base/json/json_writer.h"
#include "base/logging.h"
#include "base/values.h"
#include "content/browser/webid/delegation/sd_jwt.h"
#include "crypto/keypair.h"
#include "crypto/random.h"
#include "crypto/sha2.h"
#include "crypto/sign.h"
#include "crypto/signature_verifier.h"
#include "testing/gmock/include/gmock/gmock.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "third_party/boringssl/src/include/openssl/bn.h"
#include "third_party/boringssl/src/include/openssl/curve25519.h"
#include "third_party/boringssl/src/include/openssl/ecdsa.h"
#include "third_party/boringssl/src/include/openssl/mem.h"
#include "url/gurl.h"
#include "url/origin.h"

namespace content::sdjwt {

using ::testing::NiceMock;

namespace {
std::vector<uint8_t> TestSha256(std::string_view data) {
  std::string str = crypto::SHA256HashString(data);
  std::vector<uint8_t> result(str.begin(), str.end());
  return result;
}
}  // namespace

class JwtSignerTest : public testing::Test {
 protected:
  JwtSignerTest() = default;
  ~JwtSignerTest() override = default;

  void SetUp() override {}

  void TearDown() override {}
};

void VerifyEs256(const std::vector<uint8_t>& public_key,
                 const base::span<const uint8_t>& signature,
                 const std::string& message) {
  const size_t kMaxBytesPerBN = 32;
  EXPECT_EQ(signature.size(), 2 * kMaxBytesPerBN);
  base::span<const uint8_t> r_bytes = signature.first(kMaxBytesPerBN);
  base::span<const uint8_t> s_bytes = signature.subspan(kMaxBytesPerBN);

  bssl::UniquePtr<ECDSA_SIG> ecdsa_sig(ECDSA_SIG_new());
  EXPECT_TRUE(ecdsa_sig);

  EXPECT_TRUE(BN_bin2bn(r_bytes.data(), r_bytes.size(), ecdsa_sig->r));
  EXPECT_TRUE(BN_bin2bn(s_bytes.data(), s_bytes.size(), ecdsa_sig->s));

  uint8_t* der;
  size_t der_len;
  EXPECT_TRUE(ECDSA_SIG_to_bytes(&der, &der_len, ecdsa_sig.get()));

  // Frees memory allocated by `ECDSA_SIG_to_bytes()`.
  bssl::UniquePtr<uint8_t> delete_signature(der);

  // SAFETY: `ECDSA_SIG_to_bytes()` uses a C-style API to allocate a new buffer.
  auto signature_span = UNSAFE_BUFFERS(base::span<uint8_t>(der, der_len));
  auto der_signature = base::ToVector(signature_span);

  crypto::SignatureVerifier verifier;
  EXPECT_TRUE(verifier.VerifyInit(crypto::SignatureVerifier::ECDSA_SHA256,
                                  der_signature, public_key));

  verifier.VerifyUpdate(base::as_byte_span(message));
  EXPECT_TRUE(verifier.VerifyFinal());
}

void VerifyRs256(const std::vector<uint8_t>& public_key,
                 base::span<const uint8_t> signature,
                 const std::string& message) {
  crypto::SignatureVerifier verifier;
  EXPECT_TRUE(verifier.VerifyInit(crypto::SignatureVerifier::RSA_PKCS1_SHA256,
                                  signature, public_key));

  verifier.VerifyUpdate(base::as_byte_span(message));
  EXPECT_TRUE(verifier.VerifyFinal());
}

void VerifyEdDsa(const crypto::keypair::PublicKey& public_key,
                 base::span<const uint8_t> signature,
                 const std::string& message) {
  EXPECT_TRUE(crypto::sign::Verify(crypto::sign::SignatureKind::ED25519,
                                   public_key, base::as_byte_span(message),
                                   signature));
}

TEST_F(JwtSignerTest, JwtSigner) {
  auto private_key = crypto::keypair::PrivateKey::GenerateEcP256();
  std::vector<uint8_t> public_key = private_key.ToSubjectPublicKeyInfo();

  const std::string message = "hello wold";
  auto signer = CreateJwtSigner(std::move(private_key));
  auto signature = std::move(signer).Run(message);

  VerifyEs256(public_key, base::as_byte_span(*signature), message);
}

TEST_F(JwtSignerTest, JwtSignerRs256) {
  auto private_key = crypto::keypair::PrivateKey::GenerateRsa2048();
  std::vector<uint8_t> public_key = private_key.ToSubjectPublicKeyInfo();

  const std::string message = "hello wold";
  auto signer = CreateJwtSigner(std::move(private_key));
  auto signature = std::move(signer).Run(message);

  VerifyRs256(public_key, base::as_byte_span(*signature), message);
}

TEST_F(JwtSignerTest, ExportPublicKeyRs256) {
  auto private_key = crypto::keypair::PrivateKey::GenerateRsa2048();
  auto jwk = ExportPublicKey(private_key);
  ASSERT_TRUE(jwk);
  EXPECT_EQ(jwk->kty, "RSA");
  EXPECT_EQ(jwk->alg, "RS256");
  EXPECT_TRUE(jwk->crv.empty());
  EXPECT_FALSE(jwk->n.empty());
  EXPECT_FALSE(jwk->e.empty());
}

TEST_F(JwtSignerTest, CreateJwt) {
  auto private_key = crypto::keypair::PrivateKey::GenerateEcP256();
  std::vector<uint8_t> public_key = private_key.ToSubjectPublicKeyInfo();

  Header header;
  header.typ = "jwt";
  header.alg = "ES256";

  Payload payload;
  payload.iss = "https://issuer.example";
  payload.sub = "goto@google.com";

  Jwt issued;
  issued.header = JSONString(header.Serialize()->value());
  issued.payload = JSONString(payload.Serialize()->value());

  auto success = issued.Sign(CreateJwtSigner(std::move(private_key)));

  EXPECT_TRUE(success);

  auto signature = base::Base64UrlDecode(
      issued.signature.value(), base::Base64UrlDecodePolicy::IGNORE_PADDING);

  EXPECT_TRUE(signature);

  std::string header_base64;
  base::Base64UrlEncode(issued.header.value(),
                        base::Base64UrlEncodePolicy::OMIT_PADDING,
                        &header_base64);

  std::string payload_base64;
  base::Base64UrlEncode(issued.payload.value(),
                        base::Base64UrlEncodePolicy::OMIT_PADDING,
                        &payload_base64);

  std::string message = header_base64 + "." + payload_base64;
  VerifyEs256(public_key, base::as_byte_span(*signature), message);
}

TEST_F(JwtSignerTest, ExportPublicKeyEc256) {
  auto private_key = crypto::keypair::PrivateKey::GenerateEcP256();
  auto jwk = ExportPublicKey(private_key);
  ASSERT_TRUE(jwk);
  EXPECT_EQ(jwk->kty, "EC");
  EXPECT_EQ(jwk->alg, "ES256");
  EXPECT_EQ(jwk->crv, "P-256");
  EXPECT_FALSE(jwk->x.empty());
  EXPECT_FALSE(jwk->y.empty());
}

TEST_F(JwtSignerTest, CreateSdJwtKb) {
  auto holder_private_key = crypto::keypair::PrivateKey::GenerateEcP256();
  auto jwk = ExportPublicKey(holder_private_key);

  auto issuer_private_key = crypto::keypair::PrivateKey::GenerateEcP256();

  Header header;
  header.typ = "jwt";
  header.alg = "ES256";

  Payload payload;
  payload.iss = "https://issuer.example";

  Disclosure name;
  name.salt = Disclosure::CreateSalt();
  name.name = "name";
  name.value = "Sam";

  payload._sd = {*name.Digest(base::BindRepeating(TestSha256))};

  ConfirmationKey confirmation;
  confirmation.jwk = *jwk;
  payload.cnf = confirmation;

  auto issuer_json = ExportPublicKey(issuer_private_key);

  EXPECT_TRUE(issuer_json);
  EXPECT_TRUE(issuer_json->Serialize());

  Jwt issued;
  issued.header = JSONString(header.Serialize()->value());
  issued.payload = JSONString(payload.Serialize()->value());
  auto signer = CreateJwtSigner(std::move(issuer_private_key));
  auto success = issued.Sign(std::move(signer));

  EXPECT_TRUE(success);

  auto presentation = SdJwt::Disclose(
      {{name.name, JSONString(name.Serialize().value())}}, {"name"});
  EXPECT_TRUE(presentation);

  SdJwt sd_jwt;
  sd_jwt.jwt = issued;
  sd_jwt.disclosures = *presentation;

  std::optional<SdJwtKb> sd_jwt_kb = SdJwtKb::Create(
      sd_jwt, "https://verifier.example", "__fake_nonce__",
      base::Time::FromTimeT(1234), base::BindRepeating(TestSha256),
      CreateJwtSigner(std::move(holder_private_key)));

  EXPECT_TRUE(sd_jwt_kb);
}

TEST_F(JwtSignerTest, InvalidKeyType) {
  // Create a key that is not the supported EC, RSA or ED methods.
  auto private_key = crypto::keypair::PrivateKey::GenerateEcP384();

  // Test CreateJwtSigner with an EcP384 key type.
  auto signer = CreateJwtSigner(private_key);
  auto signature = std::move(signer).Run("message");
  EXPECT_FALSE(signature);

  // Test ExportPublicKey with an EcP384 key type.
  auto jwk = ExportPublicKey(private_key);
  EXPECT_FALSE(jwk);
}

TEST_F(JwtSignerTest, SignWithEd25519AndExportPublicKey) {
  auto private_key = crypto::keypair::PrivateKey::GenerateEd25519();

  // Test CreateJwtSigner with an Ed25519 key type.
  auto signer = CreateJwtSigner(private_key);
  auto signature = std::move(signer).Run("message");
  EXPECT_TRUE(signature);

  // Test ExportPublicKey with an Ed25519 key type.
  auto jwk = ExportPublicKey(private_key);
  EXPECT_TRUE(jwk);
}

TEST_F(JwtSignerTest, MismatchedKeyTypes) {
  // Test SignJwtRs256 with an EC key.
  auto ec_private_key = crypto::keypair::PrivateKey::GenerateEcP256();
  auto rs256_signer = CreateJwtSigner(std::move(ec_private_key));
  auto rs256_signature = std::move(rs256_signer).Run("message");
  EXPECT_TRUE(rs256_signature);

  // Test SignJwtEs256 with an RSA key.
  auto rsa_private_key = crypto::keypair::PrivateKey::GenerateRsa2048();
  auto es256_signer = CreateJwtSigner(std::move(rsa_private_key));
  auto es256_signature = std::move(es256_signer).Run("message");
  EXPECT_TRUE(es256_signature);
}

TEST_F(JwtSignerTest, JwtSignerEdDsa) {
  auto private_key = crypto::keypair::PrivateKey::GenerateEd25519();
  auto public_key_raw = private_key.ToEd25519PublicKey();
  auto public_key =
      crypto::keypair::PublicKey::FromEd25519PublicKey(public_key_raw);

  const std::string message = "hello wold";
  auto signer = CreateJwtSigner(std::move(private_key));
  auto signature = std::move(signer).Run(message);

  ASSERT_TRUE(signature);

  VerifyEdDsa(public_key, base::as_byte_span(*signature), message);
}

TEST_F(JwtSignerTest, JwtVerifierEdDsa) {
  auto private_key = crypto::keypair::PrivateKey::GenerateEd25519();
  auto jwk = ExportPublicKey(private_key);
  ASSERT_TRUE(jwk);

  const std::string message = "hello world";
  auto signer = CreateJwtSigner(std::move(private_key));
  auto signature = std::move(signer).Run(message);

  ASSERT_TRUE(signature);

  Header header;
  header.alg = "EdDSA";
  auto verifier = CreateJwtVerifier(*jwk, header);
  EXPECT_TRUE(std::move(verifier).Run(message, base::as_byte_span(*signature)));
}

TEST_F(JwtSignerTest, JwtVerifierEs256) {
  auto private_key = crypto::keypair::PrivateKey::GenerateEcP256();
  auto jwk = ExportPublicKey(private_key);
  ASSERT_TRUE(jwk);

  const std::string message = "hello world";
  auto signer = CreateJwtSigner(std::move(private_key));
  auto signature = std::move(signer).Run(message);

  ASSERT_TRUE(signature);

  Header header;
  header.alg = "ES256";
  auto verifier = CreateJwtVerifier(*jwk, header);
  EXPECT_TRUE(std::move(verifier).Run(message, base::as_byte_span(*signature)));
}

TEST_F(JwtSignerTest, JwtVerifierRs256) {
  auto private_key = crypto::keypair::PrivateKey::GenerateRsa2048();
  auto jwk = ExportPublicKey(private_key);
  ASSERT_TRUE(jwk);

  const std::string message = "hello world";
  auto signer = CreateJwtSigner(std::move(private_key));
  auto signature = std::move(signer).Run(message);

  ASSERT_TRUE(signature);

  Header header;
  header.alg = "RS256";
  auto verifier = CreateJwtVerifier(*jwk, header);
  EXPECT_TRUE(std::move(verifier).Run(message, base::as_byte_span(*signature)));
}

TEST_F(JwtSignerTest, JwtVerificationUnsupported) {
  Jwk jwk;
  jwk.kty = "unknown";

  Jwt token;
  token.header = JSONString("header");
  token.payload = JSONString("payload");
  token.signature = Base64String("signature");

  Header header;
  header.alg = "unknown";
  auto verifier = CreateJwtVerifier(jwk, header);
  EXPECT_FALSE(token.Verify(std::move(verifier)));
}

TEST_F(JwtSignerTest, ExportPublicKeyEdDsa) {
  crypto::keypair::PrivateKey private_key =
      crypto::keypair::PrivateKey::GenerateEd25519();
  std::optional<Jwk> jwk = ExportPublicKey(private_key);

  ASSERT_TRUE(jwk.has_value());
  EXPECT_EQ(jwk->kty, "OKP");
  EXPECT_EQ(jwk->crv, "Ed25519");
  EXPECT_EQ(jwk->alg, "EdDSA");
  EXPECT_FALSE(jwk->x.empty());

  std::array<uint8_t, 32> public_key_raw = private_key.ToEd25519PublicKey();
  std::string expected_x;
  base::Base64UrlEncode(public_key_raw,
                        base::Base64UrlEncodePolicy::OMIT_PADDING, &expected_x);
  EXPECT_EQ(jwk->x, expected_x);
}

}  // namespace content::sdjwt