/*
 * Copyright (c) 2010, 2024, 2025 The DragonFly Project.  All rights reserved.
 *
 * This code is derived from software contributed to The DragonFly Project
 * by Alex Hornung <ahornung@gmail.com> and
 * Michael Neumann <mneumann@ntecs.de>.
 *
 * 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.
 * 3. Neither the name of The DragonFly Project nor the names of its
 *    contributors may be used to endorse or promote products derived
 *    from this software without specific, prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS 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
 * COPYRIGHT HOLDERS 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.
 */

/*
 * This file implements initial version of device-mapper crypt target.
 */

#include <sys/systm.h>
#include <sys/endian.h>
#include <sys/bio.h>
#include <sys/kerneldump.h>
#include <sys/malloc.h>
#include <sys/md5.h>
#include <crypto/sha1.h>
#include <crypto/sha2/sha2.h>
#include <crypto/rmd160/rmd160.h>
#include <crypto/cryptoapi/cryptoapi.h>
#include <dev/disk/dm/dm.h>

#include <sys/proc.h>
#include <sys/types.h>
#include <sys/mpipe.h>

MALLOC_DEFINE(M_DMCRYPT, "dm_crypt", "Device Mapper Target Crypt");


struct target_crypt_config;

typedef void ivgen_t(struct target_crypt_config *, u_int8_t *, size_t, off_t);

typedef int ivgen_ctor_t(struct target_crypt_config *, char *, void **);
typedef int ivgen_dtor_t(struct target_crypt_config *, void *);


struct iv_generator {
	const char	*name;
	ivgen_ctor_t	*ctor;
	ivgen_dtor_t	*dtor;
	ivgen_t		*gen_iv;
};

struct essiv_ivgen_priv {
	cryptoapi_cipher_session_t	crypto_session;
	size_t			keyhash_len;
	u_int8_t		crypto_keyhash[SHA512_DIGEST_LENGTH];
};

typedef struct target_crypt_config {
	size_t	params_len;
	dm_pdev_t *pdev;
	char	*status_str;
	cryptoapi_cipher_t	crypto_cipher;
	int	crypto_klen;
	u_int8_t	crypto_key[512>>3];
	cryptoapi_cipher_session_t crypto_session;

	u_int64_t	block_offset;
	int64_t		iv_offset;
	SHA512_CTX	essivsha512_ctx;

	struct iv_generator	*ivgen;
	void	*ivgen_priv;

	struct malloc_pipe	read_mpipe;
	struct malloc_pipe	write_mpipe;

} dm_target_crypt_config_t;

struct dmtc_dump_helper {
	dm_target_crypt_config_t *priv;
	void *data;
	size_t length;
	off_t offset;

	u_char space[65536];
};

static void dmtc_init_mpipe(struct target_crypt_config *priv);
static void dmtc_destroy_mpipe(struct target_crypt_config *priv);

static cryptoapi_cipher_t
dmtc_find_crypto_cipher(const char *crypto_alg, const char *crypto_mode,
    int klen_in_bits);

#define DMTC_BUF_SIZE (MAXPHYS)

static void dmtc_bio_read_done(struct bio *bio);

static void dmtc_bio_read_decrypt_start(struct bio *bio);

static void dmtc_bio_read_decrypt_retry(void *arg1, void *arg2);

static void dmtc_bio_read_decrypt(struct bio *bio, uint8_t *data_buf,
    size_t data_buf_sz);

static void dmtc_bio_write_encrypt_start(struct bio *bio);

static void dmtc_bio_write_encrypt_retry(void *arg1, void *arg2);

static void dmtc_bio_write_encrypt(struct bio *bio, uint8_t *data_buf,
    size_t data_buf_sz);

static void dmtc_bio_write_done(struct bio *bio);

static int dmtc_bio_encdec(dm_target_crypt_config_t *priv, uint8_t *data_buf,
    int bytes, off_t offset, cryptoapi_cipher_mode mode);


static void dmtc_crypto_dump(dm_target_crypt_config_t *priv,
    struct dmtc_dump_helper *dump_helper);

static ivgen_ctor_t	essiv_ivgen_ctor;
static ivgen_dtor_t	essiv_ivgen_dtor;
static ivgen_t		essiv_ivgen;
static ivgen_t		plain_ivgen;
static ivgen_t		plain64_ivgen;

static struct iv_generator ivgens[] = {
	{ .name = "essiv", .ctor = essiv_ivgen_ctor, .dtor = essiv_ivgen_dtor,
	    .gen_iv = essiv_ivgen },
	{ .name = "plain", .ctor = NULL, .dtor = NULL, .gen_iv = plain_ivgen },
	{ .name = "plain64", .ctor = NULL, .dtor = NULL, .gen_iv = plain64_ivgen },
	{ NULL, NULL, NULL, NULL }
};

