/* NSS libncrypto implementation.

Copyright 2010 Red Hat, Inc.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice, this
   list of conditions and the following disclaimer.

2. Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY CONTRIBUTORS ``AS IS'' AND ANY EXPRESS OR IMPLIED
WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO
EVENT SHALL CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER
IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
POSSIBILITY OF SUCH DAMAGE.

Red Hat author: Miloslav Trmač <mitr@redhat.com> */
#include <config.h>

#include <pthread.h>

#include <glib.h>
#include <keyhi.h>
#include <nss.h>
#include <pk11pub.h>

#include <ncrypto/ncrypto.h>

 /* Global state */

/* FIXME: Only init NSS/connect to the daemon when actually necessary. */

static pthread_mutex_t refcount_mutex = PTHREAD_MUTEX_INITIALIZER;
static size_t refcount;

static CK_RV
ncr_open_locked (void)
{
  if (refcount == 0)
    {
      if (NSS_NoDB_Init (NULL) != SECSuccess)
	return CKR_GENERAL_ERROR;
    }
  refcount++;
  return CKR_OK;
}

CK_RV
ncr_open (void)
{
  CK_RV res;

  if (pthread_mutex_lock (&refcount_mutex) != 0)
    return CKR_GENERAL_ERROR;
  res = ncr_open_locked ();
  if (pthread_mutex_unlock (&refcount_mutex) != 0)
    return CKR_GENERAL_ERROR;

  return res;
}

static CK_RV
ensure_ncr_is_open (void)
{
  CK_RV res;

  if (pthread_mutex_lock (&refcount_mutex) != 0)
    return CKR_GENERAL_ERROR;
  if (refcount != 0)
    res = CKR_OK;
  else
    res = ncr_open_locked ();
  if (pthread_mutex_unlock (&refcount_mutex) != 0)
    return CKR_GENERAL_ERROR;

  return res;
}

CK_RV
ncr_close (void)
{
  CK_RV res;

  if (pthread_mutex_lock (&refcount_mutex) != 0)
    return CKR_GENERAL_ERROR;

  if (refcount == 0)
    res = CKR_GENERAL_ERROR;
  else
    {
      refcount--;
      res = CKR_OK;
      if (refcount == 0)
	{
	  if (NSS_Shutdown () != SECSuccess)
	    res = CKR_GENERAL_ERROR;
	}
    }

  if (pthread_mutex_unlock (&refcount_mutex) != 0)
    return CKR_GENERAL_ERROR;

  return res;
}

 /* Asymmetric keys */

struct ncr_public_key
{
  SECKEYPublicKey *key;
};

struct ncr_private_key
{
  SECKEYPrivateKey *key;
};

static CK_RV
public_key_create (struct ncr_public_key **key, CK_KEY_TYPE type,
		   const SECItem *der_key)
{
  struct ncr_public_key *k;

  k = malloc (sizeof (*k));
  if (k == NULL)
    return CKR_HOST_MEMORY;

  k->key = SECKEY_ImportDERPublicKey((SECItem *)der_key, type);
  if (k->key == NULL)
    {
      free (k);
      return CKR_GENERAL_ERROR;
    }

  *key = k;
  return CKR_OK;
}

CK_RV
ncr_public_key_create (struct ncr_public_key **key, CK_KEY_TYPE type,
		       const void *der, size_t der_size)
{
  CK_RV res;
  SECItem der_key;

  res = ensure_ncr_is_open ();
  if (res != CKR_OK)
    return res;

  g_return_val_if_fail (key != NULL, CKR_ARGUMENTS_BAD);
  g_return_val_if_fail (der != NULL, CKR_ARGUMENTS_BAD);

  der_key.data = (void *)der;
  der_key.len = der_size;
  return public_key_create (key, type, &der_key);
}

CK_RV
ncr_public_key_destroy (struct ncr_public_key *key)
{
  CK_RV res;

  res = ensure_ncr_is_open ();
  if (res != CKR_OK)
    return res;

  g_return_val_if_fail (key != NULL, CKR_KEY_HANDLE_INVALID);

  /* Note: this leaks "key->key", which the function does not free (642767). */
  SECKEY_DestroyPublicKey (key->key);
  free (key);
  return CKR_OK;
}

