// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright 2024 Google LLC
 */

#include <linux/blk-crypto.h>
#include <linux/ctype.h>
#include <linux/device-mapper.h>
#include <linux/hex.h>
#include <linux/module.h>
#include <keys/user-type.h>

#define DM_MSG_PREFIX	"inlinecrypt"

static const struct dm_inlinecrypt_cipher {
	const char *name;
	enum blk_crypto_mode_num mode_num;
} dm_inlinecrypt_ciphers[] = {
	{
		.name = "aes-xts-plain64",
		.mode_num = BLK_ENCRYPTION_MODE_AES_256_XTS,
	},
};

/**
 * struct inlinecrypt_ctx - private data of an inlinecrypt target
 * @dev: the underlying device
 * @start: starting sector of the range of @dev which this target actually maps.
 *	   For this purpose a "sector" is 512 bytes.
 * @cipher_string: the name of the encryption algorithm being used
 * @key_size: size of the encryption key in bytes
 * @iv_offset: starting offset for IVs.  IVs are generated as if the target were
 *	       preceded by @iv_offset 512-byte sectors.
 * @sector_size: crypto sector size in bytes (usually 4096)
 * @sector_bits: log2(sector_size)
 * @key_type: type of the key -- either raw or hardware-wrapped
 * @key: the encryption key to use
 * @max_dun: the maximum DUN that may be used (computed from other params)
 */
struct inlinecrypt_ctx {
	struct dm_dev *dev;
	sector_t start;
	const char *cipher_string;
	unsigned int key_size;
	u64 iv_offset;
	unsigned int sector_size;
	unsigned int sector_bits;
	enum blk_crypto_key_type key_type;
	struct blk_crypto_key key;
	u64 max_dun;
};

static const struct dm_inlinecrypt_cipher *
lookup_cipher(const char *cipher_string)
{
	int i;

	for (i = 0; i < ARRAY_SIZE(dm_inlinecrypt_ciphers); i++) {
		if (strcmp(cipher_string, dm_inlinecrypt_ciphers[i].name) == 0)
			return &dm_inlinecrypt_ciphers[i];
	}
	return NULL;
}

static void inlinecrypt_dtr(struct dm_target *ti)
{
	struct inlinecrypt_ctx *ctx = ti->private;

	if (ctx->dev) {
		if (ctx->key.size)
			blk_crypto_evict_key(ctx->dev->bdev, &ctx->key);
		dm_put_device(ti, ctx->dev);
	}
	kfree_sensitive(ctx->cipher_string);
	kfree_sensitive(ctx);
}

#ifdef CONFIG_KEYS

static bool contains_whitespace(const char *str)
{
	while (*str)
		if (isspace(*str++))
			return true;
	return false;
}

static int set_key_user(struct key *key, char *key_bytes,
			const unsigned int key_bytes_size)
{
	const struct user_key_payload *ukp;

	ukp = user_key_payload_locked(key);
	if (!ukp)
		return -EKEYREVOKED;

	if (key_bytes_size != ukp->datalen)
		return -EINVAL;

	memcpy(key_bytes, ukp->data, key_bytes_size);

	return 0;
}

static int inlinecrypt_get_keyring_key(const char *key_string, u8 *key_bytes,
					const unsigned int key_bytes_size)
{
	char *key_desc;
	int ret;
	struct key_type *type;
	struct key *key;
	int (*set_key)(struct key *key, char *key_bytes,
				   const unsigned int key_bytes_size);

	/*
	 * Reject key_string with whitespace. dm core currently lacks code for
	 * proper whitespace escaping in arguments on DM_TABLE_STATUS path.
	 */
	if (contains_whitespace(key_string)) {
		DMERR("whitespace chars not allowed in key string");
		return -EINVAL;
	}

	/* look for next ':' separating key_type from key_description */
	key_desc = strchr(key_string, ':');
	if (!key_desc || key_desc == key_string || !strlen(key_desc + 1))
		return -EINVAL;

	if (!strncmp(key_string, "logon:", key_desc - key_string + 1)) {
		type = &key_type_logon;
		set_key = set_key_user;
	} else {
		return -EINVAL;
	}

	key = request_key(type, key_desc + 1, NULL);
	if (IS_ERR(key))
		return PTR_ERR(key);

	down_read(&key->sem);

	ret = set_key(key, (char *)key_bytes, key_bytes_size);

	up_read(&key->sem);
	key_put(key);

	return ret;
}