static __inline int
clamp_value(int value, int min, int max)
{
	if (value < min)
		return min;
	if (value > max)
		return max;
	return value;
}

static __inline int
dmtc_get_nmax(void)
{
	int nmax;
	nmax = (physmem * 2 / 1000 * PAGE_SIZE) /
		(2 * DMTC_BUF_SIZE) + 1;
	return clamp_value(nmax, 2, 8 + ncpus * 2);
}

static void
dmtc_zero_mpipe_buffer(void *buffer, void *priv __unused)
{
	explicit_bzero(buffer, DMTC_BUF_SIZE);
}

static void
dmtc_init_mpipe(struct target_crypt_config *priv)
{
	int nmax = dmtc_get_nmax();
	int nmax_read = nmax;
	int nmax_write = nmax;

	kprintf("dm_target_crypt: Setting %d mpipe read buffers\n", nmax_read);
	kprintf("dm_target_crypt: Setting %d mpipe write buffers\n", nmax_write);

	mpipe_init(&priv->read_mpipe, M_DMCRYPT, DMTC_BUF_SIZE,
		nmax_read, nmax_read, MPF_NOZERO | MPF_CALLBACK,
		NULL, dmtc_zero_mpipe_buffer, NULL);

	mpipe_init(&priv->write_mpipe, M_DMCRYPT, DMTC_BUF_SIZE,
		nmax_write, nmax_write, MPF_NOZERO | MPF_CALLBACK,
		NULL, dmtc_zero_mpipe_buffer, NULL);
}

static void
dmtc_destroy_mpipe(struct target_crypt_config *priv)
{
	mpipe_done(&priv->read_mpipe);
	mpipe_done(&priv->write_mpipe);
}

/*
 * Overwrite private information (in buf) to avoid leaking it
 */
static void
dmtc_crypto_clear(void *buf, size_t len)
{
	memset(buf, 0xFF, len);
	explicit_bzero(buf, len);
}

/*
 * ESSIV IV Generator Routines
 */
static int
essiv_ivgen_ctor(struct target_crypt_config *priv, char *iv_hash, void **p_ivpriv)
{
	struct essiv_ivgen_priv *ivpriv;
	u_int8_t crypto_keyhash[SHA512_DIGEST_LENGTH];
	unsigned int klen, hashlen;
	int error;

	klen = (priv->crypto_klen >> 3);

	if (iv_hash == NULL)
		return EINVAL;

	if (strcmp(iv_hash, "sha1") == 0) {
		SHA1_CTX ctx;

		hashlen = SHA1_RESULTLEN;
		SHA1Init(&ctx);
		SHA1Update(&ctx, priv->crypto_key, klen);
		SHA1Final(crypto_keyhash, &ctx);
	} else if (strcmp(iv_hash, "sha256") == 0) {
		SHA256_CTX ctx;

		hashlen = SHA256_DIGEST_LENGTH;
		SHA256_Init(&ctx);
		SHA256_Update(&ctx, priv->crypto_key, klen);
		SHA256_Final(crypto_keyhash, &ctx);
	} else if (strcmp(iv_hash, "sha384") == 0) {
		SHA384_CTX ctx;

		hashlen = SHA384_DIGEST_LENGTH;
		SHA384_Init(&ctx);
		SHA384_Update(&ctx, priv->crypto_key, klen);
		SHA384_Final(crypto_keyhash, &ctx);
	} else if (strcmp(iv_hash, "sha512") == 0) {
		SHA512_CTX ctx;

		hashlen = SHA512_DIGEST_LENGTH;
		SHA512_Init(&ctx);
		SHA512_Update(&ctx, priv->crypto_key, klen);
		SHA512_Final(crypto_keyhash, &ctx);
	} else if (strcmp(iv_hash, "md5") == 0) {
		MD5_CTX ctx;

		hashlen = MD5_DIGEST_LENGTH;
		MD5Init(&ctx);
		MD5Update(&ctx, priv->crypto_key, klen);
		MD5Final(crypto_keyhash, &ctx);
	} else if (strcmp(iv_hash, "rmd160") == 0 ||
		   strcmp(iv_hash, "ripemd160") == 0) {
		RMD160_CTX ctx;

		hashlen = 160/8;
		RMD160Init(&ctx);
		RMD160Update(&ctx, priv->crypto_key, klen);
		RMD160Final(crypto_keyhash, &ctx);
	} else {
		return EINVAL;
	}

	/* Convert hashlen to bits */
	hashlen <<= 3;

	ivpriv = kmalloc(sizeof(struct essiv_ivgen_priv), M_DMCRYPT,
			 M_WAITOK | M_ZERO);
	memcpy(ivpriv->crypto_keyhash, crypto_keyhash, sizeof(crypto_keyhash));
	ivpriv->keyhash_len = sizeof(crypto_keyhash);
	dmtc_crypto_clear(crypto_keyhash, sizeof(crypto_keyhash));


	/*
	 * XXX: in principle we also need to check if the block size of the
	 *	cipher is a valid iv size for the block cipher.
	 */

	ivpriv->crypto_session = cryptoapi_cipher_newsession(priv->crypto_cipher);
	if (ivpriv->crypto_session == NULL) {
		kprintf("dm_target_crypt: Error during cryptoapi_cipher_newsession "
			"for essiv_ivgen\n");
		dmtc_crypto_clear(ivpriv->crypto_keyhash, ivpriv->keyhash_len);
		kfree(ivpriv, M_DMCRYPT);
		return ENOTSUP;
	}

	error = cryptoapi_cipher_setkey(ivpriv->crypto_session,
		(const uint8_t *)ivpriv->crypto_keyhash,
		hashlen / 8);

	if (error) {
		kprintf("dm_target_crypt: Error during cryptoapi_cipher_setkey "
			"for essiv_ivgen, error = %d\n",
			error);
		cryptoapi_cipher_freesession(ivpriv->crypto_session);
		dmtc_crypto_clear(ivpriv->crypto_keyhash, ivpriv->keyhash_len);
		kfree(ivpriv, M_DMCRYPT);
		return ENOTSUP;
	}

	*p_ivpriv = ivpriv;

	return 0;
}