static CK_RV
private_key_create (struct ncr_private_key **key, CK_KEY_TYPE type,
		    const SECItem *der_key, const SECItem *public_value)
{
  struct priv_key_info
  {
    SECItem version;
    SECAlgorithmID algorithm;
    SECItem private_key;
  };

  static const SEC_ASN1Template asn1_template[] =
    {
      { SEC_ASN1_SEQUENCE, 0, NULL, sizeof (struct priv_key_info) },
      { SEC_ASN1_INTEGER, offsetof (struct priv_key_info, version), NULL, 0, },
      { SEC_ASN1_INLINE, offsetof (struct priv_key_info, algorithm),
	SECOID_AlgorithmIDTemplate, 0, },
      {
	SEC_ASN1_OCTET_STRING, offsetof (struct priv_key_info, private_key),
	NULL, 0,
      },
      /* Attributes are optional and we never use them. */
      { 0, 0, NULL, 0 }
    };
  static const uint8_t wrap_key[32]; /* = { 0, }; */
  static const uint8_t wrap_iv[16]; /* = { 0, }; */
  static const uint8_t zero; /* = 0; */

  struct ncr_private_key *k;
  struct priv_key_info der_input;
  PK11SlotInfo *slot;
  SECItem der_info, *der_res, key_item, iv_item, wrapped_item;
  SECOidTag alg_tag;
  SECStatus ss;
  PK11SymKey *wrapping_key;
  PK11Context *ctx;
  int wrapped_len;

  switch (type)
    {
    case CKK_RSA:
      alg_tag = SEC_OID_PKCS1_RSA_ENCRYPTION;
      break;

    default:
      g_return_val_if_reached (CKR_ARGUMENTS_BAD);
    }

  k = malloc (sizeof (*k));
  if (k == NULL)
    return CKR_HOST_MEMORY;

  der_input.version.type = siUnsignedInteger;
  der_input.version.data = (void *)&zero;
  der_input.version.len = sizeof (zero);
  if (SECOID_SetAlgorithmID (NULL, &der_input.algorithm, alg_tag, NULL)
      != SECSuccess)
    {
      free (k);
      return CKR_GENERAL_ERROR;
    }
  der_input.private_key = *der_key;

  der_info.data = NULL;
  der_info.len = 0;
  der_res = SEC_ASN1EncodeItem (NULL, &der_info, &der_input, asn1_template);

  SECOID_DestroyAlgorithmID (&der_input.algorithm, PR_FALSE);

  if (der_res == NULL)
    {
      free (k);
      return CKR_HOST_MEMORY;
    }

  slot = PK11_GetBestSlot (CKM_AES_CBC_PAD, NULL);
  if (slot == NULL)
    goto err_der_info;

  key_item.data = (void *)wrap_key;
  key_item.len = sizeof (wrap_key);
  wrapping_key = PK11_ImportSymKeyWithFlags (slot, CKM_AES_CBC_PAD,
					     PK11_OriginUnwrap, CKA_FLAGS_ONLY,
					     &key_item,
					     CKF_ENCRYPT | CKF_UNWRAP,
					     PR_FALSE /* isPerm */, NULL);
  if (wrapping_key == NULL)
    goto err_slot;
  iv_item.data = (void *)wrap_iv;
  iv_item.len = sizeof (wrap_iv);

  ctx = PK11_CreateContextBySymKey (CKM_AES_CBC_PAD, CKA_ENCRYPT, wrapping_key,
				    &iv_item);
  if (ctx == NULL)
    goto err_wrapping_key;
  if (SECITEM_AllocItem (NULL, &wrapped_item, der_info.len + 16) == NULL)
    {
      PK11_DestroyContext (ctx, PR_TRUE);
      goto err_wrapping_key;
    }
  ss = PK11_CipherOp (ctx, wrapped_item.data, &wrapped_len, wrapped_item.len,
		      der_info.data, der_info.len);
  if (ss != SECSuccess)
    {
      PK11_DestroyContext (ctx, PR_TRUE);
      goto err_wrapped_item;
    }
  /* C_EncryptFinal is only available through this function.  "Nice.". */
  ss = PK11_DigestFinal (ctx, wrapped_item.data + wrapped_len,
			 &wrapped_item.len, wrapped_len);
  PK11_DestroyContext (ctx, PR_TRUE);
  if (ss != SECSuccess)
    goto err_wrapped_item;
  wrapped_item.len += wrapped_len;

  k->key = PK11_UnwrapPrivKey (slot, wrapping_key, CKM_AES_CBC_PAD, &iv_item,
			       &wrapped_item, NULL, (SECItem *)public_value,
			       PR_FALSE /* token */, PR_FALSE /* sensitive */,
			       type, NULL, 0, NULL);
  SECITEM_ZfreeItem (&wrapped_item, PR_FALSE);
  PK11_FreeSymKey (wrapping_key);
  PK11_FreeSlot (slot);
  PORT_Free (der_info.data);

  if (k->key == NULL)
    goto err_k;

  *key = k;
  return CKR_OK;

 err_wrapped_item:
  SECITEM_ZfreeItem (&wrapped_item, PR_FALSE);
 err_wrapping_key:
  PK11_FreeSymKey (wrapping_key);
 err_slot:
  PK11_FreeSlot (slot);
 err_der_info:
  PORT_Free (der_info.data);
 err_k:
  free (k);
  return CKR_GENERAL_ERROR;
}

