/*	$NetBSD: ssh-pkcs11.c,v 1.1 2010/11/21 17:06:01 adam Exp $	*/
/* $OpenBSD: ssh-pkcs11.c,v 1.6 2010/06/08 21:32:19 markus Exp $ */
/*
 * Copyright (c) 2010 Markus Friedl.  All rights reserved.
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <sys/types.h>
#include <sys/queue.h>
#include <stdarg.h>
#include <stdio.h>

#include <string.h>
#include <dlfcn.h>

#define CRYPTOKI_COMPAT
#include "pkcs11.h"

#include "log.h"
#include "misc.h"
#include "key.h"
#include "ssh-pkcs11.h"
#include "xmalloc.h"

struct pkcs11_slotinfo {
	CK_TOKEN_INFO		token;
	CK_SESSION_HANDLE	session;
	int			logged_in;
};

struct pkcs11_provider {
	char			*name;
	void			*handle;
	CK_FUNCTION_LIST	*function_list;
	CK_INFO			info;
	CK_ULONG		nslots;
	CK_SLOT_ID		*slotlist;
	struct pkcs11_slotinfo	*slotinfo;
	int			valid;
	int			refcount;
	TAILQ_ENTRY(pkcs11_provider) next;
};

TAILQ_HEAD(, pkcs11_provider) pkcs11_providers;

struct pkcs11_key {
	struct pkcs11_provider	*provider;
	CK_ULONG		slotidx;
	int			(*orig_finish)(RSA *rsa);
	RSA_METHOD		rsa_method;
	char			*keyid;
	int			keyid_len;
};

int pkcs11_interactive = 0;

int
pkcs11_init(int interactive)
{
	pkcs11_interactive = interactive;
	TAILQ_INIT(&pkcs11_providers);
	return (0);
}

/*
 * finalize a provider shared libarary, it's no longer usable.
 * however, there might still be keys referencing this provider,
 * so the actuall freeing of memory is handled by pkcs11_provider_unref().
 * this is called when a provider gets unregistered.
 */
static void
pkcs11_provider_finalize(struct pkcs11_provider *p)
{
	CK_RV rv;
	CK_ULONG i;

	debug("pkcs11_provider_finalize: %p refcount %d valid %d",
	    p, p->refcount, p->valid);
	if (!p->valid)
		return;
	for (i = 0; i < p->nslots; i++) {
		if (p->slotinfo[i].session &&
		    (rv = p->function_list->C_CloseSession(
		    p->slotinfo[i].session)) != CKR_OK)
			error("C_CloseSession failed: %lu", rv);
	}
	if ((rv = p->function_list->C_Finalize(NULL)) != CKR_OK)
		error("C_Finalize failed: %lu", rv);
	p->valid = 0;
	p->function_list = NULL;
#ifdef HAVE_DLOPEN
	dlclose(p->handle);
#endif
}

/*
 * remove a reference to the provider.
 * called when a key gets destroyed or when the provider is unregistered.
 */
static void
pkcs11_provider_unref(struct pkcs11_provider *p)
{
	debug("pkcs11_provider_unref: %p refcount %d", p, p->refcount);
	if (--p->refcount <= 0) {
		if (p->valid)
			error("pkcs11_provider_unref: %p still valid", p);
		xfree(p->slotlist);
		xfree(p->slotinfo);
		xfree(p);
	}
}

/* unregister all providers, keys might still point to the providers */
void
pkcs11_terminate(void)
{
	struct pkcs11_provider *p;

	while ((p = TAILQ_FIRST(&pkcs11_providers)) != NULL) {
		TAILQ_REMOVE(&pkcs11_providers, p, next);
		pkcs11_provider_finalize(p);
		pkcs11_provider_unref(p);
	}
}

/* lookup provider by name */
static struct pkcs11_provider *
pkcs11_provider_lookup(char *provider_id)
{
	struct pkcs11_provider *p;

	TAILQ_FOREACH(p, &pkcs11_providers, next) {
		debug("check %p %s", p, p->name);
		if (!strcmp(provider_id, p->name))
			return (p);
	}
	return (NULL);
}