static int
essiv_ivgen_dtor(struct target_crypt_config *priv, void *arg)
{
	struct essiv_ivgen_priv *ivpriv;

	ivpriv = (struct essiv_ivgen_priv *)arg;
	KKASSERT(ivpriv != NULL);

	cryptoapi_cipher_freesession(ivpriv->crypto_session);

	dmtc_crypto_clear(ivpriv->crypto_keyhash, ivpriv->keyhash_len);
	kfree(ivpriv, M_DMCRYPT);

	return 0;
}

static void
essiv_ivgen(dm_target_crypt_config_t *priv, u_int8_t *iv,
	    size_t iv_len, off_t sector)
{
	struct essiv_ivgen_priv *ivpriv;
	int error;

	ivpriv = priv->ivgen_priv;
	KKASSERT(ivpriv != NULL);

	bzero(iv, iv_len);
	*((off_t *)iv) = htole64(sector + priv->iv_offset);

	cryptoapi_cipher_iv iv2;
	bzero(iv2, sizeof(iv2));

	error = cryptoapi_cipher_encrypt(ivpriv->crypto_session,
			iv,
			iv_len,
			iv2,
			sizeof(iv2)
			);

	explicit_bzero(iv2, sizeof(iv2));

	if (error)
		kprintf("dm_target_crypt: essiv_ivgen, error = %d\n", error);
}


static void
plain_ivgen(dm_target_crypt_config_t *priv, u_int8_t *iv,
	    size_t iv_len, off_t sector)
{
	bzero(iv, iv_len);
	*((uint32_t *)iv) = htole32((uint32_t)(sector + priv->iv_offset));
}

static void
plain64_ivgen(dm_target_crypt_config_t *priv, u_int8_t *iv,
    size_t iv_len, off_t sector)
{
	bzero(iv, iv_len);
	*((uint64_t *)iv) = htole64((uint64_t)(sector + priv->iv_offset));
}

/*
 * Init function called from dm_table_load_ioctl.
 * cryptsetup actually passes us this:
 * aes-cbc-essiv:sha256 7997f8af... 0 /dev/ad0s0a 8
 */
static int
hex2key(char *hex, size_t key_len, u_int8_t *key)
{
	char hex_buf[3];
	size_t key_idx;

	hex_buf[2] = 0;
	for (key_idx = 0; key_idx < key_len; ++key_idx) {
		hex_buf[0] = *hex++;
		hex_buf[1] = *hex++;
		key[key_idx] = (u_int8_t)strtoul(hex_buf, NULL, 16);
	}
	hex_buf[0] = 0;
	hex_buf[1] = 0;

	return 0;
}

/**
 * Map between dm_target_crypt algorithm naming and our own cryptoapi
 * naming. It happens that they are identical, but it doesn't have to be this
 * way.
 */