static int get_key_size(char **key_string)
{
	char *colon, dummy;
	int ret;

	if (*key_string[0] != ':') {
		ret = strlen(*key_string);

		if (ret > 2 * BLK_CRYPTO_MAX_ANY_KEY_SIZE
			|| ret  % 2
			|| !ret) {
			DMERR("Invalid keysize");
			return -EINVAL;
		}
		return ret >> 1;
	}

	/* look for next ':' in key string */
	colon = strpbrk(*key_string + 1, ":");
	if (!colon)
		return -EINVAL;

	if (sscanf(*key_string + 1, "%u%c", &ret, &dummy) != 2 || dummy != ':')
		return -EINVAL;

	/* remaining key string should be :<logon|user>:<key_desc> */
	*key_string = colon;

	return ret;
}

#else

static int inlinecrypt_get_keyring_key(const char *key_string, u8 *key_bytes,
					const unsigned int key_bytes_size)
{
	return -EINVAL;
}

static int get_key_size(char **key_string)
{
	int key_hex_size = strlen(*key_string);

	if (*key_string[0] == ':')
		return -EINVAL;

	if (key_hex_size > 2 * BLK_CRYPTO_MAX_ANY_KEY_SIZE
		|| key_hex_size  % 2
		|| !key_hex_size) {
		DMERR("Invalid keysize");
		return -EINVAL;
	}

	return key_hex_size >> 1;
}

#endif /* CONFIG_KEYS */

static int inlinecrypt_get_key(const char *key_string,
				u8 key[BLK_CRYPTO_MAX_ANY_KEY_SIZE],
				const unsigned int key_size)
{
	int ret = 0;

	if (key_size > BLK_CRYPTO_MAX_ANY_KEY_SIZE) {
		DMERR("Invalid keysize");
		return -EINVAL;
	}

	/* ':' means the key is in kernel keyring, short-circuit normal key processing */
	if (key_string[0] == ':') {
		/* key string should be :<logon|user>:<key_desc> */
		ret = inlinecrypt_get_keyring_key(key_string + 1, key, key_size);
		goto out;
	}

	if (hex2bin(key, key_string, key_size) != 0)
		ret = -EINVAL;

out:
	return ret;
}

static int inlinecrypt_ctr_optional(struct dm_target *ti,
				    unsigned int argc, char **argv)
{
	struct inlinecrypt_ctx *ctx = ti->private;
	struct dm_arg_set as;
	static const struct dm_arg _args[] = {
		{0, 4, "Invalid number of feature args"},
	};
	unsigned int opt_params;
	const char *opt_string;
	bool iv_large_sectors = false;
	char dummy;
	int err;

	as.argc = argc;
	as.argv = argv;

	err = dm_read_arg_group(_args, &as, &opt_params, &ti->error);
	if (err)
		return err;

	while (opt_params--) {
		opt_string = dm_shift_arg(&as);
		if (!opt_string) {
			ti->error = "Not enough feature arguments";
			return -EINVAL;
		}
		if (str_has_prefix(opt_string, "keytype:")) {
			const char *val = opt_string + strlen("keytype:");

			if (!*val) {
				ti->error = "Invalid block key type";
				return -EINVAL;
			}

			if (!strcmp(val, "raw")) {
				ctx->key_type = BLK_CRYPTO_KEY_TYPE_RAW;
			} else if (!strcmp(val, "hw-wrapped")) {
				ctx->key_type = BLK_CRYPTO_KEY_TYPE_HW_WRAPPED;
			} else {
				ti->error = "Invalid block key type";
				return -EINVAL;
			}
		} else if (!strcmp(opt_string, "allow_discards")) {
			ti->num_discard_bios = 1;
		} else if (sscanf(opt_string, "sector_size:%u%c",
				  &ctx->sector_size, &dummy) == 1) {
			if (ctx->sector_size < SECTOR_SIZE ||
			    ctx->sector_size > 4096 ||
			    !is_power_of_2(ctx->sector_size)) {
				ti->error = "Invalid sector_size";
				return -EINVAL;
			}
		} else if (!strcmp(opt_string, "iv_large_sectors")) {
			iv_large_sectors = true;
		} else {
			ti->error = "Invalid feature arguments";
			return -EINVAL;
		}
	}

	/* dm-inlinecrypt doesn't implement iv_large_sectors=false. */
	if (ctx->sector_size != SECTOR_SIZE && !iv_large_sectors) {
		ti->error = "iv_large_sectors must be specified";
		return -EINVAL;
	}

	return 0;
}

