/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
/* vim: set ts=2 et sw=2 tw=80: */
/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this file,
 * You can obtain one at http://mozilla.org/MPL/2.0/. */

#include <climits>
#include <memory>
#include "nss.h"
#include "pk11pub.h"
#include "secutil.h"

#include "gtest/gtest.h"
#include "nss_scoped_ptrs.h"

namespace nss_test {

const std::vector<uint8_t> kValidP256Key = {
    0x30, 0x81, 0x87, 0x02, 0x01, 0x00, 0x30, 0x13, 0x06, 0x07, 0x2a, 0x86,
    0x48, 0xce, 0x3d, 0x02, 0x01, 0x06, 0x08, 0x2a, 0x86, 0x48, 0xce, 0x3d,
    0x03, 0x01, 0x07, 0x04, 0x6d, 0x30, 0x6b, 0x02, 0x01, 0x01, 0x04, 0x20,
    0xc9, 0xaf, 0xa9, 0xd8, 0x45, 0xba, 0x75, 0x16, 0x6b, 0x5c, 0x21, 0x57,
    0x67, 0xb1, 0xd6, 0x93, 0x4e, 0x50, 0xc3, 0xdb, 0x36, 0xe8, 0x9b, 0x12,
    0x7b, 0x8a, 0x62, 0x2b, 0x12, 0x0f, 0x67, 0x21, 0xa1, 0x44, 0x03, 0x42,
    0x00, 0x04, 0x60, 0xfe, 0xd4, 0xba, 0x25, 0x5a, 0x9d, 0x31, 0xc9, 0x61,
    0xeb, 0x74, 0xc6, 0x35, 0x6d, 0x68, 0xc0, 0x49, 0xb8, 0x92, 0x3b, 0x61,
    0xfa, 0x6c, 0xe6, 0x69, 0x62, 0x2e, 0x60, 0xf2, 0x9f, 0xb6, 0x79, 0x03,
    0xfe, 0x10, 0x08, 0xb8, 0xbc, 0x99, 0xa4, 0x1a, 0xe9, 0xe9, 0x56, 0x28,
    0xbc, 0x64, 0xf2, 0xf1, 0xb2, 0x0c, 0x2d, 0x7e, 0x9f, 0x51, 0x77, 0xa3,
    0xc2, 0x94, 0xd4, 0x46, 0x22, 0x99};

const std::vector<uint8_t> kValidRSAKey = {
    // 512-bit RSA private key (PKCS#8)
    0x30, 0x82, 0x01, 0x54, 0x02, 0x01, 0x00, 0x30, 0x0d, 0x06, 0x09, 0x2a,
    0x86, 0x48, 0x86, 0xf7, 0x0d, 0x01, 0x01, 0x01, 0x05, 0x00, 0x04, 0x82,
    0x01, 0x3e, 0x30, 0x82, 0x01, 0x3a, 0x02, 0x01, 0x00, 0x02, 0x41, 0x00,
    0xa2, 0x40, 0xce, 0xb5, 0x4e, 0x70, 0xdc, 0x14, 0x82, 0x5b, 0x58, 0x7d,
    0x2f, 0x5d, 0xfd, 0x46, 0x3c, 0x4b, 0x82, 0x50, 0xb6, 0x96, 0x00, 0x4a,
    0x1a, 0xca, 0xaf, 0xe4, 0x9b, 0xcf, 0x38, 0x4a, 0x46, 0xaa, 0x9f, 0xb4,
    0xd9, 0xc7, 0xee, 0x88, 0xe9, 0xef, 0x0a, 0x31, 0x5f, 0x53, 0x86, 0x8f,
    0x63, 0x68, 0x0b, 0x58, 0x34, 0x72, 0x49, 0xba, 0xed, 0xd9, 0x34, 0x15,
    0x16, 0xc4, 0xca, 0xb7, 0x02, 0x03, 0x01, 0x00, 0x01, 0x02, 0x40, 0x34,
    0xe6, 0xdc, 0x7e, 0xd0, 0xec, 0x8b, 0x55, 0x44, 0x8b, 0x73, 0xf6, 0x9d,
    0x13, 0x10, 0x19, 0x6e, 0x5f, 0x50, 0x45, 0xf0, 0xc2, 0x47, 0xa5, 0xe1,
    0xc6, 0x64, 0x43, 0x2d, 0x6a, 0x0a, 0xf7, 0xe7, 0xda, 0x40, 0xb8, 0x3a,
    0xf0, 0x47, 0xdd, 0x01, 0xf5, 0xe0, 0xa9, 0x0e, 0x47, 0xc2, 0x24, 0xd7,
    0xb5, 0x13, 0x3a, 0x35, 0x4d, 0x11, 0xaa, 0x50, 0x03, 0xb3, 0xe8, 0x54,
    0x6c, 0x99, 0x01, 0x02, 0x21, 0x00, 0xcd, 0xb2, 0xd7, 0xa7, 0x43, 0x5b,
    0xcb, 0x45, 0xe5, 0x0e, 0x86, 0xf6, 0xc1, 0x4e, 0x97, 0xed, 0x78, 0x1f,
    0x09, 0x56, 0xcd, 0x26, 0xe6, 0xf7, 0x5e, 0xd9, 0xfc, 0x88, 0x12, 0x5f,
    0x84, 0x07, 0x02, 0x21, 0x00, 0xc9, 0xee, 0x30, 0xaf, 0x6c, 0xb9, 0x5a,
    0xc9, 0xc1, 0x14, 0x9e, 0xd8, 0x4b, 0x33, 0x38, 0x48, 0x17, 0x41, 0x35,
    0x94, 0x09, 0xf3, 0x69, 0xc4, 0x97, 0xbe, 0x17, 0x7d, 0x95, 0x0f, 0xb7,
    0xd1, 0x02, 0x21, 0x00, 0x8b, 0x0e, 0xf9, 0x8d, 0x61, 0x13, 0x20, 0x63,
    0x9b, 0x0b, 0x6c, 0x20, 0x4a, 0xe4, 0xa7, 0xfe, 0xe8, 0xf3, 0x0a, 0x6c,
    0x3c, 0xfa, 0xac, 0xaf, 0xd4, 0xd6, 0xc7, 0x4a, 0xf2, 0x28, 0xd2, 0x67,
    0x02, 0x20, 0x6b, 0x0e, 0x1d, 0xbf, 0x93, 0x5b, 0xbd, 0x77, 0x43, 0x27,
    0x24, 0x83, 0xb5, 0x72, 0xa5, 0x3f, 0x0b, 0x1d, 0x26, 0x43, 0xa2, 0xf6,
    0xea, 0xb7, 0x30, 0x5f, 0xb6, 0x62, 0x7c, 0xf9, 0x85, 0x51, 0x02, 0x20,
    0x3d, 0x22, 0x63, 0x15, 0x6b, 0x32, 0x41, 0x46, 0x44, 0x78, 0xb7, 0x13,
    0xeb, 0x85, 0x4c, 0x4f, 0x6b, 0x3e, 0xf0, 0x52, 0xf0, 0x46, 0x3b, 0x65,
    0xd8, 0x21, 0x7d, 0xae, 0xc0, 0x09, 0x98, 0x34};

const std::vector<uint8_t> kInvalidLengthKey = {
    0x30, 0x1b,        // SEQUENCE(len=27)
    0x02, 0x01, 0x00,  // INT(len=1) = 0
    0x30, 0x13,        // SEQUENCE(len=19)
    0x06, 0x07,        // OID(len=7)
    // dhPublicKey (1.2.840.10046.2.1)
    0x2a, 0x86, 0x48, 0xce, 0x3e, 0x02, 0x01, 0x06, 0x08,  // OID(len=8)
    // prime256v1 (1.2.840.10045.3.1.7) */
    0x2a, 0x86, 0x48, 0xce, 0x3d, 0x03, 0x01, 0x07, 0x04,
    0x00  // OCTET STRING(len=0)
};

// Ed25519 PKCS#8 PrivateKeyInfo with non-empty algorithm parameters.
// RFC 8410 section 3 says parameters MUST be absent for Ed25519, so NSS rejects
// this at the parameters-length check in PK11_ImportPrivateKeyInfoAndReturnKey.
//
//   SEQUENCE {
//     INTEGER 0
//     SEQUENCE {
//       OBJECT IDENTIFIER 1.3.101.112 (Ed25519)
//       NULL                            <-- params.len != 0, triggers rejection
//     }
//     OCTET STRING { OCTET STRING { 32 zero bytes } }
//   }
const std::vector<uint8_t> kEd25519WithNonEmptyAlgorithmParams = {
    0x30, 0x30,                    // SEQUENCE(48)
    0x02, 0x01, 0x00,              // INTEGER(0) -- version
    0x30, 0x07,                    // SEQUENCE(7) -- AlgorithmIdentifier
    0x06, 0x03, 0x2b, 0x65, 0x70,  // OID 1.3.101.112 (Ed25519)
    0x05, 0x00,  // NULL -- parameters must be absent per RFC 8410
    0x04, 0x22,  // OCTET STRING(34) -- privateKey
    0x04, 0x20,  // OCTET STRING(32) -- CurvePrivateKey
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
};

// kValidP256Key with the optional [1] publicKey field removed. Without it
// pk11pk12.c leaves CKA_NSS_DB unset, so softoken has to derive the point
// itself in sftk_mkPrivKey().
const std::vector<uint8_t> kValidP256KeyNoPublicKey = {
    0x30, 0x41,                                            // SEQUENCE(65)
    0x02, 0x01, 0x00,                                      // INTEGER 0
    0x30, 0x13,                                            // SEQUENCE(19)
    0x06, 0x07, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x02, 0x01,  // id-ecPublicKey
    0x06, 0x08, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x03, 0x01, 0x07,  // prime256v1
    0x04, 0x27,        // OCTET STRING(39) -- privateKey
    0x30, 0x25,        // SEQUENCE(37) -- ECPrivateKey
    0x02, 0x01, 0x01,  // INTEGER 1 -- version
    0x04, 0x20,        // OCTET STRING(32) -- privateKey
    0xc9, 0xaf, 0xa9, 0xd8, 0x45, 0xba, 0x75, 0x16, 0x6b, 0x5c, 0x21,
    0x57, 0x67, 0xb1, 0xd6, 0x93, 0x4e, 0x50, 0xc3, 0xdb, 0x36, 0xe8,
    0x9b, 0x12, 0x7b, 0x8a, 0x62, 0x2b, 0x12, 0x0f, 0x67, 0x21};

const std::vector<uint8_t> kValidP256PublicKey = {
    0x04, 0x60, 0xfe, 0xd4, 0xba, 0x25, 0x5a, 0x9d, 0x31, 0xc9, 0x61,
    0xeb, 0x74, 0xc6, 0x35, 0x6d, 0x68, 0xc0, 0x49, 0xb8, 0x92, 0x3b,
    0x61, 0xfa, 0x6c, 0xe6, 0x69, 0x62, 0x2e, 0x60, 0xf2, 0x9f, 0xb6,
    0x79, 0x03, 0xfe, 0x10, 0x08, 0xb8, 0xbc, 0x99, 0xa4, 0x1a, 0xe9,
    0xe9, 0x56, 0x28, 0xbc, 0x64, 0xf2, 0xf1, 0xb2, 0x0c, 0x2d, 0x7e,
    0x9f, 0x51, 0x77, 0xa3, 0xc2, 0x94, 0xd4, 0x46, 0x22, 0x99};

// P-256 with a 31-byte privateKey OCTET STRING rather than the 32 bytes RFC
// 5915 requires, and no [1] publicKey field. The missing publicKey sends
// sftk_mkPrivKey() down the derive path, where the short scalar makes
// EC_DerivePublicKey() fail on ec_secp256r1_pt_mul()'s k->len check.
const std::vector<uint8_t> kP256ShortScalarNoPublicKey = {
    0x30, 0x40,                                            // SEQUENCE(64)
    0x02, 0x01, 0x00,                                      // INTEGER 0
    0x30, 0x13,                                            // SEQUENCE(19)
    0x06, 0x07, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x02, 0x01,  // id-ecPublicKey
    0x06, 0x08, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x03, 0x01, 0x07,  // prime256v1
    0x04, 0x26,        // OCTET STRING(38) -- privateKey
    0x30, 0x24,        // SEQUENCE(36) -- ECPrivateKey
    0x02, 0x01, 0x01,  // INTEGER 1 -- version
    0x04, 0x1f,        // OCTET STRING(31) -- one byte short
    0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa,
    0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa,
    0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa, 0xaa};

const std::vector<uint8_t> kInvalidZeroLengthKey = {
    0x30, 0x1a,        // SEQUENCE(len=26)
    0x02, 0x01, 0x00,  // INT(len=1) = 0
    0x30, 0x13,        // SEQUENCE(len=19)
    0x06, 0x07,        // OID(len=7)
    // dhPublicKey (1.2.840.10046.2.1)
    0x2a, 0x86, 0x48, 0xce, 0x3e, 0x02, 0x01, 0x06, 0x08,  // OID(len=8)
    // prime256v1 (1.2.840.10045.3.1.7) */
    0x2a, 0x86, 0x48, 0xce, 0x3d, 0x03, 0x01, 0x07, 0x04,
    0x00  // OCTET STRING(len=0)
};

class DERPrivateKeyImportTest : public ::testing::Test {
 public:
  bool ParsePrivateKey(const std::vector<uint8_t>& data, bool expect_success) {
    SECKEYPrivateKey* key = nullptr;
    SECStatus rv = SECFailure;
    std::string nick_str =
        ::testing::UnitTest::GetInstance()->current_test_info()->name() +
        std::to_string(rand());
    SECItem item = {siBuffer, const_cast<unsigned char*>(data.data()),
                    static_cast<unsigned int>(data.size())};
    SECItem nick = {
        siBuffer,
        reinterpret_cast<unsigned char*>(const_cast<char*>(nick_str.data())),
        static_cast<unsigned int>(nick_str.length())};

    ScopedPK11SlotInfo slot(PK11_GetInternalKeySlot());
    EXPECT_TRUE(slot);
    if (!slot) {
      return false;
    }

    if (PK11_NeedUserInit(slot.get())) {
      if (PK11_InitPin(slot.get(), nullptr, nullptr) != SECSuccess) {
        EXPECT_EQ(rv, SECSuccess) << "PK11_InitPin failed";
      }
    }
    rv = PK11_Authenticate(slot.get(), PR_TRUE, nullptr);
    EXPECT_EQ(rv, SECSuccess);

    /* SEC_ASN1DecoderFinish() (lib/util/secasn1d.c:3001) only reports
     * SEC_ERROR_BAD_DER when no error is already pending. */
    PORT_SetError(0);
    rv = PK11_ImportDERPrivateKeyInfoAndReturnKey(
        slot.get(), &item, &nick, nullptr, true, false, KU_ALL, &key, nullptr);
    EXPECT_EQ(rv == SECSuccess, key != nullptr);

    if (expect_success) {
      // Try to find the key via its label
      ScopedSECKEYPrivateKeyList list(PK11_ListPrivKeysInSlot(
          slot.get(), const_cast<char*>(nick_str.c_str()), nullptr));
      EXPECT_FALSE(!list);
    }

    if (key) {
      rv = PK11_DeleteTokenPrivateKey(key, true);
      EXPECT_EQ(SECSuccess, rv);

      // PK11_DeleteTokenPrivateKey leaves an errorCode set when there's
      // no cert. This is expected, so clear it.
      if (PORT_GetError() == SSL_ERROR_NO_CERTIFICATE) {
        PORT_SetError(0);
      }
    }

    return rv == SECSuccess;
  }