static cryptoapi_cipher_t
dmtc_find_crypto_cipher(const char *crypto_alg, const char *crypto_mode,
    int klen_in_bits)
{
#define ALGO_MODE_EQ(algo, mode) \
	((strcmp(crypto_alg, algo) == 0) && (strcmp(crypto_mode, mode) == 0))

	if (ALGO_MODE_EQ("aes", "cbc"))
		return cryptoapi_cipher_find("aes-cbc", klen_in_bits);

	if (ALGO_MODE_EQ("aes", "xts"))
		return cryptoapi_cipher_find("aes-xts", klen_in_bits);

	if (ALGO_MODE_EQ("twofish", "cbc"))
		return cryptoapi_cipher_find("twofish-cbc", klen_in_bits);

	if (ALGO_MODE_EQ("twofish", "xts"))
		return cryptoapi_cipher_find("twofish-xts", klen_in_bits);

	if (ALGO_MODE_EQ("serpent", "cbc"))
		return cryptoapi_cipher_find("serpent-cbc", klen_in_bits);

	if (ALGO_MODE_EQ("serpent", "xts"))
		return cryptoapi_cipher_find("serpent-xts", klen_in_bits);

	kprintf("dm_target_crypt: unsupported algo: %s and mode: %s\n",
	    crypto_alg, crypto_mode);

	return NULL;
}

static int
dm_target_crypt_init(dm_table_entry_t *table_en, int argc, char **argv)
{
	dm_target_crypt_config_t *priv;
	size_t len;
	char *crypto_alg, *crypto_mode, *iv_mode, *iv_opt, *hex_key, *dev;
	char *status_str;
	int i, klen_in_bits, error;
	uint64_t iv_offset, block_offset;

	if (argc != 5) {
		kprintf("dm_target_crypt: not enough arguments, "
			"need exactly 5\n");
		return EINVAL;
	}

	len = 0;
	for (i = 0; i < argc; i++) {
		len += strlen(argv[i]);
		len++;
	}
	/* len is strlen() of input string +1 */
	status_str = kmalloc(len, M_DMCRYPT, M_WAITOK);

	crypto_alg = strsep(&argv[0], "-");
	crypto_mode = strsep(&argv[0], "-");
	iv_opt = strsep(&argv[0], "-");
	iv_mode = strsep(&iv_opt, ":");
	hex_key = argv[1];
	iv_offset = strtouq(argv[2], NULL, 0);
	dev = argv[3];
	block_offset = strtouq(argv[4], NULL, 0);
	/* hex_key is specified in hex, so one hex char == 4 bits */
	klen_in_bits = strlen(hex_key) * 4;

#if 0
	kprintf("dm_target_crypt - new: dev=%s, crypto_alg=%s, crypto_mode=%s, "
		"iv_mode=%s, iv_opt=%s, key=%s, iv_offset=%ju, "
		"block_offset=%ju\n",
		dev, crypto_alg, crypto_mode, iv_mode, iv_opt, hex_key, iv_offset,
		block_offset);
#endif

	priv = kmalloc(sizeof(dm_target_crypt_config_t), M_DMCRYPT, M_WAITOK);

	/* Insert dmp to global pdev list */
	if ((priv->pdev = dm_pdev_insert(dev)) == NULL) {
		kprintf("dm_target_crypt: dm_pdev_insert failed\n");
		kfree(status_str, M_DMCRYPT);
		return ENOENT;
	}

	priv->crypto_cipher = dmtc_find_crypto_cipher(crypto_alg, crypto_mode, klen_in_bits);
	priv->crypto_klen = klen_in_bits;
	if (priv->crypto_cipher == NULL)
		goto notsup;

	kprintf("dm_target_crypt: using crypto_cipher: %s\n",
			cryptoapi_cipher_get_description(priv->crypto_cipher));

	/* Save length of param string */
	priv->params_len = len;
	priv->block_offset = block_offset;
	priv->iv_offset = iv_offset - block_offset;

	dm_table_add_deps(table_en, priv->pdev);

	dm_table_init_target(table_en, priv);

	error = hex2key(hex_key, priv->crypto_klen / 8,
			(u_int8_t *)priv->crypto_key);

	if (error) {
		kprintf("dm_target_crypt: hex2key failed, "
			"invalid key format\n");
		goto notsup;
	}

	/* Handle cmd */
	for(i = 0; ivgens[i].name != NULL; i++) {
		if (strcmp(iv_mode, ivgens[i].name) == 0)
			break;
	}

	if (ivgens[i].name == NULL) {
		kprintf("dm_target_crypt: iv_mode='%s' unsupported\n",
			iv_mode);
		goto notsup;
	}

	/* Call our ivgen constructor */
	if (ivgens[i].ctor != NULL) {
		error = ivgens[i].ctor(priv, iv_opt,
		    &priv->ivgen_priv);
		if (error) {
			kprintf("dm_target_crypt: ctor for '%s' failed\n",
			    ivgens[i].name);
			goto notsup;
		}
	}

	priv->ivgen = &ivgens[i];

	priv->crypto_session = cryptoapi_cipher_newsession(priv->crypto_cipher);
	if (priv->crypto_session == NULL) {
		kprintf("dm_target_crypt: Error during cryptoapi_cipher_newsession\n");
		goto notsup;
	}

	error = cryptoapi_cipher_setkey(priv->crypto_session,
		(const u_int8_t *)priv->crypto_key,
		priv->crypto_klen / 8);

	if (error) {
		kprintf("dm_target_crypt: Error during cryptoapi_cipher_setkey, "
			"error = %d\n",
			error);
		cryptoapi_cipher_freesession(priv->crypto_session);
		goto notsup;
	}

	memset(hex_key, '0', strlen(hex_key));
	if (iv_opt) {
		ksprintf(status_str, "%s-%s-%s:%s %s %ju %s %ju",
		    crypto_alg, crypto_mode, iv_mode, iv_opt,
		    hex_key, iv_offset, dev, block_offset);
	} else {
		ksprintf(status_str, "%s-%s-%s %s %ju %s %ju",
		    crypto_alg, crypto_mode, iv_mode,
		    hex_key, iv_offset, dev, block_offset);
	}
	priv->status_str = status_str;

	/* Initialize mpipes */
	dmtc_init_mpipe(priv);

	return 0;

notsup:
	kprintf("dm_target_crypt: ENOTSUP\n");
	kfree(status_str, M_DMCRYPT);
	return ENOTSUP;
}

