/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2011-2019 Pawel Jakub Dawidek <pawel@dawidek.net>
 * All rights reserved.
 *
 * 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 THE AUTHORS AND 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 THE AUTHORS OR 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.
 */

#include <sys/param.h>
#ifdef _KERNEL
#include <sys/kernel.h>
#include <sys/malloc.h>
#include <sys/sysctl.h>
#include <sys/systm.h>
#endif /* _KERNEL */
#include <sys/queue.h>
#include <sys/tree.h>

#include <geom/geom.h>

#include <geom/eli/g_eli.h>

#ifdef _KERNEL
MALLOC_DECLARE(M_ELI);

SYSCTL_DECL(_kern_geom_eli);
/*
 * The default limit (8192 keys) will allow to cache all keys for 4TB
 * provider with 512 bytes sectors and will take around 1MB of memory.
 */
static u_int g_eli_key_cache_limit = 8192;
SYSCTL_UINT(_kern_geom_eli, OID_AUTO, key_cache_limit, CTLFLAG_RDTUN,
    &g_eli_key_cache_limit, 0, "Maximum number of encryption keys to cache");
static uint64_t g_eli_key_cache_hits;
SYSCTL_UQUAD(_kern_geom_eli, OID_AUTO, key_cache_hits, CTLFLAG_RW,
    &g_eli_key_cache_hits, 0, "Key cache hits");
static uint64_t g_eli_key_cache_misses;
SYSCTL_UQUAD(_kern_geom_eli, OID_AUTO, key_cache_misses, CTLFLAG_RW,
    &g_eli_key_cache_misses, 0, "Key cache misses");

static int
g_eli_key_cmp(const struct g_eli_key *a, const struct g_eli_key *b)
{

	if (a->gek_keyno > b->gek_keyno)
		return (1);
	else if (a->gek_keyno < b->gek_keyno)
		return (-1);
	return (0);
}
#endif /* _KERNEL */

void
g_eli_key_fill(struct g_eli_softc *sc, struct g_eli_key *key, uint64_t keyno)
{
	const uint8_t *ekey;
	struct {
		char magic[4];
		uint8_t keyno[8];
	} __packed hmacdata;

	if ((sc->sc_flags & G_ELI_FLAG_ENC_IVKEY) != 0)
		ekey = sc->sc_mkey;
	else
		ekey = sc->sc_ekey;

	bcopy("ekey", hmacdata.magic, 4);
	le64enc(hmacdata.keyno, keyno);
	g_eli_crypto_hmac(ekey, G_ELI_MAXKEYLEN, (uint8_t *)&hmacdata,
	    sizeof(hmacdata), key->gek_key, 0);
	key->gek_keyno = keyno;
	key->gek_count = 0;
	key->gek_magic = G_ELI_KEY_MAGIC;
}

#ifdef _KERNEL
RB_PROTOTYPE(g_eli_key_tree, g_eli_key, gek_link, g_eli_key_cmp);
RB_GENERATE(g_eli_key_tree, g_eli_key, gek_link, g_eli_key_cmp);

static struct g_eli_key *
g_eli_key_allocate(struct g_eli_softc *sc, uint64_t keyno)
{
	struct g_eli_key *key, *ekey, keysearch;

	mtx_assert(&sc->sc_ekeys_lock, MA_OWNED);
	mtx_unlock(&sc->sc_ekeys_lock);

	key = malloc(sizeof(*key), M_ELI, M_WAITOK);
	g_eli_key_fill(sc, key, keyno);

	mtx_lock(&sc->sc_ekeys_lock);
	/*
	 * Recheck if the key wasn't added while we weren't holding the lock.
	 */
	keysearch.gek_keyno = keyno;
	ekey = RB_FIND(g_eli_key_tree, &sc->sc_ekeys_tree, &keysearch);
	if (ekey != NULL) {
		zfree(key, M_ELI);
		key = ekey;
		TAILQ_REMOVE(&sc->sc_ekeys_queue, key, gek_next);
	} else {
		RB_INSERT(g_eli_key_tree, &sc->sc_ekeys_tree, key);
		sc->sc_ekeys_allocated++;
	}
	TAILQ_INSERT_TAIL(&sc->sc_ekeys_queue, key, gek_next);

	return (key);
}

static struct g_eli_key *
g_eli_key_find_last(struct g_eli_softc *sc)
{
	struct g_eli_key *key;

	mtx_assert(&sc->sc_ekeys_lock, MA_OWNED);

	TAILQ_FOREACH(key, &sc->sc_ekeys_queue, gek_next) {
		if (key->gek_count == 0)
			break;
	}

	return (key);
}