/* unregister provider by name */
int
pkcs11_del_provider(char *provider_id)
{
	struct pkcs11_provider *p;

	if ((p = pkcs11_provider_lookup(provider_id)) != NULL) {
		TAILQ_REMOVE(&pkcs11_providers, p, next);
		pkcs11_provider_finalize(p);
		pkcs11_provider_unref(p);
		return (0);
	}
	return (-1);
}

/* openssl callback for freeing an RSA key */
static int
pkcs11_rsa_finish(RSA *rsa)
{
	struct pkcs11_key	*k11;
	int rv = -1;

	if ((k11 = RSA_get_app_data(rsa)) != NULL) {
		if (k11->orig_finish)
			rv = k11->orig_finish(rsa);
		if (k11->provider)
			pkcs11_provider_unref(k11->provider);
		if (k11->keyid)
			xfree(k11->keyid);
		xfree(k11);
	}
	return (rv);
}

/* find a single 'obj' for given attributes */
static int
pkcs11_find(struct pkcs11_provider *p, CK_ULONG slotidx, CK_ATTRIBUTE *attr,
    CK_ULONG nattr, CK_OBJECT_HANDLE *obj)
{
	CK_FUNCTION_LIST	*f;
	CK_SESSION_HANDLE	session;
	CK_ULONG		nfound = 0;
	CK_RV			rv;
	int			ret = -1;

	f = p->function_list;
	session = p->slotinfo[slotidx].session;
	if ((rv = f->C_FindObjectsInit(session, attr, nattr)) != CKR_OK) {
		error("C_FindObjectsInit failed (nattr %lu): %lu", nattr, rv);
		return (-1);
	}
	if ((rv = f->C_FindObjects(session, obj, 1, &nfound)) != CKR_OK ||
	    nfound != 1) {
		debug("C_FindObjects failed (nfound %lu nattr %lu): %lu",
		    nfound, nattr, rv);
	} else
		ret = 0;
	if ((rv = f->C_FindObjectsFinal(session)) != CKR_OK)
		error("C_FindObjectsFinal failed: %lu", rv);
	return (ret);
}

/* openssl callback doing the actual signing operation */
static int
pkcs11_rsa_private_encrypt(int flen, const u_char *from, u_char *to, RSA *rsa,
    int padding)
{
	struct pkcs11_key	*k11;
	struct pkcs11_slotinfo	*si;
	CK_FUNCTION_LIST	*f;
	CK_OBJECT_HANDLE	obj;
	CK_ULONG		tlen = 0;
	CK_RV			rv;
	CK_OBJECT_CLASS		private_key_class = CKO_PRIVATE_KEY;
	CK_BBOOL		true = CK_TRUE;
	CK_MECHANISM		mech = {
		CKM_RSA_PKCS, NULL_PTR, 0
	};
	CK_ATTRIBUTE		key_filter[] = {
		{CKA_CLASS, &private_key_class, sizeof(private_key_class) },
		{CKA_ID, NULL, 0},
		{CKA_SIGN, &true, sizeof(true) }
	};
	char			*pin, prompt[1024];
	int			rval = -1;

	if ((k11 = RSA_get_app_data(rsa)) == NULL) {
		error("RSA_get_app_data failed for rsa %p", rsa);
		return (-1);
	}
	if (!k11->provider || !k11->provider->valid) {
		error("no pkcs11 (valid) provider for rsa %p", rsa);
		return (-1);
	}
	f = k11->provider->function_list;
	si = &k11->provider->slotinfo[k11->slotidx];
	if ((si->token.flags & CKF_LOGIN_REQUIRED) && !si->logged_in) {
		if (!pkcs11_interactive) {
			error("need pin");
			return (-1);
		}
		snprintf(prompt, sizeof(prompt), "Enter PIN for '%s': ",
		    si->token.label);
		pin = read_passphrase(prompt, RP_ALLOW_EOF);
		if (pin == NULL)
			return (-1);	/* bail out */
		if ((rv = f->C_Login(si->session, CKU_USER, pin, strlen(pin)))
		    != CKR_OK) {
			xfree(pin);
			error("C_Login failed: %lu", rv);
			return (-1);
		}
		xfree(pin);
		si->logged_in = 1;
	}
	key_filter[1].pValue = k11->keyid;
	key_filter[1].ulValueLen = k11->keyid_len;
	/* try to find object w/CKA_SIGN first, retry w/o */
	if (pkcs11_find(k11->provider, k11->slotidx, key_filter, 3, &obj) < 0 &&
	    pkcs11_find(k11->provider, k11->slotidx, key_filter, 2, &obj) < 0) {
		error("cannot find private key");
	} else if ((rv = f->C_SignInit(si->session, &mech, obj)) != CKR_OK) {
		error("C_SignInit failed: %lu", rv);
	} else {
		/* XXX handle CKR_BUFFER_TOO_SMALL */
		tlen = RSA_size(rsa);
		rv = f->C_Sign(si->session, (CK_BYTE *)from, flen, to, &tlen);
		if (rv == CKR_OK) 
			rval = tlen;
		else 
			error("C_Sign failed: %lu", rv);
	}
	return (rval);
}