/* Table routine called to get params string. */
static char *
dm_target_crypt_table(void *target_config)
{
	dm_target_crypt_config_t *priv;
	char *params;

	priv = target_config;

	params = dm_alloc_string(DM_MAX_PARAMS_SIZE);

	ksnprintf(params, DM_MAX_PARAMS_SIZE, "%s",
	    priv->status_str);

	return params;
}

static int
dm_target_crypt_destroy(dm_table_entry_t *table_en)
{
	dm_target_crypt_config_t *priv;

	/*
	 * Disconnect the crypt config before unbusying the target.
	 */
	priv = table_en->target_config;
	if (priv == NULL)
		return 0;
	dm_pdev_decr(priv->pdev);

	dmtc_destroy_mpipe(priv);

	/*
	 * Clean up the crypt config
	 *
	 * Overwrite the private information before freeing memory to
	 * avoid leaking it.
	 */
	if (priv->status_str) {
		dmtc_crypto_clear(priv->status_str, strlen(priv->status_str));
		kfree(priv->status_str, M_DMCRYPT);
	}

	if ((priv->ivgen) && (priv->ivgen->dtor != NULL)) {
		priv->ivgen->dtor(priv, priv->ivgen_priv);
	}

	cryptoapi_cipher_freesession(priv->crypto_session);

	dmtc_crypto_clear(priv, sizeof(dm_target_crypt_config_t));
	kfree(priv, M_DMCRYPT);

	return 0;
}

/************************************************************************
 *			STRATEGY SUPPORT FUNCTIONS			*
 ************************************************************************
 *
 * READ PATH:	doio -> bio_read_done -> bio_read_decrypt
 * WRITE PATH:	bio_write_encrypt -> doio -> bio_write_done
 */

/**
 * Usage of the "struct bio" bio_caller_infoX fields:
 *
 * bio_caller_info1:
 * 	- always points to "priv"
 * bio_caller_info2:
 * 	- used to chain bio requests. points to next enqueued bio request
 * 	  (only within bio request queue)
 *	- orig b_data pointer (within write path)
 * bio_caller_info3:
 * 	- points to the mpipe
 */

/*
 * Start IO operation, called from dmstrategy routine.
 */