/*
 * Construct an inlinecrypt mapping:
 * <cipher> [<key>|:<key_size>:<logon>:<key_description>] <iv_offset> <dev_path> <start>
 *
 * This syntax matches dm-crypt's, but the set of supported functionality has
 * been stripped down.
 */
static int inlinecrypt_ctr(struct dm_target *ti, unsigned int argc, char **argv)
{
	struct inlinecrypt_ctx *ctx;
	const struct dm_inlinecrypt_cipher *cipher;
	u8 key_bytes[BLK_CRYPTO_MAX_ANY_KEY_SIZE];
	unsigned int dun_bytes;
	unsigned long long tmpll;
	char dummy;
	int err;

	if (argc < 5) {
		ti->error = "Not enough arguments";
		return -EINVAL;
	}

	ctx = kzalloc_obj(*ctx);
	if (!ctx) {
		ti->error = "Out of memory";
		return -ENOMEM;
	}
	ti->private = ctx;

	/* <cipher> */
	ctx->cipher_string = kstrdup(argv[0], GFP_KERNEL);
	if (!ctx->cipher_string) {
		ti->error = "Out of memory";
		err = -ENOMEM;
		goto bad;
	}
	cipher = lookup_cipher(ctx->cipher_string);
	if (!cipher) {
		ti->error = "Unsupported cipher";
		err = -EINVAL;
		goto bad;
	}

	/* <key> */
	err = get_key_size(&argv[1]);
	if (err < 0) {
		ti->error = "Cannot parse key size";
		goto bad;
	}
	ctx->key_size = err;

	err = inlinecrypt_get_key(argv[1], key_bytes, ctx->key_size);
	if (err) {
		ti->error = "Malformed key string";
		goto bad;
	}

	/* <iv_offset> */
	if (sscanf(argv[2], "%llu%c", &ctx->iv_offset, &dummy) != 1) {
		ti->error = "Invalid iv_offset sector";
		err = -EINVAL;
		goto bad;
	}

	/* <dev_path> */
	err = dm_get_device(ti, argv[3], dm_table_get_mode(ti->table),
			    &ctx->dev);
	if (err) {
		ti->error = "Device lookup failed";
		goto bad;
	}

	/* <start> */
	if (sscanf(argv[4], "%llu%c", &tmpll, &dummy) != 1 ||
	    tmpll != (sector_t)tmpll) {
		ti->error = "Invalid start sector";
		err = -EINVAL;
		goto bad;
	}
	ctx->start = tmpll;

	/* optional arguments */
	ctx->sector_size = SECTOR_SIZE;
	ctx->key_type = BLK_CRYPTO_KEY_TYPE_RAW;
	if (argc > 5) {
		err = inlinecrypt_ctr_optional(ti, argc - 5, &argv[5]);
		if (err)
			goto bad;
	}
	ctx->sector_bits = ilog2(ctx->sector_size);
	if (ti->len & ((ctx->sector_size >> SECTOR_SHIFT) - 1)) {
		ti->error = "Device size is not a multiple of sector_size";
		err = -EINVAL;
		goto bad;
	}
	if (ctx->iv_offset & ((ctx->sector_size >> SECTOR_SHIFT) - 1)) {
		ti->error = "Wrong alignment of iv_offset sector";
		err = -EINVAL;
		goto bad;
	}

	ctx->max_dun = (ctx->iv_offset + ti->len - 1) >>
		       (ctx->sector_bits - SECTOR_SHIFT);
	dun_bytes = DIV_ROUND_UP(fls64(ctx->max_dun), 8);

	err = blk_crypto_init_key(&ctx->key, key_bytes, ctx->key_size,
				  ctx->key_type, cipher->mode_num,
				  dun_bytes, ctx->sector_size,
				  BLK_CRYPTO_CFG_ALLOW_HW);
	if (err) {
		ti->error = "Error initializing blk-crypto key";
		goto bad;
	}

	err = blk_crypto_start_using_key(ctx->dev->bdev, &ctx->key);
	if (err) {
		ti->error = "Error starting to use blk-crypto";
		goto bad;
	}

	ti->num_flush_bios = 1;

	err = 0;
	goto out;

bad:
	inlinecrypt_dtr(ti);
out:
	memzero_explicit(key_bytes, sizeof(key_bytes));
	return err;
}