/* FIXME: public_value should not be necessary, it is somewhere inside "der". */
CK_RV
ncr_private_key_create (struct ncr_private_key **key, CK_KEY_TYPE type,
			const void *der, size_t der_size,
			const void *public_value, size_t public_value_size)
{
  SECItem der_key;
  SECItem public;
  CK_RV res;

  res = ensure_ncr_is_open ();
  if (res != CKR_OK)
    return res;

  g_return_val_if_fail (key != NULL, CKR_ARGUMENTS_BAD);
  g_return_val_if_fail (der != NULL, CKR_ARGUMENTS_BAD);

  der_key.data = (void *)der;
  der_key.len = der_size;
  public.data = (void *)public_value;
  public.len = public_value_size;
  return private_key_create (key, type, &der_key, &public);
}

CK_RV
ncr_private_key_destroy (struct ncr_private_key *key)
{
  CK_RV res;

  res = ensure_ncr_is_open ();
  if (res != CKR_OK)
    return res;

  g_return_val_if_fail (key != NULL, CKR_KEY_HANDLE_INVALID);

  SECKEY_DestroyPrivateKey (key->key);
  free (key);
  return CKR_OK;
}

CK_RV
ncr_public_key_create_rsa (struct ncr_public_key **key, const void *modulus,
			   size_t modulus_size, const void *public_exponent,
			   size_t public_exponent_size)
{
  struct rsa_pub_key
  {
    SECItem modulus, public_exponent;
  };

  static const SEC_ASN1Template asn1_template[] =
    {
      { SEC_ASN1_SEQUENCE, 0, NULL, sizeof (struct rsa_pub_key) },
#define INT(X) { SEC_ASN1_INTEGER, offsetof (struct rsa_pub_key, X), NULL, 0, }
      INT (modulus), INT (public_exponent),
#undef INT
      { 0, 0, NULL, 0 }
    };

  struct rsa_pub_key der_input;
  SECItem der_key;
  CK_RV res;

  res = ensure_ncr_is_open ();
  if (res != CKR_OK)
    return res;

  g_return_val_if_fail (modulus != NULL, CKR_ARGUMENTS_BAD);
  g_return_val_if_fail (public_exponent != NULL, CKR_ARGUMENTS_BAD);

#define INT(X)					\
  der_input.X.type = siUnsignedInteger;		\
  der_input.X.data = (void *)X;			\
  der_input.X.len = X##_size;

  INT (modulus);
  INT (public_exponent);
#undef INT

  der_key.data = NULL;
  der_key.len = 0;
  if (SEC_ASN1EncodeItem(NULL, &der_key, &der_input, asn1_template) == NULL)
    return CKR_HOST_MEMORY;

  res = public_key_create (key, CKK_RSA, &der_key);
  PORT_Free (der_key.data);

  return res;
}