static int
dm_target_crypt_strategy(dm_table_entry_t *table_en, struct buf *bp)
{
	struct bio *bio;

	dm_target_crypt_config_t *priv;
	priv = table_en->target_config;

	/* Get rid of stuff we can't really handle */
	if ((bp->b_cmd == BUF_CMD_READ) ||
	    (bp->b_cmd == BUF_CMD_WRITE)) {
		if (((bp->b_bcount % DEV_BSIZE) != 0) ||
		    (bp->b_bcount == 0)) {
			kprintf(
			    "dm_target_crypt_strategy: can't really "
			    "handle bp->b_bcount = %d\n",
			    bp->b_bcount);
			bp->b_error = EINVAL;
			bp->b_flags |= B_ERROR | B_INVAL;
			biodone(&bp->b_bio1);
			return 0;
		}
	}

	switch (bp->b_cmd) {
	case BUF_CMD_READ:
		bio = push_bio(&bp->b_bio1);
		bio->bio_offset = bp->b_bio1.bio_offset +
		    priv->block_offset * DEV_BSIZE;
		bio->bio_caller_info1.ptr = priv;
		bio->bio_done = dmtc_bio_read_done;
		vn_strategy(priv->pdev->pdev_vnode, bio);
		break;
	case BUF_CMD_WRITE:
		bio = push_bio(&bp->b_bio1);
		bio->bio_offset = bp->b_bio1.bio_offset +
		    priv->block_offset * DEV_BSIZE;
		bio->bio_caller_info1.ptr = priv;
		dmtc_bio_write_encrypt_start(bio);
		break;
	default:
		vn_strategy(priv->pdev->pdev_vnode, &bp->b_bio1);
		break;
	}
	return 0;
}

/************************************************************************
 * READ PATH
 ************************************************************************
 *
 * DO IO -> dmtc_bio_read_done -> dmtc_bio_read_decrypt_start ->
 * (dmtc_bio_read_decrypt_retry) -> dmtc_bio_read_decrypt -> COMPLETE
 *
 */

/*
 * Called after read BIO completes
 */
static void
dmtc_bio_read_done(struct bio *bio)
{
	struct bio *obio;

	/*
	 * If a read error occurs we shortcut the operation, otherwise
	 * go on to stage 2 (decrypt).
	 */
	if (bio->bio_buf->b_flags & B_ERROR) {
		obio = pop_bio(bio);
		biodone(obio);
	} else {
		dmtc_bio_read_decrypt_start(bio);
	}
}

/**
 * Starts decryption by allocating a buffer.
 *
 * If allocation fails, dmtc_bio_read_decrypt_retry is called.
 */
void
dmtc_bio_read_decrypt_start(struct bio *bio)
{
	dm_target_crypt_config_t *priv;
	uint8_t *data_buf;

	priv = bio->bio_caller_info1.ptr;

	KKASSERT(bio->bio_buf->b_cmd == BUF_CMD_READ);

	data_buf = mpipe_alloc_callback(&priv->read_mpipe,
	    dmtc_bio_read_decrypt_retry, bio, NULL);

	if (data_buf == NULL)
		return;

	bio->bio_caller_info3.ptr = &priv->read_mpipe;
	dmtc_bio_read_decrypt(bio, data_buf, DMTC_BUF_SIZE);
}

/**
 * Retries the allocation.
 */
void
dmtc_bio_read_decrypt_retry(void *arg1, void *arg2 __unused)
{
	dmtc_bio_read_decrypt_start(arg1);
}

/**
 * Decrypts the buffer, releases the intermediate mpipe buffer and completes the
 * bio request.
 */
void
dmtc_bio_read_decrypt(struct bio *bio, uint8_t *data_buf, size_t data_buf_sz)
{
	dm_target_crypt_config_t *priv;
	struct malloc_pipe *mpipe;

	priv = bio->bio_caller_info1.ptr;
	mpipe = bio->bio_caller_info3.ptr;

	/*
	 * Note: b_resid no good after read I/O, it will be 0, use
	 *	 b_bcount.
	 */
	int bytes = bio->bio_buf->b_bcount;

	if (__predict_false(data_buf_sz < bytes))
		panic("dmtc: Allocated data buffer is too small");

	/*
	 * Unconditionally copy in data. Never decrypt in place!
	 *
	 * For reads with bogus page we can't decrypt in place as stuff
	 * can get ripped out from under us.
	 */
	memcpy(data_buf, bio->bio_buf->b_data, bytes);

	bio->bio_buf->b_error = dmtc_bio_encdec(priv, data_buf, bytes,
	    bio->bio_offset, CRYPTOAPI_CIPHER_DECRYPT);

	if (bio->bio_buf->b_error) {
		kprintf("dm_target_crypt: dmtc_bio_read_decrypt error = %d\n",
		    bio->bio_buf->b_error);

		bio->bio_buf->b_flags |= B_ERROR;
	} else {
		memcpy(bio->bio_buf->b_data, data_buf, bytes);
	}
#if 0
	else if (bio->bio_buf->b_flags & B_HASBOGUS) {
		memcpy(bio->bio_buf->b_data, dmtc->data_buf,
		       bio->bio_buf->b_bcount);
	}
#endif

	struct bio *obio = pop_bio(bio);
	biodone(obio);

	if (mpipe)
		mpipe_free(mpipe, data_buf);
}