static void
g_eli_key_replace(struct g_eli_softc *sc, struct g_eli_key *key, uint64_t keyno)
{

	mtx_assert(&sc->sc_ekeys_lock, MA_OWNED);
	KASSERT(key->gek_magic == G_ELI_KEY_MAGIC, ("Invalid magic."));

	RB_REMOVE(g_eli_key_tree, &sc->sc_ekeys_tree, key);
	TAILQ_REMOVE(&sc->sc_ekeys_queue, key, gek_next);

	KASSERT(key->gek_count == 0, ("gek_count=%d", key->gek_count));

	g_eli_key_fill(sc, key, keyno);

	RB_INSERT(g_eli_key_tree, &sc->sc_ekeys_tree, key);
	TAILQ_INSERT_TAIL(&sc->sc_ekeys_queue, key, gek_next);
}

static void
g_eli_key_remove(struct g_eli_softc *sc, struct g_eli_key *key)
{

	mtx_assert(&sc->sc_ekeys_lock, MA_OWNED);
	KASSERT(key->gek_magic == G_ELI_KEY_MAGIC, ("Invalid magic."));
	KASSERT(key->gek_count == 0, ("gek_count=%d", key->gek_count));

	RB_REMOVE(g_eli_key_tree, &sc->sc_ekeys_tree, key);
	TAILQ_REMOVE(&sc->sc_ekeys_queue, key, gek_next);
	sc->sc_ekeys_allocated--;
	zfree(key, M_ELI);
}

void
g_eli_key_init(struct g_eli_softc *sc)
{
	uint8_t *mkey;

	mtx_lock(&sc->sc_ekeys_lock);

	mkey = sc->sc_mkey + sizeof(sc->sc_ivkey);
	if ((sc->sc_flags & G_ELI_FLAG_AUTH) == 0)
		bcopy(mkey, sc->sc_ekey, G_ELI_DATAKEYLEN);
	else {
		/*
		 * The encryption key is: ekey = HMAC_SHA512(Data-Key, 0x10)
		 */
		g_eli_crypto_hmac(mkey, G_ELI_MAXKEYLEN, "\x10", 1,
		    sc->sc_ekey, 0);
	}

	if ((sc->sc_flags & G_ELI_FLAG_SINGLE_KEY) != 0) {
		sc->sc_ekeys_total = 1;
		sc->sc_ekeys_allocated = 0;
	} else {
		off_t mediasize;
		size_t blocksize;

		if ((sc->sc_flags & G_ELI_FLAG_AUTH) != 0) {
			struct g_provider *pp;

			pp = LIST_FIRST(&sc->sc_geom->consumer)->provider;
			mediasize = pp->mediasize;
			blocksize = pp->sectorsize;
		} else {
			mediasize = sc->sc_mediasize;
			blocksize = sc->sc_sectorsize;
		}
		sc->sc_ekeys_total =
		    ((mediasize - 1) >> G_ELI_KEY_SHIFT) / blocksize + 1;
		sc->sc_ekeys_allocated = 0;
		TAILQ_INIT(&sc->sc_ekeys_queue);
		RB_INIT(&sc->sc_ekeys_tree);
		if (sc->sc_ekeys_total <= g_eli_key_cache_limit) {
			uint64_t keyno;

			for (keyno = 0; keyno < sc->sc_ekeys_total; keyno++)
				(void)g_eli_key_allocate(sc, keyno);
			KASSERT(sc->sc_ekeys_total == sc->sc_ekeys_allocated,
			    ("sc_ekeys_total=%ju != sc_ekeys_allocated=%ju",
			    (uintmax_t)sc->sc_ekeys_total,
			    (uintmax_t)sc->sc_ekeys_allocated));
		}
	}

	mtx_unlock(&sc->sc_ekeys_lock);
}

void
g_eli_key_destroy(struct g_eli_softc *sc)
{

	mtx_lock(&sc->sc_ekeys_lock);
	if ((sc->sc_flags & G_ELI_FLAG_SINGLE_KEY) != 0) {
		explicit_bzero(sc->sc_ekey, sizeof(sc->sc_ekey));
	} else {
		struct g_eli_key *key;

		while ((key = TAILQ_FIRST(&sc->sc_ekeys_queue)) != NULL)
			g_eli_key_remove(sc, key);
		TAILQ_INIT(&sc->sc_ekeys_queue);
		RB_INIT(&sc->sc_ekeys_tree);
	}
	mtx_unlock(&sc->sc_ekeys_lock);
}