  SECStatus BuildPrivateKeyInfoAndImportIt(SECOidTag algTag) {
    ScopedPK11SlotInfo slot(PK11_GetInternalSlot());
    EXPECT_TRUE(slot);
    if (!slot) {
      return SECFailure;
    }

    ScopedPLArenaPool arena(PORT_NewArena(DER_DEFAULT_CHUNKSIZE));
    EXPECT_TRUE(arena);
    if (!arena) {
      return SECFailure;
    }

    SECKEYPrivateKeyInfo* pki = (SECKEYPrivateKeyInfo*)PORT_ArenaZAlloc(
        arena.get(), sizeof(SECKEYPrivateKeyInfo));
    EXPECT_TRUE(pki);
    if (!pki) {
      return SECFailure;
    }

    pki->version.data = (unsigned char*)PORT_ArenaAlloc(arena.get(), 1);
    EXPECT_TRUE(pki->version.data);
    if (!pki->version.data) {
      return SECFailure;
    }

    pki->version.data[0] = 0x00;
    pki->version.len = 1;

    EXPECT_EQ(
        SECOID_SetAlgorithmID(arena.get(), &pki->algorithm, algTag, nullptr),
        SECSuccess);

    // Empty private key
    pki->privateKey.data = (unsigned char*)PORT_ArenaAlloc(arena.get(), 1);
    EXPECT_TRUE(pki->privateKey.data);
    if (!pki->privateKey.data) {
      return SECFailure;
    }
    pki->privateKey.len = 0;

    SECKEYPrivateKey* privk = nullptr;
    PORT_SetError(0);
    SECStatus rv = PK11_ImportPrivateKeyInfoAndReturnKey(
        slot.get(), pki, nullptr, nullptr, PR_FALSE, PR_FALSE, KU_ALL, &privk,
        nullptr);

    if (privk) {
      SECKEY_DestroyPrivateKey(privk);
    }
    return rv;
  }
};

TEST_F(DERPrivateKeyImportTest, ImportPrivateRSAKey) {
  EXPECT_TRUE(ParsePrivateKey(kValidRSAKey, true));
  EXPECT_FALSE(PORT_GetError()) << PORT_GetError();
}

TEST_F(DERPrivateKeyImportTest, ImportRSAPSSKey) {
  // An id-RSASSA-PSS (1.2.840.113549.1.1.10) PrivateKeyInfo wraps an ordinary
  // RSAPrivateKey, byte-identical to rsaEncryption; only the algorithm OID
  // differs (and NSS represents both as CKK_RSA). Reuse the valid RSA key and
  // swap only the OID's terminal arc (1 -> 10) to exercise the PSS path.
  std::vector<uint8_t> pss_key = kValidRSAKey;
  ASSERT_EQ(pss_key[19], 0x01);  // rsaEncryption OID terminal arc
  pss_key[19] = 0x0a;            // id-RSASSA-PSS
  EXPECT_TRUE(ParsePrivateKey(pss_key, true));
  EXPECT_FALSE(PORT_GetError()) << PORT_GetError();
}

TEST_F(DERPrivateKeyImportTest, ImportEcdsaKey) {
  EXPECT_TRUE(ParsePrivateKey(kValidP256Key, true));
  EXPECT_FALSE(PORT_GetError()) << PORT_GetError();
}

// The input has no public key, so CKA_NSS_DB is left unset,
// and CKA_TOKEN is false, so the point derivation happens during the import.
TEST_F(DERPrivateKeyImportTest, ImportEcdsaKeyWithoutPublicKey) {
  ScopedPK11SlotInfo slot(PK11_GetInternalSlot());
  ASSERT_TRUE(slot);

  SECItem item = {siBuffer,
                  const_cast<unsigned char*>(kValidP256KeyNoPublicKey.data()),
                  static_cast<unsigned int>(kValidP256KeyNoPublicKey.size())};
  SECKEYPrivateKey* privk = nullptr;
  PORT_SetError(0);
  ASSERT_EQ(SECSuccess, PK11_ImportDERPrivateKeyInfoAndReturnKey(
                            slot.get(), &item, nullptr, nullptr, PR_FALSE,
                            PR_FALSE, KU_ALL, &privk, nullptr));
  ScopedSECKEYPrivateKey key(privk);
  ASSERT_TRUE(key);

  SECItem point = {siBuffer, nullptr, 0};
  ASSERT_EQ(SECSuccess, PK11_ReadRawAttribute(PK11_TypePrivKey, key.get(),
                                              CKA_NSS_DB, &point));
  EXPECT_EQ(kValidP256PublicKey,
            std::vector<uint8_t>(point.data, point.data + point.len));
  SECITEM_FreeItem(&point, PR_FALSE);
}

// Same derive path, but the scalar is a byte short, so EC_DerivePublicKey()
// fails.
TEST_F(DERPrivateKeyImportTest, DeriveFailsOnShortScalar) {
  ScopedPK11SlotInfo slot(PK11_GetInternalSlot());
  ASSERT_TRUE(slot);

  SECItem item = {
      siBuffer, const_cast<unsigned char*>(kP256ShortScalarNoPublicKey.data()),
      static_cast<unsigned int>(kP256ShortScalarNoPublicKey.size())};
  SECKEYPrivateKey* privk = nullptr;
  PORT_SetError(0);
  EXPECT_EQ(SECFailure, PK11_ImportDERPrivateKeyInfoAndReturnKey(
                            slot.get(), &item, nullptr, nullptr, PR_FALSE,
                            PR_FALSE, KU_ALL, &privk, nullptr));
  EXPECT_EQ(nullptr, privk);
  EXPECT_EQ(PORT_GetError(), SEC_ERROR_BAD_DATA);
}

TEST_F(DERPrivateKeyImportTest, ImportInvalidPrivateKey) {
  EXPECT_FALSE(ParsePrivateKey(kInvalidLengthKey, false));
  EXPECT_EQ(PORT_GetError(), SEC_ERROR_BAD_DER) << PORT_GetError();
}

TEST_F(DERPrivateKeyImportTest, ImportZeroLengthPrivateKey) {
  EXPECT_FALSE(ParsePrivateKey(kInvalidZeroLengthKey, false));
  EXPECT_EQ(PORT_GetError(), SEC_ERROR_BAD_KEY) << PORT_GetError();
}

TEST_F(DERPrivateKeyImportTest, ImportZeroLengthMLDSAPrivateKey) {
  EXPECT_EQ(BuildPrivateKeyInfoAndImportIt(SEC_OID_ML_DSA_44), SECFailure);
  EXPECT_EQ(PORT_GetError(), SEC_ERROR_BAD_KEY);
}

TEST_F(DERPrivateKeyImportTest, ImportEd25519WithNonEmptyAlgorithmParams) {
  EXPECT_FALSE(ParsePrivateKey(kEd25519WithNonEmptyAlgorithmParams, false));
  EXPECT_EQ(PORT_GetError(), SEC_ERROR_UNSUPPORTED_KEYALG);
}

}  // namespace nss_test
