// Copyright (c) 2013 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <openssl/bn.h>
#include <openssl/dsa.h>
#include <openssl/ecdsa.h>
#include <openssl/err.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/rsa.h>
#include <openssl/x509.h>

#include "base/android/build_info.h"
#include "base/android/jni_android.h"
#include "base/android/jni_array.h"
#include "base/android/scoped_java_ref.h"
#include "base/basictypes.h"
#include "base/bind.h"
#include "base/callback.h"
#include "base/compiler_specific.h"
#include "base/files/file_path.h"
#include "base/files/file_util.h"
#include "base/files/scoped_file.h"
#include "base/strings/string_number_conversions.h"
#include "base/strings/string_util.h"
#include "crypto/openssl_util.h"
#include "jni/AndroidKeyStoreTestUtil_jni.h"
#include "net/android/keystore.h"
#include "net/android/keystore_openssl.h"
#include "net/base/test_data_directory.h"
#include "net/ssl/scoped_openssl_types.h"
#include "testing/gtest/include/gtest/gtest.h"

// Technical note:
//
// This source file not only checks that signing with
// RawSignDigestWithPrivateKey() works correctly, it also verifies that
// the generated signature matches 100% of what OpenSSL generates when
// calling RSA_sign(NID_md5_sha1,...), DSA_sign(0, ...) or
// ECDSA_sign(0, ...).
//
// That's crucial to ensure that this function can later be used to
// implement client certificate support. More specifically, that it is
// possible to create a custom EVP_PKEY that uses
// RawSignDigestWithPrivateKey() internally to perform RSA/DSA/ECDSA
// signing, as invoked by the OpenSSL code at
// openssl/ssl/s3_clnt.c:ssl3_send_client_verify().
//
// For more details, read the comments in AndroidKeyStore.java.
//
// Finally, it also checks that using the EVP_PKEY generated with
// GetOpenSSLPrivateKeyWrapper() works correctly.