CK_RV
ncr_private_key_create_rsa (struct ncr_private_key **key, const void *modulus,
			    size_t modulus_size, const void *public_exponent,
			    size_t public_exponent_size,
			    const void *private_exponent,
			    size_t private_exponent_size, const void *prime_1,
			    size_t prime_1_size, const void *prime_2,
			    size_t prime_2_size, const void *exponent_1,
			    size_t exponent_1_size, const void *exponent_2,
			    size_t exponent_2_size, const void *coefficient,
			    size_t coefficient_size)
{
  struct rsa_pub_key
  {
    SECItem version;
    SECItem modulus, public_exponent, private_exponent, prime_1, prime_2;
    SECItem exponent_1, exponent_2, coefficient;
  };

  static const SEC_ASN1Template asn1_template[] =
    {
      { SEC_ASN1_SEQUENCE, 0, NULL, sizeof (struct rsa_pub_key) },
#define INT(X) { SEC_ASN1_INTEGER, offsetof (struct rsa_pub_key, X), NULL, 0, }
      INT (version),
      INT (modulus), INT (public_exponent), INT (private_exponent),
      INT (prime_1), INT (prime_2), INT (exponent_1), INT (exponent_2),
      INT (coefficient),
#undef INT
      { 0, 0, NULL, 0 }
    };
  static const uint8_t zero; /* = 0; */

  struct rsa_pub_key der_input;
  SECItem der_key;
  CK_RV res;

  res = ensure_ncr_is_open ();
  if (res != CKR_OK)
    return res;

  g_return_val_if_fail (modulus != NULL, CKR_ARGUMENTS_BAD);
  g_return_val_if_fail (public_exponent != NULL, CKR_ARGUMENTS_BAD);

  der_input.version.type = siUnsignedInteger;
  der_input.version.data = (void *)&zero;
  der_input.version.len = sizeof (zero);
#define INT(X)					\
  der_input.X.type = siUnsignedInteger;		\
  der_input.X.data = (void *)X;			\
  der_input.X.len = X##_size;

  INT (modulus);
  INT (public_exponent);
  INT (private_exponent);
  INT (prime_1);
  INT (prime_2);
  INT (exponent_1);
  INT (exponent_2);
  INT (coefficient);
#undef INT

  der_key.data = NULL;
  der_key.len = 0;
  if (SEC_ASN1EncodeItem (NULL, &der_key, &der_input, asn1_template) == NULL)
    return CKR_HOST_MEMORY;

  res = private_key_create (key, CKK_RSA, &der_key, &der_input.modulus);
  PORT_Free (der_key.data);

  return res;
}

 /* Asymmetric operations */

CK_RV
ncr_public_key_encrypt (CK_MECHANISM_TYPE mech, struct ncr_public_key *key,
			void *dest, size_t *dest_size_ptr, const void *src,
			size_t src_size)
{
  size_t dest_size;
  SECStatus ss;
  CK_RV res;

  res = ensure_ncr_is_open ();
  if (res != CKR_OK)
    return res;

  g_return_val_if_fail (key != NULL, CKR_KEY_HANDLE_INVALID);
  g_return_val_if_fail (dest_size_ptr != NULL, CKR_ARGUMENTS_BAD);

  /* This is correct for RSA, at least. */
  dest_size = SECKEY_SignatureLen (key->key);
  if (dest_size == 0)
    return CKR_GENERAL_ERROR;
  if (dest == NULL)
    {
      *dest_size_ptr = dest_size;
      return CKR_OK;
    }
  if (*dest_size_ptr < dest_size)
    {
      *dest_size_ptr = dest_size;
      return CKR_BUFFER_TOO_SMALL;
    }
  *dest_size_ptr = dest_size;

  g_return_val_if_fail (src != NULL || src_size == 0, CKR_ARGUMENTS_BAD);

  switch (mech)
    {
    case CKM_RSA_PKCS:
      ss = PK11_PubEncryptPKCS1 (key->key, dest, (void *)src, src_size, NULL);
      break;

    case CKM_RSA_X_509:
      ss = PK11_PubEncryptRaw (key->key, dest, (void *)src, src_size, NULL);
      break;

    default:
      g_return_val_if_reached (CKR_MECHANISM_INVALID);
    }
  if (ss != SECSuccess)
    return CKR_GENERAL_ERROR;

  return CKR_OK;
}

