// SPDX-License-Identifier: GPL-2.0
/*
 * Copyright (C) 2022 Microchip.
 */

#include <linux/device.h>
#include <linux/kernel.h>
#include <linux/kthread.h>
#include <linux/module.h>
#include <linux/rtc.h>
#include <linux/tee_drv.h>

#define RTC_INFO_VERSION			0x1

#define TA_RTC_FEATURE_CORRECTION		BIT(0)
#define TA_RTC_FEATURE_ALARM			BIT(1)
#define TA_RTC_FEATURE_WAKEUP_ALARM		BIT(2)

enum rtc_optee_pta_cmd {
	/*  PTA_CMD_RTC_GET_INFO - Get RTC information
	 *
	 * [out]        memref[0]  RTC buffer memory reference containing a struct pta_rtc_info
	 */
	PTA_CMD_RTC_GET_INF = 0x0,

	/*
	 * PTA_CMD_RTC_GET_TIME - Get time from RTC
	 *
	 * [out]    memref[0]  RTC buffer memory reference containing a struct pta_rtc_time
	 */
	PTA_CMD_RTC_GET_TIME = 0x1,

	/*
	 * PTA_CMD_RTC_SET_TIME - Set time from RTC
	 *
	 * [in]     memref[0]  RTC buffer memory reference containing a struct pta_rtc_time to be
	 *                     used as RTC time
	 */
	PTA_CMD_RTC_SET_TIME = 0x2,

	/*
	 * PTA_CMD_RTC_GET_OFFSET - Get RTC offset
	 *
	 * [out]    value[0].a  RTC offset (signed 32bit value)
	 */
	PTA_CMD_RTC_GET_OFFSET = 0x3,

	/*
	 * PTA_CMD_RTC_SET_OFFSET - Set RTC offset
	 *
	 * [in]     value[0].a  RTC offset to be set (signed 32bit value)
	 */
	PTA_CMD_RTC_SET_OFFSET = 0x4,

	/*
	 * PTA_CMD_RTC_READ_ALARM - Read RTC alarm
	 *
	 * [out]     memref[0]  RTC buffer memory reference containing a struct pta_rtc_alarm
	 */
	PTA_CMD_RTC_READ_ALARM = 0x5,

	/*
	 * PTA_CMD_RTC_SET_ALARM - Set RTC alarm
	 *
	 * [in]     memref[0]  RTC buffer memory reference containing a struct pta_rtc_alarm to be
	 *                     used as RTC alarm
	 */
	PTA_CMD_RTC_SET_ALARM = 0x6,

	/*
	 * PTA_CMD_RTC_ENABLE_ALARM - Enable Alarm
	 *
	 * [in]     value[0].a  RTC IRQ flag (uint32_t), 0 to disable the alarm, 1 to enable
	 */
	PTA_CMD_RTC_ENABLE_ALARM = 0x7,

	/*
	 * PTA_CMD_RTC_WAIT_ALARM - Get alarm event
	 *
	 * [out]     value[0].a  RTC wait alarm return status (uint32_t):
	 *                       - 0: No alarm event
	 *                       - 1: Alarm event occurred
	 *                       - 2: Alarm event canceled
	 */
	PTA_CMD_RTC_WAIT_ALARM = 0x8,

	/*
	 * PTA_CMD_RTC_CANCEL_WAIT - Cancel wait for alarm event
	 */
	PTA_CMD_RTC_CANCEL_WAIT = 0x9,

	/*
	 * PTA_CMD_RTC_SET_WAKE_ALARM_STATUS - Set RTC wake alarm status flag
	 *
	 * [in]     value[0].a RTC IRQ wake alarm flag (uint32_t), 0 to disable the wake up
	 *                     capability, 1 to enable.
	 */
	PTA_CMD_RTC_SET_WAKE_ALARM_STATUS = 0xA,
};

enum rtc_wait_alarm_status {
	WAIT_ALARM_RESET = 0x0,
	WAIT_ALARM_ALARM_OCCURRED = 0x1,
	WAIT_ALARM_CANCELED = 0x2,
};