static int
pkcs11_rsa_private_decrypt(int flen, const u_char *from, u_char *to, RSA *rsa,
    int padding)
{
	return (-1);
}

/* redirect private key operations for rsa key to pkcs11 token */
static int
pkcs11_rsa_wrap(struct pkcs11_provider *provider, CK_ULONG slotidx,
    CK_ATTRIBUTE *keyid_attrib, RSA *rsa)
{
	struct pkcs11_key	*k11;
	const RSA_METHOD	*def = RSA_get_default_method();

	k11 = xcalloc(1, sizeof(*k11));
	k11->provider = provider;
	provider->refcount++;	/* provider referenced by RSA key */
	k11->slotidx = slotidx;
	/* identify key object on smartcard */
	k11->keyid_len = keyid_attrib->ulValueLen;
	k11->keyid = xmalloc(k11->keyid_len);
	memcpy(k11->keyid, keyid_attrib->pValue, k11->keyid_len);
	k11->orig_finish = def->finish;
	memcpy(&k11->rsa_method, def, sizeof(k11->rsa_method));
	k11->rsa_method.name = "pkcs11";
	k11->rsa_method.rsa_priv_enc = pkcs11_rsa_private_encrypt;
	k11->rsa_method.rsa_priv_dec = pkcs11_rsa_private_decrypt;
	k11->rsa_method.finish = pkcs11_rsa_finish;
	RSA_set_method(rsa, &k11->rsa_method);
	RSA_set_app_data(rsa, k11);
	return (0);
}

/* remove trailing spaces */
static void
rmspace(char *buf, size_t len)
{
	size_t i;

	if (!len)
		return;
	for (i = len - 1;  i > 0; i--)
		if (i == len - 1 || buf[i] == ' ')
			buf[i] = '\0';
		else
			break;
}

/*
 * open a pkcs11 session and login if required.
 * if pin == NULL we delay login until key use
 */
static int
pkcs11_open_session(struct pkcs11_provider *p, CK_ULONG slotidx, char *pin)
{
	CK_RV			rv;
	CK_FUNCTION_LIST	*f;
	CK_SESSION_HANDLE	session;
	int			login_required;

	f = p->function_list;
	login_required = p->slotinfo[slotidx].token.flags & CKF_LOGIN_REQUIRED;
	if (pin && login_required && !strlen(pin)) {
		error("pin required");
		return (-1);
	}
	if ((rv = f->C_OpenSession(p->slotlist[slotidx], CKF_RW_SESSION|
	    CKF_SERIAL_SESSION, NULL, NULL, &session))
	    != CKR_OK) {
		error("C_OpenSession failed: %lu", rv);
		return (-1);
	}
	if (login_required && pin) {
		if ((rv = f->C_Login(session, CKU_USER, pin, strlen(pin)))
		    != CKR_OK) {
			error("C_Login failed: %lu", rv);
			if ((rv = f->C_CloseSession(session)) != CKR_OK)
				error("C_CloseSession failed: %lu", rv);
			return (-1);
		}
		p->slotinfo[slotidx].logged_in = 1;
	}
	p->slotinfo[slotidx].session = session;
	return (0);
}

/*
 * lookup public keys for token in slot identified by slotidx,
 * add 'wrapped' public keys to the 'keysp' array and increment nkeys.
 * keysp points to an (possibly empty) array with *nkeys keys.
 */