CK_RV
ncr_private_key_decrypt (CK_MECHANISM_TYPE mech, struct ncr_private_key *key,
			 void *dest, size_t *dest_size_ptr, const void *src,
			 size_t src_size)
{
  unsigned out_len;
  SECStatus ss;
  CK_RV res;

  res = ensure_ncr_is_open ();
  if (res != CKR_OK)
    return res;

  g_return_val_if_fail (key != NULL, CKR_KEY_HANDLE_INVALID);
  g_return_val_if_fail (dest_size_ptr != NULL, CKR_ARGUMENTS_BAD);

  g_return_val_if_fail (src != NULL || src_size == 0, CKR_ARGUMENTS_BAD);

  /* Don't change the value if the PrivDecrypt function doesn't touch it. */
  out_len = *dest_size_ptr;
  switch (mech)
    {
    case CKM_RSA_PKCS:
      ss = PK11_PrivDecryptPKCS1 (key->key, dest, &out_len, *dest_size_ptr,
				  (void *)src, src_size);
      break;

    case CKM_RSA_X_509:
      ss = PK11_PrivDecryptRaw (key->key, dest, &out_len, *dest_size_ptr,
				(void *)src, src_size);
      break;

    default:
      g_return_val_if_reached (CKR_MECHANISM_INVALID);
    }
  *dest_size_ptr = out_len;
  if (ss != SECSuccess)
    return CKR_GENERAL_ERROR;

  return CKR_OK;
}

CK_RV
ncr_private_key_sign (CK_MECHANISM_TYPE mech, struct ncr_private_key *key,
		      void *dest, size_t *dest_size_ptr, const void *src,
		      size_t src_size)
{
  size_t dest_size;
  SECItem src_item, dest_item;
  CK_RV res;

  res = ensure_ncr_is_open ();
  if (res != CKR_OK)
    return res;

  g_return_val_if_fail (key != NULL, CKR_KEY_HANDLE_INVALID);
  g_return_val_if_fail (dest_size_ptr != NULL, CKR_ARGUMENTS_BAD);

  g_return_val_if_fail (mech == PK11_MapSignKeyType (key->key->keyType),
			CKR_MECHANISM_INVALID);

  dest_size = PK11_SignatureLen (key->key);
  if (dest_size == 0)
    return CKR_GENERAL_ERROR;
  if (dest == NULL)
    {
      *dest_size_ptr = dest_size;
      return CKR_OK;
    }
  if (*dest_size_ptr < dest_size)
    {
      *dest_size_ptr = dest_size;
      return CKR_BUFFER_TOO_SMALL;
    }

  g_return_val_if_fail (src != NULL || src_size == 0, CKR_ARGUMENTS_BAD);

  src_item.data = (void *)src;
  src_item.len = src_size;
  dest_item.data = dest;
  dest_item.len = *dest_size_ptr;

  if (PK11_Sign (key->key, &dest_item, &src_item) != SECSuccess)
    return CKR_GENERAL_ERROR;

  *dest_size_ptr = dest_item.len;
  return CKR_OK;
}

CK_RV
ncr_public_key_verify (CK_MECHANISM_TYPE mech, struct ncr_public_key *key,
		       const void *signature, size_t signature_size,
		       const void *src, size_t src_size)
{
  SECItem sig_item, src_item;
  CK_RV res;

  res = ensure_ncr_is_open ();
  if (res != CKR_OK)
    return res;

  g_return_val_if_fail (key != NULL, CKR_KEY_HANDLE_INVALID);
  g_return_val_if_fail (signature != NULL || signature_size == 0,
			CKR_ARGUMENTS_BAD);
  g_return_val_if_fail (src != NULL || src_size == 0, CKR_ARGUMENTS_BAD);

  g_return_val_if_fail (mech == PK11_MapSignKeyType (key->key->keyType),
			CKR_MECHANISM_INVALID);

  sig_item.data = (void *)signature;
  sig_item.len = signature_size;
  src_item.data = (void *)src;
  src_item.len = src_size;

  return PK11_Verify (key->key, &sig_item, &src_item, NULL) == SECSuccess
    ? CKR_OK : CKR_GENERAL_ERROR;
}