struct optee_rtc_time {
	u32 tm_sec;
	u32 tm_min;
	u32 tm_hour;
	u32 tm_mday;
	u32 tm_mon;
	u32 tm_year;
	u32 tm_wday;
};

struct optee_rtc_alarm {
	u8 enabled;
	u8 pending;
	struct optee_rtc_time time;
};

struct optee_rtc_info {
	u64 version;
	u64 features;
	struct optee_rtc_time range_min;
	struct optee_rtc_time range_max;
};

/**
 * struct optee_rtc - OP-TEE RTC private data
 * @dev:		OP-TEE based RTC device.
 * @ctx:		OP-TEE context handler.
 * @session_id:		RTC TA session identifier.
 * @session2_id:	RTC wait alarm session identifier.
 * @shm:		Memory pool shared with RTC device.
 * @features:		Bitfield of RTC features
 * @alarm_task:		RTC wait alamr task.
 * @rtc:		RTC device.
 */
struct optee_rtc {
	struct device *dev;
	struct tee_context *ctx;
	u32 session_id;
	u32 session2_id;
	struct tee_shm *shm;
	u64 features;
	struct task_struct *alarm_task;
	struct rtc_device *rtc;
};

static int optee_rtc_readtime(struct device *dev, struct rtc_time *tm)
{
	struct optee_rtc *priv = dev_get_drvdata(dev);
	struct tee_ioctl_invoke_arg inv_arg = {0};
	struct optee_rtc_time *optee_tm;
	struct tee_param param[4] = {0};
	int ret;

	inv_arg.func = PTA_CMD_RTC_GET_TIME;
	inv_arg.session = priv->session_id;
	inv_arg.num_params = 4;

	/* Fill invoke cmd params */
	param[0].attr = TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_OUTPUT;
	param[0].u.memref.shm = priv->shm;
	param[0].u.memref.size = sizeof(struct optee_rtc_time);

	ret = tee_client_invoke_func(priv->ctx, &inv_arg, param);
	if (ret < 0 || inv_arg.ret != 0)
		return ret ? ret : -EPROTO;

	optee_tm = tee_shm_get_va(priv->shm, 0);
	if (IS_ERR(optee_tm))
		return PTR_ERR(optee_tm);

	if (param[0].u.memref.size != sizeof(*optee_tm))
		return -EPROTO;

	tm->tm_sec = optee_tm->tm_sec;
	tm->tm_min = optee_tm->tm_min;
	tm->tm_hour = optee_tm->tm_hour;
	tm->tm_mday = optee_tm->tm_mday;
	tm->tm_mon = optee_tm->tm_mon;
	tm->tm_year = optee_tm->tm_year - 1900;
	tm->tm_wday = optee_tm->tm_wday;
	tm->tm_yday = rtc_year_days(tm->tm_mday, tm->tm_mon, tm->tm_year);

	return 0;
}

static int optee_rtc_settime(struct device *dev, struct rtc_time *tm)
{
	struct optee_rtc *priv = dev_get_drvdata(dev);
	struct tee_ioctl_invoke_arg inv_arg = {0};
	struct tee_param param[4] = {0};
	struct optee_rtc_time *optee_tm;
	int ret;

	inv_arg.func = PTA_CMD_RTC_SET_TIME;
	inv_arg.session = priv->session_id;
	inv_arg.num_params = 4;

	param[0].attr = TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_INPUT;
	param[0].u.memref.shm = priv->shm;
	param[0].u.memref.size = sizeof(struct optee_rtc_time);

	optee_tm = tee_shm_get_va(priv->shm, 0);
	if (IS_ERR(optee_tm))
		return PTR_ERR(optee_tm);

	optee_tm->tm_min = tm->tm_min;
	optee_tm->tm_sec = tm->tm_sec;
	optee_tm->tm_hour = tm->tm_hour;
	optee_tm->tm_mday = tm->tm_mday;
	optee_tm->tm_mon = tm->tm_mon;
	optee_tm->tm_year = tm->tm_year + 1900;
	optee_tm->tm_wday = tm->tm_wday;

	ret = tee_client_invoke_func(priv->ctx, &inv_arg, param);
	if (ret < 0 || inv_arg.ret != 0)
		return ret ? ret : -EPROTO;

	return 0;
}