static int
pkcs11_fetch_keys(struct pkcs11_provider *p, CK_ULONG slotidx, Key ***keysp,
    int *nkeys)
{
	Key			*key;
	RSA			*rsa;
	int			i;
	CK_RV			rv;
	CK_OBJECT_HANDLE	obj;
	CK_ULONG		nfound;
	CK_SESSION_HANDLE	session;
	CK_FUNCTION_LIST	*f;
	CK_OBJECT_CLASS		pubkey_class = CKO_PUBLIC_KEY;
	CK_ATTRIBUTE		pubkey_filter[] = {
		{ CKA_CLASS, &pubkey_class, sizeof(pubkey_class) }
	};
	CK_ATTRIBUTE		attribs[] = {
		{ CKA_ID, NULL, 0 },
		{ CKA_MODULUS, NULL, 0 },
		{ CKA_PUBLIC_EXPONENT, NULL, 0 }
	};

	f = p->function_list;
	session = p->slotinfo[slotidx].session;
	/* setup a filter the looks for public keys */
	if ((rv = f->C_FindObjectsInit(session, pubkey_filter, 1)) != CKR_OK) {
		error("C_FindObjectsInit failed: %lu", rv);
		return (-1);
	}
	while (1) {
		/* XXX 3 attributes in attribs[] */
		for (i = 0; i < 3; i++) {
			attribs[i].pValue = NULL;
			attribs[i].ulValueLen = 0;
		}
		if ((rv = f->C_FindObjects(session, &obj, 1, &nfound)) != CKR_OK
		    || nfound == 0)
			break;
		/* found a key, so figure out size of the attributes */
		if ((rv = f->C_GetAttributeValue(session, obj, attribs, 3))
		    != CKR_OK) {
			error("C_GetAttributeValue failed: %lu", rv);
			continue;
		}
		/* check that none of the attributes are zero length */
		if (attribs[0].ulValueLen == 0 ||
		    attribs[1].ulValueLen == 0 ||
		    attribs[2].ulValueLen == 0) {
			continue;
		}
		/* allocate buffers for attributes */
		for (i = 0; i < 3; i++)
			attribs[i].pValue = xmalloc(attribs[i].ulValueLen);
		/* retrieve ID, modulus and public exponent of RSA key */
		if ((rv = f->C_GetAttributeValue(session, obj, attribs, 3))
		    != CKR_OK) {
			error("C_GetAttributeValue failed: %lu", rv);
		} else if ((rsa = RSA_new()) == NULL) {
			error("RSA_new failed");
		} else {
			rsa->n = BN_bin2bn(attribs[1].pValue,
			    attribs[1].ulValueLen, NULL);
			rsa->e = BN_bin2bn(attribs[2].pValue,
			    attribs[2].ulValueLen, NULL);
			if (rsa->n && rsa->e &&
			    pkcs11_rsa_wrap(p, slotidx, &attribs[0], rsa) == 0) {
				key = key_new(KEY_UNSPEC);
				key->rsa = rsa;
				key->type = KEY_RSA;
				key->flags |= KEY_FLAG_EXT;
				/* expand key array and add key */
				*keysp = xrealloc(*keysp, *nkeys + 1,
				    sizeof(Key *));
				(*keysp)[*nkeys] = key;
				*nkeys = *nkeys + 1;
				debug("have %d keys", *nkeys);
			} else {
				RSA_free(rsa);
			}
		}
		for (i = 0; i < 3; i++)
			xfree(attribs[i].pValue);
	}
	if ((rv = f->C_FindObjectsFinal(session)) != CKR_OK)
		error("C_FindObjectsFinal failed: %lu", rv);
	return (0);
}

