// SPDX-License-Identifier: GPL-2.0
/*
 * UFS OP-TEE based RPMB Driver
 *
 * Copyright (C) 2025 Micron Technology, Inc.
 * Copyright (C) 2025 Qualcomm Technologies, Inc.
 *
 * Authors:
 *	Bean Huo <beanhuo@micron.com>
 *	Can Guo <can.guo@oss.qualcomm.com>
 */

#include <linux/module.h>
#include <linux/device.h>
#include <linux/kernel.h>
#include <linux/types.h>
#include <linux/rpmb.h>
#include <linux/string.h>
#include <linux/list.h>
#include <ufs/ufshcd.h>
#include <linux/unaligned.h>
#include "ufshcd-priv.h"

#define UFS_RPMB_SEC_PROTOCOL		0xEC	/* JEDEC UFS application */
#define UFS_RPMB_SEC_PROTOCOL_ID	0x01	/* JEDEC UFS RPMB protocol ID, CDB byte3 */

static const struct bus_type ufs_rpmb_bus_type = {
	.name = "ufs_rpmb",
};

/* UFS RPMB device structure */
struct ufs_rpmb_dev {
	u8 region_id;
	struct device dev;
	struct rpmb_dev *rdev;
	struct ufs_hba *hba;
	struct list_head node;
};

static int ufs_sec_submit(struct ufs_hba *hba, u16 spsp, void *buffer, size_t len, bool send)
{
	struct scsi_device *sdev = hba->ufs_rpmb_wlun;
	u8 cdb[12] = { };

	cdb[0] = send ? SECURITY_PROTOCOL_OUT : SECURITY_PROTOCOL_IN;
	cdb[1] = UFS_RPMB_SEC_PROTOCOL;
	put_unaligned_be16(spsp, &cdb[2]);
	put_unaligned_be32(len, &cdb[6]);

	return scsi_execute_cmd(sdev, cdb, send ? REQ_OP_DRV_OUT : REQ_OP_DRV_IN,
				buffer, len, /*timeout=*/30 * HZ, 0, NULL);
}

/* UFS RPMB route frames implementation */
static int ufs_rpmb_route_frames(struct device *dev, u8 *req, unsigned int req_len, u8 *resp,
					unsigned int resp_len)
{
	struct ufs_rpmb_dev *ufs_rpmb = dev_get_drvdata(dev);
	struct rpmb_frame *frm_out = (struct rpmb_frame *)req;
	bool need_result_read = true;
	u16 req_type, protocol_id;
	struct ufs_hba *hba;
	int ret;

	if (!ufs_rpmb) {
		dev_err(dev, "Missing driver data\n");
		return -ENODEV;
	}

	hba = ufs_rpmb->hba;

	req_type = be16_to_cpu(frm_out->req_resp);

	switch (req_type) {
	case RPMB_PROGRAM_KEY:
		if (req_len != sizeof(struct rpmb_frame) || resp_len != sizeof(struct rpmb_frame))
			return -EINVAL;
		break;
	case RPMB_GET_WRITE_COUNTER:
		if (req_len != sizeof(struct rpmb_frame) || resp_len != sizeof(struct rpmb_frame))
			return -EINVAL;
		need_result_read = false;
		break;
	case RPMB_WRITE_DATA:
		if (req_len % sizeof(struct rpmb_frame) || resp_len != sizeof(struct rpmb_frame))
			return -EINVAL;
		break;
	case RPMB_READ_DATA:
		if (req_len != sizeof(struct rpmb_frame) || resp_len % sizeof(struct rpmb_frame))
			return -EINVAL;
		need_result_read = false;
		break;
	default:
		dev_err(dev, "Unknown request type=0x%04x\n", req_type);
		return -EINVAL;
	}

	protocol_id = ufs_rpmb->region_id << 8 | UFS_RPMB_SEC_PROTOCOL_ID;

	ret = ufs_sec_submit(hba, protocol_id, req, req_len, true);
	if (ret) {
		dev_err(dev, "Command failed with ret=%d\n", ret);
		return ret;
	}

	if (need_result_read) {
		struct rpmb_frame *frm_resp = (struct rpmb_frame *)resp;

		memset(frm_resp, 0, sizeof(*frm_resp));
		frm_resp->req_resp = cpu_to_be16(RPMB_RESULT_READ);
		ret = ufs_sec_submit(hba, protocol_id, resp, resp_len, true);
		if (ret) {
			dev_err(dev, "Result read request failed with ret=%d\n", ret);
			return ret;
		}
	}

	if (!ret) {
		ret = ufs_sec_submit(hba, protocol_id, resp, resp_len, false);
		if (ret)
			dev_err(dev, "Response read failed with ret=%d\n", ret);
	}

	return ret;
}