static int optee_rtc_readoffset(struct device *dev, long *offset)
{
	struct optee_rtc *priv = dev_get_drvdata(dev);
	struct tee_ioctl_invoke_arg inv_arg = {0};
	struct tee_param param[4] = {0};
	int ret;

	if (!(priv->features & TA_RTC_FEATURE_CORRECTION))
		return -EOPNOTSUPP;

	inv_arg.func = PTA_CMD_RTC_GET_OFFSET;
	inv_arg.session = priv->session_id;
	inv_arg.num_params = 4;

	param[0].attr = TEE_IOCTL_PARAM_ATTR_TYPE_VALUE_OUTPUT;

	ret = tee_client_invoke_func(priv->ctx, &inv_arg, param);
	if (ret < 0 || inv_arg.ret != 0)
		return ret ? ret : -EPROTO;

	*offset = param[0].u.value.a;

	return 0;
}

static int optee_rtc_setoffset(struct device *dev, long offset)
{
	struct optee_rtc *priv = dev_get_drvdata(dev);
	struct tee_ioctl_invoke_arg inv_arg = {0};
	struct tee_param param[4] = {0};
	int ret;

	if (!(priv->features & TA_RTC_FEATURE_CORRECTION))
		return -EOPNOTSUPP;

	inv_arg.func = PTA_CMD_RTC_SET_OFFSET;
	inv_arg.session = priv->session_id;
	inv_arg.num_params = 4;

	param[0].attr = TEE_IOCTL_PARAM_ATTR_TYPE_VALUE_INPUT;
	param[0].u.value.a = offset;

	ret = tee_client_invoke_func(priv->ctx, &inv_arg, param);
	if (ret < 0 || inv_arg.ret != 0)
		return ret ? ret : -EPROTO;

	return 0;
}

static int optee_rtc_read_alarm(struct device *dev, struct rtc_wkalrm *alarm)
{
	struct optee_rtc *priv = dev_get_drvdata(dev);
	struct tee_ioctl_invoke_arg inv_arg = {0};
	struct optee_rtc_alarm *optee_alarm;
	struct tee_param param[1] = {0};
	int ret;

	if (!(priv->features & TA_RTC_FEATURE_ALARM))
		return -EOPNOTSUPP;

	inv_arg.func = PTA_CMD_RTC_READ_ALARM;
	inv_arg.session = priv->session_id;
	inv_arg.num_params = 1;

	/* Fill invoke cmd params */
	param[0].attr = TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_OUTPUT;
	param[0].u.memref.shm = priv->shm;
	param[0].u.memref.size = sizeof(struct optee_rtc_alarm);

	ret = tee_client_invoke_func(priv->ctx, &inv_arg, param);
	if (ret < 0 || inv_arg.ret != 0)
		return ret ? ret : -EPROTO;

	optee_alarm = tee_shm_get_va(priv->shm, 0);
	if (IS_ERR(optee_alarm))
		return PTR_ERR(optee_alarm);

	if (param[0].u.memref.size != sizeof(*optee_alarm))
		return -EPROTO;

	alarm->enabled = optee_alarm->enabled;
	alarm->pending = optee_alarm->pending;
	alarm->time.tm_sec = optee_alarm->time.tm_sec;
	alarm->time.tm_min = optee_alarm->time.tm_min;
	alarm->time.tm_hour = optee_alarm->time.tm_hour;
	alarm->time.tm_mday = optee_alarm->time.tm_mday;
	alarm->time.tm_mon = optee_alarm->time.tm_mon;
	alarm->time.tm_year = optee_alarm->time.tm_year - 1900;
	alarm->time.tm_wday = optee_alarm->time.tm_wday;
	alarm->time.tm_yday = rtc_year_days(alarm->time.tm_mday,
					    alarm->time.tm_mon,
					    alarm->time.tm_year);

	return 0;
}