namespace net {
namespace android {

namespace {

typedef base::android::ScopedJavaLocalRef<jobject> ScopedJava;

JNIEnv* InitEnv() {
  JNIEnv* env = base::android::AttachCurrentThread();
  static bool inited = false;
  if (!inited) {
    RegisterNativesImpl(env);
    inited = true;
  }
  return env;
}

// Returns true if running on an Android version older than 4.2
bool IsOnAndroidOlderThan_4_2(void) {
  const int kAndroid42ApiLevel = 17;
  int level = base::android::BuildInfo::GetInstance()->sdk_int();
  return level < kAndroid42ApiLevel;
}

// Implements the callback expected by ERR_print_errors_cb().
// used by GetOpenSSLErrorString below.
int openssl_print_error_callback(const char* msg, size_t msglen, void* u) {
  std::string* result = reinterpret_cast<std::string*>(u);
  result->append(msg, msglen);
  return 1;
}

// Retrieves the OpenSSL error as a string
std::string GetOpenSSLErrorString(void) {
  std::string result;
  ERR_print_errors_cb(openssl_print_error_callback, &result);
  return result;
}

// Resize a string to |size| bytes of data, then return its data buffer
// address cast as an 'unsigned char*', as expected by OpenSSL functions.
// |str| the target string.
// |size| the number of bytes to write into the string.
// Return the string's new buffer in memory, as an 'unsigned char*'
// pointer.
unsigned char* OpenSSLWriteInto(std::string* str, size_t size) {
  return reinterpret_cast<unsigned char*>(WriteInto(str, size + 1));
}

// Load a given private key file into an EVP_PKEY.
// |filename| is the key file path.
// Returns a new EVP_PKEY on success, NULL on failure.
EVP_PKEY* ImportPrivateKeyFile(const char* filename) {
  // Load file in memory.
  base::FilePath certs_dir = GetTestCertsDirectory();
  base::FilePath file_path = certs_dir.AppendASCII(filename);
  base::ScopedFILE handle(base::OpenFile(file_path, "rb"));
  if (!handle.get()) {
    LOG(ERROR) << "Could not open private key file: " << filename;
    return NULL;
  }
  // Assume it is PEM_encoded. Load it as an EVP_PKEY.
  EVP_PKEY* pkey = PEM_read_PrivateKey(handle.get(), NULL, NULL, NULL);
  if (!pkey) {
    LOG(ERROR) << "Could not load public key file: " << filename
               << ", " << GetOpenSSLErrorString();
    return NULL;
  }
  return pkey;
}

// Convert a private key into its PKCS#8 encoded representation.
// |pkey| is the EVP_PKEY handle for the private key.
// |pkcs8| will receive the PKCS#8 bytes.
// Returns true on success, false otherwise.
bool GetPrivateKeyPkcs8Bytes(const crypto::ScopedEVP_PKEY& pkey,
                             std::string* pkcs8) {
  // Convert to PKCS#8 object.
  ScopedPKCS8_PRIV_KEY_INFO p8_info(EVP_PKEY2PKCS8(pkey.get()));
  if (!p8_info.get()) {
    LOG(ERROR) << "Can't get PKCS#8 private key from EVP_PKEY: "
               << GetOpenSSLErrorString();
    return false;
  }

  // Then convert it
  int len = i2d_PKCS8_PRIV_KEY_INFO(p8_info.get(), NULL);
  unsigned char* p = OpenSSLWriteInto(pkcs8, static_cast<size_t>(len));
  i2d_PKCS8_PRIV_KEY_INFO(p8_info.get(), &p);
  return true;
}

bool ImportPrivateKeyFileAsPkcs8(const char* filename,
                                 std::string* pkcs8) {
  crypto::ScopedEVP_PKEY pkey(ImportPrivateKeyFile(filename));
  if (!pkey.get())
    return false;
  return GetPrivateKeyPkcs8Bytes(pkey, pkcs8);
}

// Same as ImportPrivateKey, but for public ones.
EVP_PKEY* ImportPublicKeyFile(const char* filename) {
  // Load file as PEM data.
  base::FilePath certs_dir = GetTestCertsDirectory();
  base::FilePath file_path = certs_dir.AppendASCII(filename);
  base::ScopedFILE handle(base::OpenFile(file_path, "rb"));
  if (!handle.get()) {
    LOG(ERROR) << "Could not open public key file: " << filename;
    return NULL;
  }
  EVP_PKEY* pkey = PEM_read_PUBKEY(handle.get(), NULL, NULL, NULL);
  if (!pkey) {
    LOG(ERROR) << "Could not load public key file: " << filename
               << ", " << GetOpenSSLErrorString();
    return NULL;
  }
  return pkey;
}

// Retrieve a JNI local ref from encoded PKCS#8 data.
ScopedJava GetPKCS8PrivateKeyJava(PrivateKeyType key_type,
                                  const std::string& pkcs8_key) {
  JNIEnv* env = InitEnv();
  base::android::ScopedJavaLocalRef<jbyteArray> bytes(
      base::android::ToJavaByteArray(
          env,
          reinterpret_cast<const uint8*>(pkcs8_key.data()),
          pkcs8_key.size()));

  ScopedJava key(
      Java_AndroidKeyStoreTestUtil_createPrivateKeyFromPKCS8(
          env, key_type, bytes.obj()));

  return key;
}

const char kTestRsaKeyFile[] = "android-test-key-rsa.pem";

// The RSA test hash must be 36 bytes exactly.
const char kTestRsaHash[] = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ";

// Retrieve a JNI local ref for our test RSA key.
ScopedJava GetRSATestKeyJava() {
  std::string key;
  if (!ImportPrivateKeyFileAsPkcs8(kTestRsaKeyFile, &key))
    return ScopedJava();
  return GetPKCS8PrivateKeyJava(PRIVATE_KEY_TYPE_RSA, key);
}

const char kTestEcdsaKeyFile[] = "android-test-key-ecdsa.pem";
const char kTestEcdsaPublicKeyFile[] = "android-test-key-ecdsa-public.pem";

// The test hash for ECDSA keys must be 20 bytes exactly.
const char kTestEcdsaHash[] = "0123456789ABCDEFGHIJ";

// Retrieve a JNI local ref for our test ECDSA key.
ScopedJava GetECDSATestKeyJava() {
  std::string key;
  if (!ImportPrivateKeyFileAsPkcs8(kTestEcdsaKeyFile, &key))
    return ScopedJava();
  return GetPKCS8PrivateKeyJava(PRIVATE_KEY_TYPE_ECDSA, key);
}

// Call this function to verify that one message signed with our
// test ECDSA private key is correct. Since ECDSA signing introduces
// random elements in the signature, it is not possible to compare
// signature bits directly. However, one can use the public key
// to do the check.
bool VerifyTestECDSASignature(const base::StringPiece& message,
                              const base::StringPiece& signature) {
  crypto::ScopedEVP_PKEY pkey(ImportPublicKeyFile(kTestEcdsaPublicKeyFile));
  if (!pkey.get())
    return false;
  crypto::ScopedEC_KEY pub_key(EVP_PKEY_get1_EC_KEY(pkey.get()));
  if (!pub_key.get()) {
    LOG(ERROR) << "Could not get ECDSA public key: "
               << GetOpenSSLErrorString();
    return false;
  }

  const unsigned char* digest =
      reinterpret_cast<const unsigned char*>(message.data());
  int digest_len = static_cast<int>(message.size());
  const unsigned char* sigbuf =
      reinterpret_cast<const unsigned char*>(signature.data());
  int siglen = static_cast<int>(signature.size());

  int ret = ECDSA_verify(
      0, digest, digest_len, sigbuf, siglen, pub_key.get());
  if (ret != 1) {
    LOG(ERROR) << "ECDSA_verify() failed: " << GetOpenSSLErrorString();
    return false;
  }
  return true;
}

// Sign a message with OpenSSL, return the result as a string.
// |message| is the message to be signed.
// |openssl_key| is an OpenSSL EVP_PKEY to use.
// |result| receives the result.
// Returns true on success, false otherwise.
bool SignWithOpenSSL(const base::StringPiece& message,
                     EVP_PKEY* openssl_key,
                     std::string* result) {
  const unsigned char* digest =
      reinterpret_cast<const unsigned char*>(message.data());
  unsigned int digest_len = static_cast<unsigned int>(message.size());
  std::string signature;
  size_t signature_size;
  size_t max_signature_size;
  int key_type = EVP_PKEY_id(openssl_key);
  switch (key_type) {
    case EVP_PKEY_RSA:
    {
      crypto::ScopedRSA rsa(EVP_PKEY_get1_RSA(openssl_key));
      if (!rsa.get()) {
        LOG(ERROR) << "Could not get RSA from EVP_PKEY: "
                   << GetOpenSSLErrorString();
        return false;
      }
      // With RSA, the signature will always be RSA_size() bytes.
      max_signature_size = static_cast<size_t>(RSA_size(rsa.get()));
      unsigned char* p = OpenSSLWriteInto(&signature,
                                          max_signature_size);
      unsigned int p_len = 0;
      int ret = RSA_sign(
          NID_md5_sha1, digest, digest_len, p, &p_len, rsa.get());
      if (ret != 1) {
        LOG(ERROR) << "RSA_sign() failed: " << GetOpenSSLErrorString();
        return false;
      }
      signature_size = static_cast<size_t>(p_len);
      break;
    }
    case EVP_PKEY_EC:
    {
      crypto::ScopedEC_KEY ecdsa(EVP_PKEY_get1_EC_KEY(openssl_key));
      if (!ecdsa.get()) {
        LOG(ERROR) << "Could not get EC_KEY from EVP_PKEY: "
                   << GetOpenSSLErrorString();
        return false;
      }
      // Note, the actual signature can be smaller than ECDSA_size()
      max_signature_size = ECDSA_size(ecdsa.get());
      unsigned char* p = OpenSSLWriteInto(&signature,
                                          max_signature_size);
      unsigned int p_len = 0;
      // Note: first parameter is ignored by function.
      int ret = ECDSA_sign(
          0, digest, digest_len, p, &p_len, ecdsa.get());
      if (ret != 1) {
        LOG(ERROR) << "ECDSA_sign() fialed: " << GetOpenSSLErrorString();
        return false;
      }
      signature_size = static_cast<size_t>(p_len);
      break;
    }
    default:
      LOG(WARNING) << "Invalid OpenSSL key type: " << key_type;
      return false;
  }

  if (signature_size == 0) {
    LOG(ERROR) << "Signature is empty!";
    return false;
  }
  if (signature_size > max_signature_size) {
    LOG(ERROR) << "Signature size mismatch, actual " << signature_size
                << ", expected <= " << max_signature_size;
    return false;
  }
  signature.resize(signature_size);
  result->swap(signature);
  return true;
}

// Check that a generated signature for a given message matches
// OpenSSL output byte-by-byte.
// |message| is the input message.
// |signature| is the generated signature for the message.
// |openssl_key| is a raw EVP_PKEY for the same private key than the
// one which was used to generate the signature.
// Returns true on success, false otherwise.
bool CompareSignatureWithOpenSSL(const base::StringPiece& message,
                                 const base::StringPiece& signature,
                                 EVP_PKEY* openssl_key) {
  std::string openssl_signature;
  SignWithOpenSSL(message, openssl_key, &openssl_signature);

  if (signature.size() != openssl_signature.size()) {
    LOG(ERROR) << "Signature size mismatch, actual "
               << signature.size() << ", expected "
               << openssl_signature.size();
    return false;
  }
  for (size_t n = 0; n < signature.size(); ++n) {
    if (openssl_signature[n] != signature[n]) {
      LOG(ERROR) << "Signature byte mismatch at index " << n
                 << "actual " << signature[n] << ", expected "
                 << openssl_signature[n];
      LOG(ERROR) << "Actual signature  : "
                 << base::HexEncode(signature.data(), signature.size());
      LOG(ERROR) << "Expected signature: "
                 << base::HexEncode(openssl_signature.data(),
                                    openssl_signature.size());
      return false;
    }
  }
  return true;
}

// Sign a message with our platform API.
//
// |android_key| is a JNI reference to the platform PrivateKey object.
// |openssl_key| is a pointer to an OpenSSL key object for the exact
// same key content.
// |message| is a message.
// |result| will receive the result.
void DoKeySigning(jobject android_key,
                  EVP_PKEY* openssl_key,
                  const base::StringPiece& message,
                  std::string* result) {
  // First, get the platform signature.
  std::vector<uint8> android_signature;
  ASSERT_TRUE(
      RawSignDigestWithPrivateKey(android_key,
                                  message,
                                  &android_signature));

  result->assign(
      reinterpret_cast<const char*>(&android_signature[0]),
      android_signature.size());
}

// Sign a message with our OpenSSL EVP_PKEY wrapper around platform
// APIS.
//
// |android_key| is a JNI reference to the platform PrivateKey object.
// |openssl_key| is a pointer to an OpenSSL key object for the exact
// same key content.
// |message| is a message.
// |result| will receive the result.
void DoKeySigningWithWrapper(EVP_PKEY* wrapper_key,
                             EVP_PKEY* openssl_key,
                             const base::StringPiece& message,
                             std::string* result) {
  // First, get the platform signature.
  std::string wrapper_signature;
  SignWithOpenSSL(message, wrapper_key, &wrapper_signature);
  ASSERT_NE(0U, wrapper_signature.size());

  result->assign(
      reinterpret_cast<const char*>(&wrapper_signature[0]),
      wrapper_signature.size());
}

}  // namespace

TEST(AndroidKeyStore,GetRSAKeyModulus) {
  crypto::OpenSSLErrStackTracer err_trace(FROM_HERE);
  InitEnv();

  // Load the test RSA key.
  crypto::ScopedEVP_PKEY pkey(ImportPrivateKeyFile(kTestRsaKeyFile));
  ASSERT_TRUE(pkey.get());

  // Convert it to encoded PKCS#8 bytes.
  std::string pkcs8_data;
  ASSERT_TRUE(GetPrivateKeyPkcs8Bytes(pkey, &pkcs8_data));

  // Create platform PrivateKey object from it.
  ScopedJava key_java = GetPKCS8PrivateKeyJava(PRIVATE_KEY_TYPE_RSA,
                                                pkcs8_data);
  ASSERT_FALSE(key_java.is_null());

  // Retrieve the corresponding modulus through JNI
  std::vector<uint8> modulus_java;
  ASSERT_TRUE(GetRSAKeyModulus(key_java.obj(), &modulus_java));

  // Create an OpenSSL BIGNUM from it.
  crypto::ScopedBIGNUM bn(
      BN_bin2bn(reinterpret_cast<const unsigned char*>(&modulus_java[0]),
                static_cast<int>(modulus_java.size()),
                NULL));
  ASSERT_TRUE(bn.get());

  // Compare it to the one in the RSA key, they must be identical.
  crypto::ScopedRSA rsa(EVP_PKEY_get1_RSA(pkey.get()));
  ASSERT_TRUE(rsa.get()) << GetOpenSSLErrorString();

  ASSERT_EQ(0, BN_cmp(bn.get(), rsa.get()->n));
}

TEST(AndroidKeyStore,GetPrivateKeyTypeRSA) {
  crypto::OpenSSLErrStackTracer err_trace(FROM_HERE);

  ScopedJava rsa_key = GetRSATestKeyJava();
  ASSERT_FALSE(rsa_key.is_null());
  EXPECT_EQ(PRIVATE_KEY_TYPE_RSA,
            GetPrivateKeyType(rsa_key.obj()));
}

TEST(AndroidKeyStore,SignWithPrivateKeyRSA) {
  ScopedJava rsa_key = GetRSATestKeyJava();
  ASSERT_FALSE(rsa_key.is_null());

  if (IsOnAndroidOlderThan_4_2()) {
    LOG(INFO) << "This test can't run on Android < 4.2";
    return;
  }

  crypto::ScopedEVP_PKEY openssl_key(ImportPrivateKeyFile(kTestRsaKeyFile));
  ASSERT_TRUE(openssl_key.get());

  std::string message = kTestRsaHash;
  ASSERT_EQ(36U, message.size());

  std::string signature;
  DoKeySigning(rsa_key.obj(), openssl_key.get(), message, &signature);
  ASSERT_TRUE(
      CompareSignatureWithOpenSSL(message, signature, openssl_key.get()));
  // All good.
}

TEST(AndroidKeyStore,SignWithWrapperKeyRSA) {
  crypto::OpenSSLErrStackTracer err_tracer(FROM_HERE);

  ScopedJava rsa_key = GetRSATestKeyJava();
  ASSERT_FALSE(rsa_key.is_null());

  crypto::ScopedEVP_PKEY wrapper_key(
      GetOpenSSLPrivateKeyWrapper(rsa_key.obj()));
  ASSERT_TRUE(wrapper_key.get() != NULL);

  crypto::ScopedEVP_PKEY openssl_key(ImportPrivateKeyFile(kTestRsaKeyFile));
  ASSERT_TRUE(openssl_key.get());

  // Check that RSA_size() works properly on the wrapper key.
  EXPECT_EQ(EVP_PKEY_size(openssl_key.get()),
            EVP_PKEY_size(wrapper_key.get()));

  // Message size must be 36 for RSA_sign(NID_md5_sha1,...) to return
  // without an error.
  std::string message = kTestRsaHash;
  ASSERT_EQ(36U, message.size());

  std::string signature;
  DoKeySigningWithWrapper(wrapper_key.get(),
                          openssl_key.get(),
                          message,
                          &signature);
  ASSERT_TRUE(
      CompareSignatureWithOpenSSL(message, signature, openssl_key.get()));
}

TEST(AndroidKeyStore,GetPrivateKeyTypeECDSA) {
  crypto::OpenSSLErrStackTracer err_trace(FROM_HERE);

  ScopedJava ecdsa_key = GetECDSATestKeyJava();
  ASSERT_FALSE(ecdsa_key.is_null());
  EXPECT_EQ(PRIVATE_KEY_TYPE_ECDSA,
            GetPrivateKeyType(ecdsa_key.obj()));
}

TEST(AndroidKeyStore,SignWithPrivateKeyECDSA) {
  ScopedJava ecdsa_key = GetECDSATestKeyJava();
  ASSERT_FALSE(ecdsa_key.is_null());

  crypto::ScopedEVP_PKEY openssl_key(ImportPrivateKeyFile(kTestEcdsaKeyFile));
  ASSERT_TRUE(openssl_key.get());

  std::string message = kTestEcdsaHash;
  std::string signature;
  DoKeySigning(ecdsa_key.obj(), openssl_key.get(), message, &signature);
  ASSERT_TRUE(VerifyTestECDSASignature(message, signature));
}

TEST(AndroidKeyStore, SignWithWrapperKeyECDSA) {
  crypto::OpenSSLErrStackTracer err_tracer(FROM_HERE);

  ScopedJava ecdsa_key = GetECDSATestKeyJava();
  ASSERT_FALSE(ecdsa_key.is_null());

  crypto::ScopedEVP_PKEY wrapper_key(
      GetOpenSSLPrivateKeyWrapper(ecdsa_key.obj()));
  ASSERT_TRUE(wrapper_key.get());

  crypto::ScopedEVP_PKEY openssl_key(ImportPrivateKeyFile(kTestEcdsaKeyFile));
  ASSERT_TRUE(openssl_key.get());

  // Check that ECDSA size works correctly on the wrapper.
  EXPECT_EQ(EVP_PKEY_size(openssl_key.get()),
            EVP_PKEY_size(wrapper_key.get()));

  std::string message = kTestEcdsaHash;
  std::string signature;
  DoKeySigningWithWrapper(wrapper_key.get(),
                          openssl_key.get(),
                          message,
                          &signature);
  ASSERT_TRUE(VerifyTestECDSASignature(message, signature));
}

}  // namespace android
}  // namespace net