#ifdef HAVE_DLOPEN
/* register a new provider, fails if provider already exists */
int
pkcs11_add_provider(char *provider_id, char *pin, Key ***keyp)
{
	int nkeys, need_finalize = 0;
	struct pkcs11_provider *p = NULL;
	void *handle = NULL;
	CK_RV (*getfunctionlist)(CK_FUNCTION_LIST **);
	CK_RV rv;
	CK_FUNCTION_LIST *f = NULL;
	CK_TOKEN_INFO *token;
	CK_ULONG i;

	*keyp = NULL;
	if (pkcs11_provider_lookup(provider_id) != NULL) {
		error("provider already registered: %s", provider_id);
		goto fail;
	}
	/* open shared pkcs11-libarary */
	if ((handle = dlopen(provider_id, RTLD_NOW)) == NULL) {
		error("dlopen %s failed: %s", provider_id, dlerror());
		goto fail;
	}
	if ((getfunctionlist = dlsym(handle, "C_GetFunctionList")) == NULL) {
		error("dlsym(C_GetFunctionList) failed: %s", dlerror());
		goto fail;
	}
	p = xcalloc(1, sizeof(*p));
	p->name = xstrdup(provider_id);
	p->handle = handle;
	/* setup the pkcs11 callbacks */
	if ((rv = (*getfunctionlist)(&f)) != CKR_OK) {
		error("C_GetFunctionList failed: %lu", rv);
		goto fail;
	}
	p->function_list = f;
	if ((rv = f->C_Initialize(NULL)) != CKR_OK) {
		error("C_Initialize failed: %lu", rv);
		goto fail;
	}
	need_finalize = 1;
	if ((rv = f->C_GetInfo(&p->info)) != CKR_OK) {
		error("C_GetInfo failed: %lu", rv);
		goto fail;
	}
	rmspace(p->info.manufacturerID, sizeof(p->info.manufacturerID));
	rmspace(p->info.libraryDescription, sizeof(p->info.libraryDescription));
	debug("manufacturerID <%s> cryptokiVersion %d.%d"
	    " libraryDescription <%s> libraryVersion %d.%d",
	    p->info.manufacturerID,
	    p->info.cryptokiVersion.major,
	    p->info.cryptokiVersion.minor,
	    p->info.libraryDescription,
	    p->info.libraryVersion.major,
	    p->info.libraryVersion.minor);
	if ((rv = f->C_GetSlotList(CK_TRUE, NULL, &p->nslots)) != CKR_OK) {
		error("C_GetSlotList failed: %lu", rv);
		goto fail;
	}
	if (p->nslots == 0) {
		error("no slots");
		goto fail;
	}
	p->slotlist = xcalloc(p->nslots, sizeof(CK_SLOT_ID));
	if ((rv = f->C_GetSlotList(CK_TRUE, p->slotlist, &p->nslots))
	    != CKR_OK) {
		error("C_GetSlotList failed: %lu", rv);
		goto fail;
	}
	p->slotinfo = xcalloc(p->nslots, sizeof(struct pkcs11_slotinfo));
	p->valid = 1;
	nkeys = 0;
	for (i = 0; i < p->nslots; i++) {
		token = &p->slotinfo[i].token;
		if ((rv = f->C_GetTokenInfo(p->slotlist[i], token))
		    != CKR_OK) {
			error("C_GetTokenInfo failed: %lu", rv);
			continue;
		}
		rmspace(token->label, sizeof(token->label));
		rmspace(token->manufacturerID, sizeof(token->manufacturerID));
		rmspace(token->model, sizeof(token->model));
		rmspace(token->serialNumber, sizeof(token->serialNumber));
		debug("label <%s> manufacturerID <%s> model <%s> serial <%s>"
		    " flags 0x%lx",
		    token->label, token->manufacturerID, token->model,
		    token->serialNumber, token->flags);
		/* open session, login with pin and retrieve public keys */
		if (pkcs11_open_session(p, i, pin) == 0)
			pkcs11_fetch_keys(p, i, keyp, &nkeys);
	}
	if (nkeys > 0) {
		TAILQ_INSERT_TAIL(&pkcs11_providers, p, next);
		p->refcount++;	/* add to provider list */
		return (nkeys);
	}
	error("no keys");
	/* don't add the provider, since it does not have any keys */
fail:
	if (need_finalize && (rv = f->C_Finalize(NULL)) != CKR_OK)
		error("C_Finalize failed: %lu", rv);
	if (p) {
		if (p->slotlist)
			xfree(p->slotlist);
		if (p->slotinfo)
			xfree(p->slotinfo);
		xfree(p);
	}
	if (handle)
		dlclose(handle);
	return (-1);
}
#else
int
pkcs11_add_provider(char *provider_id, char *pin, Key ***keyp)
{
	error("dlopen() not supported");
	return (-1);
}
#endif