static int optee_rtc_set_alarm(struct device *dev, struct rtc_wkalrm *alarm)
{
	struct optee_rtc *priv = dev_get_drvdata(dev);
	struct tee_ioctl_invoke_arg inv_arg = {0};
	struct optee_rtc_alarm *optee_alarm;
	struct tee_param param[1] = {0};
	int ret;

	if (!(priv->features & TA_RTC_FEATURE_ALARM))
		return -EOPNOTSUPP;

	inv_arg.func = PTA_CMD_RTC_SET_ALARM;
	inv_arg.session = priv->session_id;
	inv_arg.num_params = 1;

	param[0].attr = TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_INPUT;
	param[0].u.memref.shm = priv->shm;
	param[0].u.memref.size = sizeof(struct optee_rtc_alarm);

	optee_alarm = tee_shm_get_va(priv->shm, 0);
	if (IS_ERR(optee_alarm))
		return PTR_ERR(optee_alarm);

	optee_alarm->enabled = alarm->enabled;
	optee_alarm->pending = alarm->pending;
	optee_alarm->time.tm_sec = alarm->time.tm_sec;
	optee_alarm->time.tm_min = alarm->time.tm_min;
	optee_alarm->time.tm_hour = alarm->time.tm_hour;
	optee_alarm->time.tm_mday = alarm->time.tm_mday;
	optee_alarm->time.tm_mon = alarm->time.tm_mon;
	optee_alarm->time.tm_year = alarm->time.tm_year + 1900;
	optee_alarm->time.tm_wday = alarm->time.tm_wday;

	ret = tee_client_invoke_func(priv->ctx, &inv_arg, param);
	if (ret < 0 || inv_arg.ret != 0)
		return ret ? ret : -EPROTO;

	return 0;
}

static int optee_rtc_enable_alarm(struct device *dev, unsigned int enabled)
{
	struct optee_rtc *priv = dev_get_drvdata(dev);
	struct tee_ioctl_invoke_arg inv_arg = {0};
	struct tee_param param[1] = {0};
	int ret;

	if (!(priv->features & TA_RTC_FEATURE_ALARM))
		return -EOPNOTSUPP;

	inv_arg.func = PTA_CMD_RTC_ENABLE_ALARM;
	inv_arg.session = priv->session_id;
	inv_arg.num_params = 1;

	param[0].attr = TEE_IOCTL_PARAM_ATTR_TYPE_VALUE_INPUT;
	param[0].u.value.a = (bool)enabled;

	ret = tee_client_invoke_func(priv->ctx, &inv_arg, param);
	if (ret < 0 || inv_arg.ret != 0)
		return ret ? ret : -EPROTO;

	return 0;
}

static const struct rtc_class_ops optee_rtc_ops = {
	.read_time		= optee_rtc_readtime,
	.set_time		= optee_rtc_settime,
	.set_offset		= optee_rtc_setoffset,
	.read_offset		= optee_rtc_readoffset,
	.read_alarm		= optee_rtc_read_alarm,
	.set_alarm		= optee_rtc_set_alarm,
	.alarm_irq_enable	= optee_rtc_enable_alarm,
};

static int optee_rtc_wait_alarm(struct device *dev, int *return_status)
{
	struct optee_rtc *priv = dev_get_drvdata(dev);
	struct tee_ioctl_invoke_arg inv_arg = {0};
	struct tee_param param[1] = {0};
	int ret;

	if (!(priv->features & TA_RTC_FEATURE_ALARM))
		return -EOPNOTSUPP;

	inv_arg.func = PTA_CMD_RTC_WAIT_ALARM;
	inv_arg.session = priv->session2_id;
	inv_arg.num_params = 1;

	param[0].attr = TEE_IOCTL_PARAM_ATTR_TYPE_VALUE_OUTPUT;

	ret = tee_client_invoke_func(priv->ctx, &inv_arg, param);
	if (ret < 0 || inv_arg.ret != 0)
		return ret ? ret : -EPROTO;

	*return_status = param[0].u.value.a;

	return 0;
}