/************************************************************************
 * WRITE PATH
 ************************************************************************
 *
 * dmtc_bio_write_encrypt_start -> (dmtc_bio_write_encrypt_retry) ->
 * dmtc_bio_write_encrypt -> DO IO -> dmtc_bio_write_done -> COMPLETE
 *
 */

/**
 * Allocates a mpipe buffer in order to proceed with encryption.
 *
 * If that fails, dmtc_bio_write_encrypt_retry will be called.
 */
void
dmtc_bio_write_encrypt_start(struct bio *bio)
{
	dm_target_crypt_config_t *priv;
	uint8_t *data_buf;

	priv = bio->bio_caller_info1.ptr;

	KKASSERT(bio->bio_buf->b_cmd == BUF_CMD_WRITE);

	data_buf = mpipe_alloc_callback(&priv->write_mpipe,
	    dmtc_bio_write_encrypt_retry, bio, NULL);

	if (data_buf == NULL)
		return;

	bio->bio_caller_info3.ptr = &priv->write_mpipe;
	dmtc_bio_write_encrypt(bio, data_buf, DMTC_BUF_SIZE);
}

/**
 * Retries the allocation.
 */
void
dmtc_bio_write_encrypt_retry(void *arg1, void *arg2 __unused)
{
	dmtc_bio_write_encrypt_start(arg1);
}

/**
 *
 * Encrypts the buffer and sends a write request for the encrypted buffer to the
 * underlying device via vn_strategy. Once this completes, dmtc_bio_write_done
 * will be called.
 */
void
dmtc_bio_write_encrypt(struct bio *bio, uint8_t *data_buf, size_t data_buf_sz)
{
	dm_target_crypt_config_t *priv = bio->bio_caller_info1.ptr;

	/*
	 * Use b_bcount for consistency
	 */
	int bytes = bio->bio_buf->b_bcount;

	if (__predict_false(data_buf_sz < bytes))
		panic("dmtc: Allocated data buffer is too small");

	memcpy(data_buf, bio->bio_buf->b_data, bytes);

	bio->bio_buf->b_error = dmtc_bio_encdec(priv, data_buf, bytes,
	    bio->bio_offset, CRYPTOAPI_CIPHER_ENCRYPT);

	if (bio->bio_buf->b_error) {
		kprintf("dm_target_crypt: dmtc_bio_write_encrypt error = %d\n",
		    bio->bio_buf->b_error);

		mpipe_free(&priv->write_mpipe, data_buf);
		bio->bio_buf->b_flags |= B_ERROR;
		struct bio *obio = pop_bio(bio);
		biodone(obio);
	} else {
		bio->bio_caller_info2.ptr = bio->bio_buf->b_data; /* orig_buf */
		bio->bio_buf->b_data = data_buf;
		bio->bio_done = dmtc_bio_write_done;

		/*
		 * Note: We just copy back data_buf to bio->bio_buf->b_data
		 * and then call vn_stategy() on it.
		 *
		 * We have to allocate our own buffer, copy bio->bio_buf->b_data
		 * into it, modify the copy, rewire bio->bio_buf->b_data to our
		 * own buffer, call vn_stategy() and once it finished with
		 * dmtc_bio_write_done(), free our buffer.
		 */
		vn_strategy(priv->pdev->pdev_vnode, bio);
	}
}

/**
 * The write completed, so releases the intermediate buffer.
 */
void
dmtc_bio_write_done(struct bio *bio)
{
	dm_target_crypt_config_t *priv;
	uint8_t *data_buf;
	struct malloc_pipe *mpipe;

	priv = bio->bio_caller_info1.ptr;
	data_buf = bio->bio_buf->b_data;
	mpipe = bio->bio_caller_info3.ptr;

	if (mpipe)
		mpipe_free(mpipe, data_buf);

	// Restore original bio buffer
	bio->bio_buf->b_data = bio->bio_caller_info2.ptr;

	struct bio *obio = pop_bio(bio);
	biodone(obio);
}

/************************************************************************
 * COMMON
 ************************************************************************/

/**
 * Encrypts or decrypts `data_buf`.
 */