void
g_eli_key_resize(struct g_eli_softc *sc)
{
	uint64_t new_ekeys_total;
	off_t mediasize;
	size_t blocksize;

	if ((sc->sc_flags & G_ELI_FLAG_SINGLE_KEY) != 0) {
		return;
	}

	mtx_lock(&sc->sc_ekeys_lock);

	if ((sc->sc_flags & G_ELI_FLAG_AUTH) != 0) {
		struct g_provider *pp;

		pp = LIST_FIRST(&sc->sc_geom->consumer)->provider;
		mediasize = pp->mediasize;
		blocksize = pp->sectorsize;
	} else {
		mediasize = sc->sc_mediasize;
		blocksize = sc->sc_sectorsize;
	}
	new_ekeys_total = ((mediasize - 1) >> G_ELI_KEY_SHIFT) / blocksize + 1;
	/* We only allow to grow. */
	KASSERT(new_ekeys_total >= sc->sc_ekeys_total,
	    ("new_ekeys_total=%ju < sc_ekeys_total=%ju",
	    (uintmax_t)new_ekeys_total, (uintmax_t)sc->sc_ekeys_total));
	if (new_ekeys_total <= g_eli_key_cache_limit) {
		uint64_t keyno;

		for (keyno = sc->sc_ekeys_total; keyno < new_ekeys_total;
		    keyno++) {
			(void)g_eli_key_allocate(sc, keyno);
		}
		KASSERT(new_ekeys_total == sc->sc_ekeys_allocated,
		    ("new_ekeys_total=%ju != sc_ekeys_allocated=%ju",
		    (uintmax_t)new_ekeys_total,
		    (uintmax_t)sc->sc_ekeys_allocated));
	}

	sc->sc_ekeys_total = new_ekeys_total;

	mtx_unlock(&sc->sc_ekeys_lock);
}

/*
 * Select encryption key. If G_ELI_FLAG_SINGLE_KEY is present we only have one
 * key available for all the data. If the flag is not present select the key
 * based on data offset.
 */
uint8_t *
g_eli_key_hold(struct g_eli_softc *sc, off_t offset, size_t blocksize)
{
	struct g_eli_key *key, keysearch;
	uint64_t keyno;

	if ((sc->sc_flags & G_ELI_FLAG_SINGLE_KEY) != 0)
		return (sc->sc_ekey);

	/* We switch key every 2^G_ELI_KEY_SHIFT blocks. */
	keyno = (offset >> G_ELI_KEY_SHIFT) / blocksize;

	KASSERT(keyno < sc->sc_ekeys_total,
	    ("%s: keyno=%ju >= sc_ekeys_total=%ju",
	    __func__, (uintmax_t)keyno, (uintmax_t)sc->sc_ekeys_total));

	keysearch.gek_keyno = keyno;

	if (sc->sc_ekeys_total == sc->sc_ekeys_allocated) {
		/* We have all the keys, so avoid some overhead. */
		key = RB_FIND(g_eli_key_tree, &sc->sc_ekeys_tree, &keysearch);
		KASSERT(key != NULL, ("No key %ju found.", (uintmax_t)keyno));
		KASSERT(key->gek_magic == G_ELI_KEY_MAGIC,
		    ("Invalid key magic."));
		return (key->gek_key);
	}

	mtx_lock(&sc->sc_ekeys_lock);
	key = RB_FIND(g_eli_key_tree, &sc->sc_ekeys_tree, &keysearch);
	if (key != NULL) {
		g_eli_key_cache_hits++;
		TAILQ_REMOVE(&sc->sc_ekeys_queue, key, gek_next);
		TAILQ_INSERT_TAIL(&sc->sc_ekeys_queue, key, gek_next);
	} else {
		/*
		 * No key in cache, find the least recently unreferenced key
		 * or allocate one if we haven't reached our limit yet.
		 */
		if (sc->sc_ekeys_allocated < g_eli_key_cache_limit) {
			key = g_eli_key_allocate(sc, keyno);
		} else {
			g_eli_key_cache_misses++;
			key = g_eli_key_find_last(sc);
			if (key != NULL) {
				g_eli_key_replace(sc, key, keyno);
			} else {
				/* All keys are referenced? Allocate one. */
				key = g_eli_key_allocate(sc, keyno);
			}
		}
	}
	key->gek_count++;
	mtx_unlock(&sc->sc_ekeys_lock);

	KASSERT(key->gek_magic == G_ELI_KEY_MAGIC, ("Invalid key magic."));

	return (key->gek_key);
}

void
g_eli_key_drop(struct g_eli_softc *sc, uint8_t *rawkey)
{
	struct g_eli_key *key = (struct g_eli_key *)rawkey;

	if ((sc->sc_flags & G_ELI_FLAG_SINGLE_KEY) != 0)
		return;

	KASSERT(key->gek_magic == G_ELI_KEY_MAGIC, ("Invalid key magic."));

	if (sc->sc_ekeys_total == sc->sc_ekeys_allocated)
		return;

	mtx_lock(&sc->sc_ekeys_lock);
	KASSERT(key->gek_count > 0, ("key->gek_count=%d", key->gek_count));
	key->gek_count--;
	while (sc->sc_ekeys_allocated > g_eli_key_cache_limit) {
		key = g_eli_key_find_last(sc);
		if (key == NULL)
			break;
		g_eli_key_remove(sc, key);
	}
	mtx_unlock(&sc->sc_ekeys_lock);
}
#endif /* _KERNEL */