static int optee_rtc_cancel_wait_alarm(struct device *dev)
{
	struct optee_rtc *priv = dev_get_drvdata(dev);
	struct tee_ioctl_invoke_arg inv_arg = {0};
	struct tee_param param[1] = {0};
	int ret;

	if (!(priv->features & TA_RTC_FEATURE_ALARM))
		return -EOPNOTSUPP;

	inv_arg.func = PTA_CMD_RTC_CANCEL_WAIT;
	inv_arg.session = priv->session_id;
	inv_arg.num_params = 0;

	ret = tee_client_invoke_func(priv->ctx, &inv_arg, param);
	if (ret < 0 || inv_arg.ret != 0)
		return ret ? ret : -EPROTO;

	return 0;
}

static int optee_rtc_set_alarm_wake_status(struct device *dev, bool status)
{
	struct optee_rtc *priv = dev_get_drvdata(dev);
	struct tee_ioctl_invoke_arg inv_arg = {0};
	struct tee_param param[1] = {0};
	int ret;

	if (!(priv->features & TA_RTC_FEATURE_ALARM))
		return -EOPNOTSUPP;

	inv_arg.func = PTA_CMD_RTC_SET_WAKE_ALARM_STATUS;
	inv_arg.session = priv->session_id;
	inv_arg.num_params = 1;

	param[0].attr = TEE_IOCTL_PARAM_ATTR_TYPE_VALUE_INPUT;
	param[0].u.value.a = status;

	ret = tee_client_invoke_func(priv->ctx, &inv_arg, param);

	if (ret < 0 || inv_arg.ret != 0)
		return ret ? ret : -EPROTO;

	return 0;
}

static int optee_rtc_handle_alarm_event(void *data)
{
	struct optee_rtc *priv = (struct optee_rtc *)data;
	int wait_alarm_return_status = 0;
	int ret;

	while (!kthread_should_stop()) {
		ret = optee_rtc_wait_alarm(priv->dev, &wait_alarm_return_status);
		if (ret) {
			dev_err(priv->dev, "Failed to wait for alarm: %d\n", ret);
			return ret;
		}
		switch (wait_alarm_return_status) {
		case WAIT_ALARM_ALARM_OCCURRED:
			dev_dbg(priv->dev, "Alarm occurred\n");
			rtc_update_irq(priv->rtc, 1, RTC_IRQF | RTC_AF);
			break;
		case WAIT_ALARM_CANCELED:
			dev_dbg(priv->dev, "Alarm canceled\n");
			break;
		default:
			dev_warn(priv->dev, "Unknown return status: %d\n",
				 wait_alarm_return_status);
			break;
		}
	}

	return 0;
}

static int optee_rtc_read_info(struct device *dev, struct rtc_device *rtc,
			       u64 *features)
{
	struct optee_rtc *priv = dev_get_drvdata(dev);
	struct tee_ioctl_invoke_arg inv_arg = {0};
	struct tee_param param[4] = {0};
	struct optee_rtc_info *info;
	struct optee_rtc_time *tm;
	int ret;

	inv_arg.func = PTA_CMD_RTC_GET_INF;
	inv_arg.session = priv->session_id;
	inv_arg.num_params = 4;

	param[0].attr = TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_OUTPUT;
	param[0].u.memref.shm = priv->shm;
	param[0].u.memref.size = sizeof(*info);

	ret = tee_client_invoke_func(priv->ctx, &inv_arg, param);
	if (ret < 0 || inv_arg.ret != 0)
		return ret ? ret : -EPROTO;

	info = tee_shm_get_va(priv->shm, 0);
	if (IS_ERR(info))
		return PTR_ERR(info);

	if (param[0].u.memref.size != sizeof(*info))
		return -EPROTO;

	if (info->version != RTC_INFO_VERSION)
		return -EPROTO;

	*features = info->features;

	tm = &info->range_min;
	rtc->range_min = mktime64(tm->tm_year, tm->tm_mon, tm->tm_mday, tm->tm_hour, tm->tm_min,
				  tm->tm_sec);
	tm = &info->range_max;
	rtc->range_max = mktime64(tm->tm_year, tm->tm_mon, tm->tm_mday, tm->tm_hour, tm->tm_min,
				  tm->tm_sec);

	return 0;
}