int
dmtc_bio_encdec(dm_target_crypt_config_t *priv, uint8_t *data_buf, int bytes,
    off_t offset, cryptoapi_cipher_mode mode)
{
	cryptoapi_cipher_iv iv;
	int sectors = bytes / DEV_BSIZE;    /* Number of sectors */
	off_t isector = offset / DEV_BSIZE; /* ivgen salt base? */
	int error = 0;

	KKASSERT((sectors * DEV_BSIZE) == bytes);

	for (int i = 0; i < sectors; i++) {
		/*
		 * Note: last argument is used to generate salt(?) and
		 * is a 64 bit value, but the original code passed an
		 *	 int.  Changing it now will break pre-existing
		 *	 crypt volumes.
		 */
		priv->ivgen->gen_iv(priv, iv, sizeof(iv), isector + i);

		error = cryptoapi_cipher_crypt(priv->crypto_session,
				data_buf + i * DEV_BSIZE,
				DEV_BSIZE, iv, sizeof(iv), mode);

		if (error) {
			break;
		}
	}

	// TODO: required?
	explicit_bzero(&iv, sizeof(iv));

	return (error);
}
/* DUMPING MAGIC */

extern int tsleep_crypto_dump;

static int
dm_target_crypt_dump(dm_table_entry_t *table_en, void *data, size_t length, off_t offset)
{
	static struct dmtc_dump_helper dump_helper;
	dm_target_crypt_config_t *priv;
	static int first_call = 1;

	priv = table_en->target_config;

	if (first_call) {
		first_call = 0;
		dump_reactivate_cpus();
	}

	/* Magically enable tsleep */
	tsleep_crypto_dump = 1;

	/*
	 * 0 length means flush buffers and return
	 */
	if (length == 0) {
		if (priv->pdev->pdev_vnode->v_rdev == NULL) {
			tsleep_crypto_dump = 0;
			return ENXIO;
		}
		dev_ddump(priv->pdev->pdev_vnode->v_rdev,
		    data, 0, offset, 0);
		tsleep_crypto_dump = 0;
		return 0;
	}

	bzero(&dump_helper, sizeof(dump_helper));
	dump_helper.priv = priv;
	dump_helper.data = data;
	dump_helper.length = length;
	dump_helper.offset = offset +
	    priv->block_offset * DEV_BSIZE;
	dmtc_crypto_dump(priv, &dump_helper);

	dump_helper.offset = dm_pdev_correct_dump_offset(priv->pdev,
	    dump_helper.offset);

	dev_ddump(priv->pdev->pdev_vnode->v_rdev,
	    dump_helper.space, 0, dump_helper.offset,
	    dump_helper.length);

	tsleep_crypto_dump = 0;
	return 0;
}

static void
dmtc_crypto_dump(dm_target_crypt_config_t *priv, struct dmtc_dump_helper *dump_helper)
{
	int bytes = dump_helper->length;

	KKASSERT(dump_helper->length <= 65536);

	memcpy(dump_helper->space, dump_helper->data, bytes);

	int error = dmtc_bio_encdec(priv, dump_helper->space, bytes, dump_helper->offset,
			CRYPTOAPI_CIPHER_ENCRYPT);
	if (error != 0) {
		kprintf("dm_target_crypt: dmtc_crypto_dump = %d\n",
		error);
	}
}

static int
dmtc_mod_handler(module_t mod, int type, void *unused)
{
	dm_target_t *dmt = NULL;
	int err = 0;

	switch (type) {
	case MOD_LOAD:
		if ((dmt = dm_target_lookup("crypt")) != NULL) {
			kprintf("dm_target_crypt: Error: crypt target already defined\n");
			dm_target_unbusy(dmt);
			return EEXIST;
		}
		dmt = dm_target_alloc("crypt");
		dmt->version[0] = 1;
		dmt->version[1] = 6;
		dmt->version[2] = 0;
		dmt->init = &dm_target_crypt_init;
		dmt->destroy = &dm_target_crypt_destroy;
		dmt->strategy = &dm_target_crypt_strategy;
		dmt->table = &dm_target_crypt_table;
		dmt->dump = &dm_target_crypt_dump;

		err = dm_target_insert(dmt);
		if (!err)
			kprintf("dm_target_crypt: Successfully initialized\n");
		break;

	case MOD_UNLOAD:
		err = dm_target_remove("crypt");
		if (err == 0) {
			kprintf("dm_target_crypt: unloaded\n");
		}
		break;
	}

	return err;
}

DM_TARGET_MODULE(dm_target_crypt, dmtc_mod_handler);
MODULE_DEPEND(dm_target_crypt, crypto, 1, 1, 1);