// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0 OR ISC

#include <gtest/gtest.h>
#include <openssl/evp.h>
#include <openssl/obj.h>

#include "../test/file_test.h"

// ML-DSA parameter sets
struct MLDSAParamSet {
  const char name[20];
  const int nid;
};

static const struct MLDSAParamSet kMLDSAs[] = {{"MLDSA44", NID_MLDSA44},
                                               {"MLDSA65", NID_MLDSA65},
                                               {"MLDSA87", NID_MLDSA87}};

class MLDSATest : public testing::TestWithParam<MLDSAParamSet> {};

INSTANTIATE_TEST_SUITE_P(All, MLDSATest, testing::ValuesIn(kMLDSAs),
                         [](const testing::TestParamInfo<MLDSAParamSet> &params)
                             -> std::string { return params.param.name; });

TEST_P(MLDSATest, ExpandedKeyValidation) {
  const MLDSAParamSet ps = GetParam();

  // This test verifies that we reject invalid extended keys, because they can
  // cause undefined behavior including producing unverifiable signatures.
  //
  // Test vectors are generated by make_corrupted_key_tests.cc which uses
  // internal ML-DSA functions to corrupt keys in specific ways.

  FileTestGTest("crypto/evp_extra/mldsa_corrupted_key_tests.txt",
                [&](FileTest *t) {
                  std::string param_set;
                  ASSERT_TRUE(t->GetInstruction(&param_set, "ParamSet"));

                  // Skip test vectors for other parameter sets
                  if (param_set != ps.name) {
                    t->SkipCurrent();
                    return;
                  }

                  std::vector<uint8_t> corrupted_key;
                  ASSERT_TRUE(t->GetBytes(&corrupted_key, "CorruptedKey"));

                  // Try to import the corrupted key - it should fail
                  bssl::UniquePtr<EVP_PKEY> corrupted_pkey(
                      EVP_PKEY_pqdsa_new_raw_private_key(
                          ps.nid, corrupted_key.data(), corrupted_key.size()));

                  EXPECT_FALSE(corrupted_pkey.get())
                      << "Imported corrupted " << ps.name << " key";
                });
}