static int optee_ctx_match(struct tee_ioctl_version_data *ver, const void *data)
{
	return (ver->impl_id == TEE_IMPL_ID_OPTEE);
}

static int optee_rtc_probe(struct tee_client_device *rtc_device)
{
	struct device *dev = &rtc_device->dev;
	struct tee_ioctl_open_session_arg sess2_arg = {0};
	struct tee_ioctl_open_session_arg sess_arg = {0};
	struct optee_rtc *priv;
	struct rtc_device *rtc;
	struct tee_shm *shm;
	int ret, err;

	priv = devm_kzalloc(dev, sizeof(*priv), GFP_KERNEL);
	if (!priv)
		return -ENOMEM;

	rtc = devm_rtc_allocate_device(dev);
	if (IS_ERR(rtc))
		return PTR_ERR(rtc);

	priv->rtc = rtc;

	/* Open context with TEE driver */
	priv->ctx = tee_client_open_context(NULL, optee_ctx_match, NULL, NULL);
	if (IS_ERR(priv->ctx))
		return -ENODEV;

	/* Open first session with rtc Pseudo Trusted App */
	export_uuid(sess_arg.uuid, &rtc_device->id.uuid);
	sess_arg.clnt_login = TEE_IOCTL_LOGIN_REE_KERNEL;

	ret = tee_client_open_session(priv->ctx, &sess_arg, NULL);
	if (ret < 0 || sess_arg.ret != 0) {
		dev_err(dev, "tee_client_open_session failed, err: %x\n", sess_arg.ret);
		err = -EINVAL;
		goto out_ctx;
	}
	priv->session_id = sess_arg.session;

	/*
	 * Shared memory is used for passing an instance of either struct optee_rtc_info,
	 * struct optee_rtc_time or struct optee_rtc_alarm to OP-TEE service.
	 * The former is by definition large enough to cover both parameter cases.
	 */
	shm = tee_shm_alloc_kernel_buf(priv->ctx, sizeof(struct optee_rtc_info));
	if (IS_ERR(shm)) {
		dev_err(priv->dev, "tee_shm_alloc_kernel_buf failed\n");
		err = PTR_ERR(shm);
		goto out_sess;
	}

	priv->shm = shm;
	priv->dev = dev;
	dev_set_drvdata(dev, priv);

	rtc->ops = &optee_rtc_ops;

	err = optee_rtc_read_info(dev, rtc, &priv->features);
	if (err) {
		dev_err(dev, "Failed to get RTC features from OP-TEE\n");
		goto out_shm;
	}

	/* Handle feature's related setup before registering to rtc framework */
	if (priv->features & TA_RTC_FEATURE_ALARM) {
		priv->alarm_task = kthread_create(optee_rtc_handle_alarm_event,
						  priv, "rtc_alarm_evt");
		if (IS_ERR(priv->alarm_task)) {
			dev_err(dev, "Failed to create alarm thread\n");
			err = PTR_ERR(priv->alarm_task);
			goto out_shm;
		}

		/*
		 * In case of supported alarm feature on optee side, we create a kthread
		 * that will, in a new optee session, call a PTA interface "rtc_wait_alarm".
		 * This call return in case of alarm and in case of canceled alarm.
		 * The new optee session is therefore only needed in this case as we cannot
		 * use the same session for parallel calls to optee PTA.
		 * Hence one session is reserved to wait for alarms and the other to make
		 * standard calls to RTC PTA.
		 */

		/* Open second session with rtc Trusted App */
		export_uuid(sess2_arg.uuid, &rtc_device->id.uuid);
		sess2_arg.clnt_login = TEE_IOCTL_LOGIN_REE_KERNEL;

		ret = tee_client_open_session(priv->ctx, &sess2_arg, NULL);
		if (ret < 0 || sess2_arg.ret != 0) {
			dev_err(dev, "tee_client_open_session failed, err: %x\n", sess2_arg.ret);
			err = -EINVAL;
			goto out_thrd;
		}
		priv->session2_id = sess2_arg.session;

		if (priv->features & TA_RTC_FEATURE_WAKEUP_ALARM)
			device_init_wakeup(dev, true);
	}

	err = devm_rtc_register_device(rtc);
	if (err)
		goto out_wk;

	/*
	 * We must clear those bits after registering because registering a rtc_device
	 * will set them if it sees that .set_offset and .set_alarm are provided.
	 */
	if (!(priv->features & TA_RTC_FEATURE_CORRECTION))
		clear_bit(RTC_FEATURE_CORRECTION, rtc->features);
	if (!(priv->features & TA_RTC_FEATURE_ALARM))
		clear_bit(RTC_FEATURE_ALARM, rtc->features);

	/* Start the thread after the rtc is setup */
	if (priv->alarm_task) {
		wake_up_process(priv->alarm_task);
		dev_dbg(dev, "Wait alarm thread successfully started\n");
	}

	return 0;
out_wk:
	if (priv->features & TA_RTC_FEATURE_ALARM) {
		device_init_wakeup(dev, false);
		tee_client_close_session(priv->ctx, priv->session2_id);
	}
out_thrd:
	if (priv->features & TA_RTC_FEATURE_ALARM)
		kthread_stop(priv->alarm_task);
out_shm:
	tee_shm_free(priv->shm);
out_sess:
	tee_client_close_session(priv->ctx, priv->session_id);
out_ctx:
	tee_client_close_context(priv->ctx);

	return err;
}