static int inlinecrypt_map(struct dm_target *ti, struct bio *bio)
{
	const struct inlinecrypt_ctx *ctx = ti->private;
	sector_t sector_in_target;
	u64 dun[BLK_CRYPTO_DUN_ARRAY_SIZE] = {};

	bio_set_dev(bio, ctx->dev->bdev);

	/*
	 * If the bio is a device-level request which doesn't target a specific
	 * sector, there's nothing more to do.
	 */
	if (bio_sectors(bio) == 0)
		return DM_MAPIO_REMAPPED;

	/*
	 * The bio should never have an encryption context already, since
	 * dm-inlinecrypt doesn't pass through any inline encryption
	 * capabilities to the layer above it.
	 */
	if (WARN_ON_ONCE(bio_has_crypt_ctx(bio)))
		return DM_MAPIO_KILL;

	/* Map the bio's sector to the underlying device. (512-byte sectors) */
	sector_in_target = dm_target_offset(ti, bio->bi_iter.bi_sector);
	bio->bi_iter.bi_sector = ctx->start + sector_in_target;
	/*
	 * If the bio doesn't have any data (e.g. if it's a DISCARD request),
	 * there's nothing more to do.
	 */
	if (!bio_has_data(bio))
		return DM_MAPIO_REMAPPED;

	/* Calculate the DUN and enforce data-unit (crypto sector) alignment. */
	dun[0] = ctx->iv_offset + sector_in_target; /* 512-byte sectors */
	if (dun[0] & ((ctx->sector_size >> SECTOR_SHIFT) - 1))
		return DM_MAPIO_KILL;
	dun[0] >>= ctx->sector_bits - SECTOR_SHIFT; /* crypto sectors */

	/*
	 * This check isn't necessary as we should have calculated max_dun
	 * correctly, but be safe.
	 */
	if (WARN_ON_ONCE(dun[0] > ctx->max_dun))
		return DM_MAPIO_KILL;

	bio_crypt_set_ctx(bio, &ctx->key, dun, GFP_NOIO);

	/*
	 * Since we've added an encryption context to the bio and
	 * blk-crypto-fallback may be needed to process it, it's necessary to
	 * use the fallback-aware bio submission code rather than
	 * unconditionally returning DM_MAPIO_REMAPPED.
	 *
	 * To get the correct accounting for a dm target in the case where
	 * __blk_crypto_submit_bio() doesn't take ownership of the bio (returns
	 * true), call __blk_crypto_submit_bio() directly and return
	 * DM_MAPIO_REMAPPED in that case, rather than relying on
	 * blk_crypto_submit_bio() which calls submit_bio() in that case.
	 *
	 * TODO: blk-crypto fallback write slow-path currently double-accounts
	 * IO in vmstat, as encrypted bios are submitted via submit_bio().
	 * This does not affect data correctness. Consider fixing this if
	 * a cleaner accounting model for derived bios is introduced.
	 */
	if (__blk_crypto_submit_bio(bio))
		return DM_MAPIO_REMAPPED;
	return DM_MAPIO_SUBMITTED;
}

