spark/src/lib/utils/crypto/PrivateKey.cpp

#ifndef _WIN32
#include "PrivateKey.h"

#include "CryptoInit.h"
#include "PublicKey.h"
#include "iostreambio.h"
#include "lib/logging/logging.h"
#include "lib/utils/bts_assert.h"
#include "lib/utils/helper_macros.h"
#include <cstring>
#include <tbb/spin_mutex.h>

namespace openssl {
#include <openssl/bio.h>
#include <openssl/bn.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/rand.h>
#include <openssl/rsa.h>
} // namespace openssl

namespace bts {
namespace crypto {

struct PrivateKey::Pimp {
   BTS_NO_COPY(Pimp);

   const PasswordCallback *password;
   PasswordCallback default_password;
   PublicKeyPtr public_key;

   Pimp(const PasswordCallback *pass)
       : password(pass), default_password(), public_key() {
      initialize();
      if (!password)
         password = &default_password;
   }

   ~Pimp() {}
   static void initialize() { CryptoInit::initialize(); }
};

int password_callback(char *buf, int size, int rwflag, void *u) {
   UNUSED(rwflag);
   std::string data;
   if (u) {
      PasswordCallback *pwd = reinterpret_cast<PasswordCallback *>(u);
      pwd->request_password(data, data);
   } else {
      PasswordCallback ddefault;
      ddefault.request_password(data, data);
   }
   int sz = data.size();
   if ((int)data.size() > size)
      sz = size;

   std::memcpy(buf, data.c_str(), sz);
   return sz;
}

const PublicKeyPtr PrivateKey::public_key() const {
   if (!me->public_key)
      me->public_key = PublicKey::create(*this);
   return me->public_key;
}

const PrivateKeyPtr
PrivateKey::create(int bits, AlgoType::type at, const PasswordCallback *pass) {
   openssl::EVP_PKEY *key = openssl::EVP_PKEY_new();
   if (at == AlgoType::RSA) {
      openssl::RSA *rsa = openssl::RSA_new();
      openssl::BIGNUM *bn = openssl::BN_new();
      openssl::BN_set_word(bn, RSA_F4);
      int keyok = openssl::RSA_generate_key_ex(rsa, bits, bn, 0);
      openssl::BN_free(bn);
      bts_assert_msg(keyok != -1, "Unable to create key");
      int i = openssl::EVP_PKEY_assign_RSA(key, rsa);
      bts_assert_msg(1 == i, "Unable to create key");
   } else
      bts_assert_msg(false, "only support RSA");
   return PrivateKeyPtr(new PrivateKey(key, pass));
}

const PrivateKeyPtr PrivateKey::read_file(
    const std::string &filename, const PasswordCallback *pass) {
   PrivateKey::Pimp::initialize();
   PrivateKeyPtr ret;
   openssl::BIO *bio = openssl::BIO_new_file(filename.c_str(), "r");
   bts_assert_msg(bio, "can't load keyfile");
   openssl::EVP_PKEY *key = openssl::PEM_read_bio_PrivateKey(
       bio, 0, password_callback, (void *)pass);
   if (!key)
      LOG(Gen, Error, "Invalid private key passphrase");
   else
      ret.reset(new PrivateKey(key, pass));
   return ret;
}

const PrivateKeyPtr PrivateKey::read_stream(
    const std::istream &stream, const PasswordCallback *pass) {
   PrivateKey::Pimp::initialize();
   PrivateKeyPtr ret;
   const_cast<std::istream &>(stream).seekg(0, std::ios::end);
   int length = const_cast<std::istream &>(stream).tellg();
   const_cast<std::istream &>(stream).seekg(0, std::ios::beg);

   char *data = new char[length];
   const_cast<std::istream &>(stream).read(data, length);
   openssl::BIO *bio = openssl::BIO_new(openssl::BIO_s_mem());
   BIO_write(bio, data, length);
   delete[] data;

   openssl::EVP_PKEY *key = openssl::PEM_read_bio_PrivateKey(
       bio, 0, password_callback, (void *)pass);
   if (!key)
      LOG(Gen, Error, "Invalid private key passphrase");
   else
      ret.reset(new PrivateKey(key, pass));
   return ret;
}

PrivateKey::PrivateKey(openssl::EVP_PKEY *key, const PasswordCallback *pass)
    : me(new Pimp(pass)), _key(key) {}

PrivateKey::~PrivateKey() {
   if (_key)
      EVP_PKEY_free(_key);
}

typedef int encrypt_decrypt_init(openssl::EVP_PKEY_CTX *ctx);
typedef int encrypt_decrypt(
    openssl::EVP_PKEY_CTX *ctx,
    unsigned char *out,
    size_t *outlen,
    const unsigned char *in,
    size_t inlen);

bool encrypt_decrypt_function(
    const std::iostream &in_dat,
    std::iostream &out_data,
    encrypt_decrypt_init *init_function,
    encrypt_decrypt *function,
    openssl::EVP_PKEY *key,
    bool encrypt) {
   UNUSED(encrypt);
   openssl::EVP_PKEY_CTX *ctx = openssl::EVP_PKEY_CTX_new(key, 0);
   bts_assert(ctx);
   if (init_function)
      bts_assert(init_function(ctx));
   bts_assert(openssl::EVP_PKEY_CTX_set_rsa_padding(ctx, RSA_PKCS1_PADDING));

   char *name, *header;
   unsigned char *in, *out;
   long inlen;
   size_t outlen;
   int i = 0;
   iostreambio biostream(const_cast<std::iostream *>(&in_dat));
   if (!encrypt) {
      auto bio = biostream.getBIO();
      i = openssl::PEM_read_bio(bio, &name, &header, &in, &inlen);
   } else {
      std::iostream seeker(in_dat.rdbuf());
      seeker.seekg(0, std::ios::end);
      inlen = seeker.tellg();
      seeker.seekg(0, std::ios::beg);

      in = new unsigned char[inlen];
      i = openssl::BIO_read(biostream.getBIO(), in, inlen);
   }
   // BIO_flush(biostream.getBIO());
   BIO_free(biostream.getBIO());

   bts_assert_msg(i > 0, "file not read, too short?");
   bts_assert(function(ctx, NULL, &outlen, in, inlen)); // to get size
   out = new unsigned char[outlen];
   bts_assert(function(ctx, out, &outlen, in, inlen));
   // if(encrypt)
   //   delete[] in;

   std::string header_name = "BTS", header_header = "";
   iostreambio biostream_out(const_cast<std::iostream *>(&out_data));
   if (encrypt)
      i = openssl::PEM_write_bio(
          biostream_out.getBIO(), header_name.c_str(),
          const_cast<char *>(header_header.c_str()), out, outlen);
   else
      i = openssl::BIO_write(biostream_out.getBIO(), out, outlen);
   // BIO_flush(biostream_out.getBIO());
   BIO_free(biostream_out.getBIO());
   // delete[] out;

   return (i > 0);
}

int private_encrypt_func(
    openssl::EVP_PKEY_CTX *ctx,
    unsigned char *out,
    size_t *outlen,
    const unsigned char *in,
    size_t inlen) {
   openssl::EVP_PKEY *pk = openssl::EVP_PKEY_CTX_get0_pkey(ctx);
   auto rsa = openssl::EVP_PKEY_get0_RSA(pk);
   if (out == 0) {
      *outlen = openssl::RSA_size(rsa);
      return -1;
   } else {
      return RSA_private_encrypt(
          inlen, const_cast<unsigned char *>(in), out, rsa, RSA_PKCS1_PADDING);
   }
}

int public_decrypt_func(
    openssl::EVP_PKEY_CTX *ctx,
    unsigned char *out,
    size_t *outlen,
    const unsigned char *in,
    size_t inlen) {
   openssl::EVP_PKEY *pk = openssl::EVP_PKEY_CTX_get0_pkey(ctx);
   auto rsa = openssl::EVP_PKEY_get0_RSA(pk);
   if (out == 0) {
      *outlen = openssl::RSA_size(rsa);
      return -1;
   } else {
      int i = RSA_public_decrypt(
          inlen, const_cast<unsigned char *>(in), out, rsa, RSA_PKCS1_PADDING);
      out[i] = '\0';
      *outlen = i;
      return i;
   }
}

bool PrivateKey::private_encrypt(
    const std::iostream &in_data, std::iostream &out_data) const {
   return encrypt_decrypt_function(
       in_data, out_data, 0, private_encrypt_func, _key, true);
}

bool PrivateKey::private_decrypt(
    const std::iostream &in_data, std::iostream &out_data) const {
   return encrypt_decrypt_function(
       in_data, out_data, openssl::EVP_PKEY_decrypt_init,
       openssl::EVP_PKEY_decrypt, _key, false);
}

bool PrivateKey::public_encrypt(
    const std::iostream &in_data, std::iostream &out_data) const {
   return encrypt_decrypt_function(
       in_data, out_data, openssl::EVP_PKEY_encrypt_init,
       openssl::EVP_PKEY_encrypt, _key, true);
}

bool PrivateKey::public_decrypt(
    const std::iostream &in_data, std::iostream &out_data) const {
   return encrypt_decrypt_function(
       in_data, out_data, 0, public_decrypt_func, _key, false);
}

bool PrivateKey::encrypt(
    const std::iostream &in_data, std::iostream &out_data) const {
   return private_encrypt(in_data, out_data);
}

bool PrivateKey::decrypt(
    const std::iostream &in_data, std::iostream &out_data) const {
   return private_decrypt(in_data, out_data);
}

void PrivateKey::write(openssl::BIO *bio, CipherType::type cipher_type) const {
   // see rsa.c ln 213
   const openssl::EVP_CIPHER *ciph =
       openssl::EVP_get_cipherbyname(CipherType::name[(int)cipher_type]);
   bts_assert_msg(
       cipher_type == CipherType::None || ciph, "unknown cipher type");
   int nbytes = openssl::PEM_write_bio_PKCS8PrivateKey(
       bio, _key, ciph, 0, 0, password_callback, (void *)me->password);
   UNUSED(nbytes);
   int err = BIO_flush(bio);
   UNUSED(err);
}

void PrivateKey::write_file(
    const std::string &filename, CipherType::type cipher_type) const {
   openssl::BIO *bio = openssl::BIO_new_file(filename.c_str(), "w+");
   write(bio, cipher_type);
   BIO_free(bio);
}

void PrivateKey::write_stream(
    std::iostream &stream, CipherType::type cipher_type) const {
   iostreambio biostream_out(&stream);
   write(biostream_out.getBIO(), cipher_type);
   BIO_free(biostream_out.getBIO());
}

} // namespace crypto
} // namespace bts
#endif