static void optee_rtc_remove(struct tee_client_device *rtc_device)
{
	struct device *dev = &rtc_device->dev;
	struct optee_rtc *priv = dev_get_drvdata(dev);

	if (priv->features & TA_RTC_FEATURE_ALARM) {
		optee_rtc_cancel_wait_alarm(dev);
		kthread_stop(priv->alarm_task);
		device_init_wakeup(dev, false);
		tee_client_close_session(priv->ctx, priv->session2_id);
	}

	tee_shm_free(priv->shm);
	tee_client_close_session(priv->ctx, priv->session_id);
	tee_client_close_context(priv->ctx);
}

static int optee_rtc_suspend(struct device *dev)
{
	int res = optee_rtc_set_alarm_wake_status(dev, device_may_wakeup(dev));

	if (res) {
		dev_err(dev, "Unable to transmit wakeup information to optee rtc\n");
		return res;
	}

	return 0;
}

static DEFINE_SIMPLE_DEV_PM_OPS(optee_rtc_pm_ops, optee_rtc_suspend, NULL);

static const struct tee_client_device_id optee_rtc_id_table[] = {
	{UUID_INIT(0xf389f8c8, 0x845f, 0x496c,
		   0x8b, 0xbe, 0xd6, 0x4b, 0xd2, 0x4c, 0x92, 0xfd)},
	{}
};

MODULE_DEVICE_TABLE(tee, optee_rtc_id_table);

static struct tee_client_driver optee_rtc_driver = {
	.id_table	= optee_rtc_id_table,
	.probe		= optee_rtc_probe,
	.remove		= optee_rtc_remove,
	.driver		= {
		.name		= "optee_rtc",
		.pm		= pm_sleep_ptr(&optee_rtc_pm_ops),
	},
};

module_tee_client_driver(optee_rtc_driver);

MODULE_LICENSE("GPL v2");
MODULE_AUTHOR("Clément Léger <clement.leger@bootlin.com>");
MODULE_DESCRIPTION("OP-TEE based RTC driver");