static void inlinecrypt_status(struct dm_target *ti, status_type_t type,
			       unsigned int status_flags, char *result,
			       unsigned int maxlen)
{
	const struct inlinecrypt_ctx *ctx = ti->private;
	unsigned int sz = 0;
	int num_feature_args = 0;

	switch (type) {
	case STATUSTYPE_INFO:
	case STATUSTYPE_IMA:
		result[0] = '\0';
		break;

	case STATUSTYPE_TABLE:
		/*
		 * Warning: like dm-crypt, dm-inlinecrypt includes the key in
		 * the returned table.  Userspace is responsible for redacting
		 * the key when needed.
		 */
		DMEMIT("%s %*phN %u %llu %s %llu", ctx->cipher_string,
		       ctx->key.size, ctx->key.bytes,
		       ctx->key_type, ctx->iv_offset,
		       ctx->dev->name, ctx->start);
		num_feature_args += !!ti->num_discard_bios;
		if (ctx->sector_size != SECTOR_SIZE)
			num_feature_args += 2;
		if (num_feature_args != 0) {
			DMEMIT(" %d", num_feature_args);
			if (ti->num_discard_bios)
				DMEMIT(" allow_discards");
			if (ctx->sector_size != SECTOR_SIZE) {
				DMEMIT(" sector_size:%u", ctx->sector_size);
				DMEMIT(" iv_large_sectors");
			}
		}
		break;
	}
}

static int inlinecrypt_prepare_ioctl(struct dm_target *ti,
				     struct block_device **bdev, unsigned int cmd,
				     unsigned long arg, bool *forward)
{
	const struct inlinecrypt_ctx *ctx = ti->private;
	const struct dm_dev *dev = ctx->dev;

	*bdev = dev->bdev;

	/* Only pass ioctls through if the device sizes match exactly. */
	return ctx->start != 0 || ti->len != bdev_nr_sectors(dev->bdev);
}

static int inlinecrypt_iterate_devices(struct dm_target *ti,
				       iterate_devices_callout_fn fn,
				       void *data)
{
	const struct inlinecrypt_ctx *ctx = ti->private;

	return fn(ti, ctx->dev, ctx->start, ti->len, data);
}

#ifdef CONFIG_BLK_DEV_ZONED
static int inlinecrypt_report_zones(struct dm_target *ti,
				    struct dm_report_zones_args *args,
				    unsigned int nr_zones)
{
	const struct inlinecrypt_ctx *ctx = ti->private;

	return dm_report_zones(ctx->dev->bdev, ctx->start,
			ctx->start + dm_target_offset(ti, args->next_sector),
			args, nr_zones);
}
#else
#define inlinecrypt_report_zones NULL
#endif

static void inlinecrypt_io_hints(struct dm_target *ti,
				 struct queue_limits *limits)
{
	const struct inlinecrypt_ctx *ctx = ti->private;
	const unsigned int sector_size = ctx->sector_size;

	limits->logical_block_size =
		max_t(unsigned int, limits->logical_block_size, sector_size);
	limits->physical_block_size =
		max_t(unsigned int, limits->physical_block_size, sector_size);
	limits->io_min = max_t(unsigned int, limits->io_min, sector_size);
	limits->dma_alignment = limits->logical_block_size - 1;
}

static struct target_type inlinecrypt_target = {
	.name			= "inlinecrypt",
	.version		= {1, 0, 0},
	/*
	 * Do not set DM_TARGET_PASSES_CRYPTO, since dm-inlinecrypt consumes the
	 * crypto capability itself.
	 */
	.features		= DM_TARGET_ZONED_HM,
	.module			= THIS_MODULE,
	.ctr			= inlinecrypt_ctr,
	.dtr			= inlinecrypt_dtr,
	.map			= inlinecrypt_map,
	.status			= inlinecrypt_status,
	.prepare_ioctl		= inlinecrypt_prepare_ioctl,
	.iterate_devices	= inlinecrypt_iterate_devices,
	.report_zones		= inlinecrypt_report_zones,
	.io_hints		= inlinecrypt_io_hints,
};

module_dm(inlinecrypt);

MODULE_AUTHOR("Eric Biggers <ebiggers@google.com>");
MODULE_AUTHOR("Linlin Zhang <linlin.zhang@oss.qualcomm.com>");
MODULE_DESCRIPTION(DM_NAME " target for inline encryption");
MODULE_LICENSE("GPL");