static void ufs_rpmb_device_release(struct device *dev)
{
	struct ufs_rpmb_dev *ufs_rpmb = dev_get_drvdata(dev);

	rpmb_dev_unregister(ufs_rpmb->rdev);
}

/* UFS RPMB device registration */
int ufs_rpmb_probe(struct ufs_hba *hba)
{
	struct ufs_rpmb_dev *ufs_rpmb, *it, *tmp;
	struct rpmb_dev *rdev;
	char *cid = NULL;
	int region;
	u32 cap;
	int ret;

	if (!hba->ufs_rpmb_wlun || hba->dev_info.b_advanced_rpmb_en) {
		dev_info(hba->dev, "Skip OP-TEE RPMB registration\n");
		return -ENODEV;
	}

	/* Check if device_id is available */
	if (!hba->dev_info.device_id) {
		dev_err(hba->dev, "UFS Device ID not available\n");
		return -EINVAL;
	}

	INIT_LIST_HEAD(&hba->rpmbs);

	struct rpmb_descr descr = {
		.type = RPMB_TYPE_UFS,
		.route_frames = ufs_rpmb_route_frames,
		.reliable_wr_count = hba->dev_info.rpmb_io_size,
	};

	for (region = 0; region < ARRAY_SIZE(hba->dev_info.rpmb_region_size); region++) {
		cap = hba->dev_info.rpmb_region_size[region];
		if (!cap)
			continue;

		ufs_rpmb = devm_kzalloc(hba->dev, sizeof(*ufs_rpmb), GFP_KERNEL);
		if (!ufs_rpmb) {
			ret = -ENOMEM;
			goto err_out;
		}

		ufs_rpmb->hba = hba;
		ufs_rpmb->dev.parent = &hba->ufs_rpmb_wlun->sdev_gendev;
		ufs_rpmb->dev.bus = &ufs_rpmb_bus_type;
		ufs_rpmb->dev.release = ufs_rpmb_device_release;
		dev_set_name(&ufs_rpmb->dev, "ufs_rpmb%d", region);

		/* Set driver data BEFORE device_register */
		dev_set_drvdata(&ufs_rpmb->dev, ufs_rpmb);

		ret = device_register(&ufs_rpmb->dev);
		if (ret) {
			dev_err(hba->dev, "Failed to register UFS RPMB device %d\n", region);
			put_device(&ufs_rpmb->dev);
			goto err_out;
		}

		/* Create unique ID by appending region number to device_id */
		cid = kasprintf(GFP_KERNEL, "%s-R%d", hba->dev_info.device_id, region);
		if (!cid) {
			device_unregister(&ufs_rpmb->dev);
			ret = -ENOMEM;
			goto err_out;
		}

		descr.dev_id = cid;
		descr.dev_id_len = strlen(cid);
		descr.capacity = cap;

		/* Register RPMB device */
		rdev = rpmb_dev_register(&ufs_rpmb->dev, &descr);
		if (IS_ERR(rdev)) {
			dev_err(hba->dev, "Failed to register UFS RPMB device.\n");
			device_unregister(&ufs_rpmb->dev);
			ret = PTR_ERR(rdev);
			goto err_out;
		}

		kfree(cid);
		cid = NULL;

		ufs_rpmb->rdev = rdev;
		ufs_rpmb->region_id = region;

		list_add_tail(&ufs_rpmb->node, &hba->rpmbs);

		dev_info(hba->dev, "UFS RPMB region %d registered (capacity=%u)\n", region, cap);
	}

	return 0;
err_out:
	kfree(cid);
	list_for_each_entry_safe(it, tmp, &hba->rpmbs, node) {
		list_del(&it->node);
		device_unregister(&it->dev);
	}

	return ret;
}

/* UFS RPMB remove handler */
void ufs_rpmb_remove(struct ufs_hba *hba)
{
	struct ufs_rpmb_dev *ufs_rpmb, *tmp;

	if (list_empty(&hba->rpmbs))
		return;

	/* Remove all registered RPMB devices */
	list_for_each_entry_safe(ufs_rpmb, tmp, &hba->rpmbs, node) {
		dev_info(hba->dev, "Removing UFS RPMB region %d\n", ufs_rpmb->region_id);
		/* Remove from list first */
		list_del(&ufs_rpmb->node);
		/* Unregister device */
		device_unregister(&ufs_rpmb->dev);
	}

	dev_info(hba->dev, "All UFS RPMB devices unregistered\n");
}

MODULE_LICENSE("GPL v2");
MODULE_DESCRIPTION("OP-TEE UFS RPMB driver");