// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries.
 */

#define pr_fmt(fmt) KBUILD_MODNAME ": " fmt

#include <linux/slab.h>
#include <linux/tee.h>
#include <linux/platform_device.h>
#include <linux/xarray.h>

#include "qcomtee.h"

static int find_qtee_object(struct qcomtee_object **object, unsigned long id,
                            struct qcomtee_context_data *ctxdata)
{
        int err = 0;

        guard(rcu)();
        /* Object release is RCU protected. */
        *object = idr_find(&ctxdata->qtee_objects_idr, id);
        if (!qcomtee_object_get(*object))
                err = -EINVAL;

        return err;
}

static void del_qtee_object(unsigned long id,
                            struct qcomtee_context_data *ctxdata)
{
        struct qcomtee_object *object;

        scoped_guard(mutex, &ctxdata->qtee_lock)
                object = idr_remove(&ctxdata->qtee_objects_idr, id);

        qcomtee_object_put(object);
}

/**
 * qcomtee_context_add_qtee_object() - Add a QTEE object to the context.
 * @param: TEE parameter representing @object.
 * @object: QTEE object.
 * @ctx: context to add the object.
 *
 * It assumes @object is %QCOMTEE_OBJECT_TYPE_TEE and the caller has already
 * issued qcomtee_object_get() for @object.
 *
 * Return: On success, returns 0; on failure, returns < 0.
 */
int qcomtee_context_add_qtee_object(struct tee_param *param,
                                    struct qcomtee_object *object,
                                    struct tee_context *ctx)
{
        int ret;
        struct qcomtee_context_data *ctxdata = ctx->data;

        scoped_guard(mutex, &ctxdata->qtee_lock)
                ret = idr_alloc(&ctxdata->qtee_objects_idr, object, 0, 0,
                                GFP_KERNEL);
        if (ret < 0)
                return ret;

        param->u.objref.id = ret;
        /* QTEE Object: QCOMTEE_OBJREF_FLAG_TEE set. */
        param->u.objref.flags = QCOMTEE_OBJREF_FLAG_TEE;

        return 0;
}

/* Retrieve the QTEE object added with qcomtee_context_add_qtee_object(). */
int qcomtee_context_find_qtee_object(struct qcomtee_object **object,
                                     struct tee_param *param,
                                     struct tee_context *ctx)
{
        struct qcomtee_context_data *ctxdata = ctx->data;

        return find_qtee_object(object, param->u.objref.id, ctxdata);
}

/**
 * qcomtee_context_del_qtee_object() - Delete a QTEE object from the context.
 * @param: TEE parameter representing @object.
 * @ctx: context for deleting the object.
 *
 * The @param has been initialized by qcomtee_context_add_qtee_object().
 */
void qcomtee_context_del_qtee_object(struct tee_param *param,
                                     struct tee_context *ctx)
{
        struct qcomtee_context_data *ctxdata = ctx->data;
        /* 'qtee_objects_idr' stores QTEE objects only. */
        if (param->u.objref.flags & QCOMTEE_OBJREF_FLAG_TEE)
                del_qtee_object(param->u.objref.id, ctxdata);
}

/**
 * qcomtee_objref_to_arg() - Convert OBJREF parameter to QTEE argument.
 * @arg: QTEE argument.
 * @param: TEE parameter.
 * @ctx: context in which the conversion should happen.
 *
 * It assumes @param is an OBJREF.
 * It does not set @arg.type; the caller should initialize it to a correct
 * &enum qcomtee_arg_type value. It gets the object's refcount in @arg;
 * the caller should manage to put it afterward.
 *
 * Return: On success, returns 0; on failure, returns < 0.
 */
int qcomtee_objref_to_arg(struct qcomtee_arg *arg, struct tee_param *param,
                          struct tee_context *ctx)
{
        int err = -EINVAL;

        arg->o = NULL_QCOMTEE_OBJECT;
        /* param is a NULL object: */
        if (param->u.objref.id == TEE_OBJREF_NULL)
                return 0;

        /* param is a callback object: */
        if (param->u.objref.flags & QCOMTEE_OBJREF_FLAG_USER)
                err =  qcomtee_user_param_to_object(&arg->o, param, ctx);
        /* param is a QTEE object: */
        else if (param->u.objref.flags & QCOMTEE_OBJREF_FLAG_TEE)
                err = qcomtee_context_find_qtee_object(&arg->o, param, ctx);
        /* param is a memory object: */
        else if (param->u.objref.flags & QCOMTEE_OBJREF_FLAG_MEM)
                err = qcomtee_memobj_param_to_object(&arg->o, param, ctx);

        /*
         * For callback objects, call qcomtee_object_get() to keep a temporary
         * copy for the driver, as these objects are released asynchronously
         * and may disappear even before returning from QTEE.
         *
         *  - For direct object invocations, the matching put is called in
         *    qcomtee_object_invoke() when parsing the QTEE response.
         *  - For callback responses, put is called in qcomtee_user_object_notify()
         *    after QTEE has received its copies.
         */

        if (!err && (typeof_qcomtee_object(arg->o) == QCOMTEE_OBJECT_TYPE_CB))
                qcomtee_object_get(arg->o);

        return err;
}

/**
 * qcomtee_objref_from_arg() - Convert QTEE argument to OBJREF param.
 * @param: TEE parameter.
 * @arg: QTEE argument.
 * @ctx: context in which the conversion should happen.
 *
 * It assumes @arg is of %QCOMTEE_ARG_TYPE_IO or %QCOMTEE_ARG_TYPE_OO.
 * It does not set @param.attr; the caller should initialize it to a
 * correct type.
 *
 * Return: On success, returns 0; on failure, returns < 0.
 */
int qcomtee_objref_from_arg(struct tee_param *param, struct qcomtee_arg *arg,
                            struct tee_context *ctx)
{
        struct qcomtee_object *object = arg->o;

        switch (typeof_qcomtee_object(object)) {
        case QCOMTEE_OBJECT_TYPE_NULL:
                param->u.objref.id = TEE_OBJREF_NULL;

                return 0;
        case QCOMTEE_OBJECT_TYPE_CB:
                /* object is a callback object: */
                if (is_qcomtee_user_object(object))
                        return qcomtee_user_param_from_object(param, object,
                                                              ctx);
                /* object is a memory object: */
                else if (is_qcomtee_memobj_object(object))
                        return qcomtee_memobj_param_from_object(param, object,
                                                               ctx);

                break;
        case QCOMTEE_OBJECT_TYPE_TEE:
                return qcomtee_context_add_qtee_object(param, object, ctx);

        case QCOMTEE_OBJECT_TYPE_ROOT:
        default:
                break;
        }

        return -EINVAL;
}

/**
 * qcomtee_params_to_args() - Convert TEE parameters to QTEE arguments.
 * @u: QTEE arguments.
 * @params: TEE parameters.
 * @num_params: number of elements in the parameter array.
 * @ctx: context in which the conversion should happen.
 *
 * It assumes @u has at least @num_params + 1 entries and has been initialized
 * with %QCOMTEE_ARG_TYPE_INV as &struct qcomtee_arg.type.
 *
 * Return: On success, returns 0; on failure, returns < 0.
 */
static int qcomtee_params_to_args(struct qcomtee_arg *u,
                                  struct tee_param *params, int num_params,
                                  struct tee_context *ctx)
{
        int i;

        for (i = 0; i < num_params; i++) {
                switch (params[i].attr) {
                case TEE_IOCTL_PARAM_ATTR_TYPE_UBUF_INPUT:
                case TEE_IOCTL_PARAM_ATTR_TYPE_UBUF_OUTPUT:
                        u[i].flags = QCOMTEE_ARG_FLAGS_UADDR;
                        u[i].b.uaddr = params[i].u.ubuf.uaddr;
                        u[i].b.size = params[i].u.ubuf.size;

                        if (params[i].attr ==
                            TEE_IOCTL_PARAM_ATTR_TYPE_UBUF_INPUT)
                                u[i].type = QCOMTEE_ARG_TYPE_IB;
                        else /* TEE_IOCTL_PARAM_ATTR_TYPE_UBUF_OUTPUT */
                                u[i].type = QCOMTEE_ARG_TYPE_OB;

                        break;
                case TEE_IOCTL_PARAM_ATTR_TYPE_OBJREF_INPUT:
                        u[i].type = QCOMTEE_ARG_TYPE_IO;
                        if (qcomtee_objref_to_arg(&u[i], &params[i], ctx))
                                goto out_failed;

                        break;
                case TEE_IOCTL_PARAM_ATTR_TYPE_OBJREF_OUTPUT:
                        u[i].type = QCOMTEE_ARG_TYPE_OO;
                        u[i].o = NULL_QCOMTEE_OBJECT;
                        break;
                default:
                        goto out_failed;
                }
        }

        return 0;

out_failed:
        /* Undo qcomtee_objref_to_arg(). */
        for (i--; i >= 0; i--) {
                if (u[i].type != QCOMTEE_ARG_TYPE_IO)
                        continue;

                qcomtee_user_object_set_notify(u[i].o, false);
                /* See docs for qcomtee_objref_to_arg() for double put. */
                if (typeof_qcomtee_object(u[i].o) == QCOMTEE_OBJECT_TYPE_CB)
                        qcomtee_object_put(u[i].o);

                qcomtee_object_put(u[i].o);
        }

        return -EINVAL;
}

/**
 * qcomtee_params_from_args() - Convert QTEE arguments to TEE parameters.
 * @params: TEE parameters.
 * @u: QTEE arguments.
 * @num_params: number of elements in the parameter array.
 * @ctx: context in which the conversion should happen.
 *
 * @u should have already been initialized by qcomtee_params_to_args().
 * This also represents the end of a QTEE invocation that started with
 * qcomtee_params_to_args() by releasing %QCOMTEE_ARG_TYPE_IO objects.
 *
 * Return: On success, returns 0; on failure, returns < 0.
 */
static int qcomtee_params_from_args(struct tee_param *params,
                                    struct qcomtee_arg *u, int num_params,
                                    struct tee_context *ctx)
{
        int i, np;

        qcomtee_arg_for_each(np, u) {
                switch (u[np].type) {
                case QCOMTEE_ARG_TYPE_OB:
                        /* TEE_IOCTL_PARAM_ATTR_TYPE_UBUF_OUTPUT */
                        params[np].u.ubuf.size = u[np].b.size;

                        break;
                case QCOMTEE_ARG_TYPE_IO:
                        /* IEE_IOCTL_PARAM_ATTR_TYPE_OBJREF_INPUT */
                        qcomtee_object_put(u[np].o);

                        break;
                case QCOMTEE_ARG_TYPE_OO:
                        /* TEE_IOCTL_PARAM_ATTR_TYPE_OBJREF_OUTPUT */
                        if (qcomtee_objref_from_arg(&params[np], &u[np], ctx))
                                goto out_failed;

                        break;
                case QCOMTEE_ARG_TYPE_IB:
                default:
                        break;
                }
        }

        return 0;

out_failed:
        /* Undo qcomtee_objref_from_arg(). */
        for (i = 0; i < np; i++) {
                if (params[i].attr == TEE_IOCTL_PARAM_ATTR_TYPE_OBJREF_OUTPUT)
                        qcomtee_context_del_qtee_object(&params[i], ctx);
        }

        /* Release any IO and OO objects not processed. */
        for (; i < num_params && u[i].type; i++) {
                if (u[i].type == QCOMTEE_ARG_TYPE_OO ||
                    u[i].type == QCOMTEE_ARG_TYPE_IO)
                        qcomtee_object_put(u[i].o);
        }

        return -EINVAL;
}

/* TEE Device Ops. */

static int qcomtee_params_check(struct tee_param *params, int num_params)
{
        int io = 0, oo = 0, ib = 0, ob = 0;
        int i;

        /* QTEE can accept 64 arguments. */
        if (num_params > QCOMTEE_ARGS_MAX)
                return -EINVAL;

        /* Supported parameter types. */
        for (i = 0; i < num_params; i++) {
                switch (params[i].attr) {
                case TEE_IOCTL_PARAM_ATTR_TYPE_UBUF_INPUT:
                        ib++;
                        break;
                case TEE_IOCTL_PARAM_ATTR_TYPE_UBUF_OUTPUT:
                        ob++;
                        break;
                case TEE_IOCTL_PARAM_ATTR_TYPE_OBJREF_INPUT:
                        io++;
                        break;
                case TEE_IOCTL_PARAM_ATTR_TYPE_OBJREF_OUTPUT:
                        oo++;
                        break;
                default:
                        return -EINVAL;
                }
        }

        /*  QTEE can accept 16 arguments of each supported types. */
        if (io > QCOMTEE_ARGS_PER_TYPE || oo > QCOMTEE_ARGS_PER_TYPE ||
            ib > QCOMTEE_ARGS_PER_TYPE || ob > QCOMTEE_ARGS_PER_TYPE)
                return -EINVAL;

        return 0;
}

/* Check if an operation on ROOT_QCOMTEE_OBJECT from userspace is permitted. */
static int qcomtee_root_object_check(u32 op, struct tee_param *params,
                                     int num_params)
{
        /* Some privileged operations recognized by QTEE. */
        if (op == QCOMTEE_ROOT_OP_NOTIFY_DOMAIN_CHANGE ||
            op == QCOMTEE_ROOT_OP_ADCI_ACCEPT ||
            op == QCOMTEE_ROOT_OP_ADCI_SHUTDOWN)
                return -EINVAL;

        /*
         * QCOMTEE_ROOT_OP_REG_WITH_CREDENTIALS is to register with QTEE
         * by passing a credential object as input OBJREF. TEE_OBJREF_NULL as a
         * credential object represents a privileged client for QTEE and
         * is used by the kernel only.
         */
        if (op == QCOMTEE_ROOT_OP_REG_WITH_CREDENTIALS && num_params == 2) {
                if (params[0].attr == TEE_IOCTL_PARAM_ATTR_TYPE_OBJREF_INPUT &&
                    params[1].attr == TEE_IOCTL_PARAM_ATTR_TYPE_OBJREF_OUTPUT) {
                        if (params[0].u.objref.id == TEE_OBJREF_NULL)
                                return -EINVAL;
                }
        }

        return 0;
}

/**
 * qcomtee_object_invoke() - Invoke a QTEE object.
 * @ctx: TEE context.
 * @arg: ioctl arguments.
 * @params: parameters for the object.
 *
 * Return: On success, returns 0; on failure, returns < 0.
 */
static int qcomtee_object_invoke(struct tee_context *ctx,
                                 struct tee_ioctl_object_invoke_arg *arg,
                                 struct tee_param *params)
{
        struct qcomtee_context_data *ctxdata = ctx->data;
        struct qcomtee_object *object;
        int i, ret, result;

        if (qcomtee_params_check(params, arg->num_params))
                return -EINVAL;

        /* First, handle reserved operations: */
        if (arg->op == QCOMTEE_MSG_OBJECT_OP_RELEASE) {
                del_qtee_object(arg->id, ctxdata);

                return 0;
        }

        /* Otherwise, invoke a QTEE object: */
        struct qcomtee_object_invoke_ctx *oic __free(kfree) =
                qcomtee_object_invoke_ctx_alloc(ctx);
        if (!oic)
                return -ENOMEM;

        /* +1 for ending QCOMTEE_ARG_TYPE_INV. */
        struct qcomtee_arg *u __free(kfree) = kzalloc_objs(*u,
                                                           arg->num_params + 1);
        if (!u)
                return -ENOMEM;

        /* Get an object to invoke. */
        if (arg->id == TEE_OBJREF_NULL) {
                /* Use ROOT if TEE_OBJREF_NULL is invoked. */
                if (qcomtee_root_object_check(arg->op, params, arg->num_params))
                        return -EINVAL;

                object = ROOT_QCOMTEE_OBJECT;
        } else if (find_qtee_object(&object, arg->id, ctxdata)) {
                return -EINVAL;
        }

        ret = qcomtee_params_to_args(u, params, arg->num_params, ctx);
        if (ret)
                goto out;

        ret = qcomtee_object_do_invoke(oic, object, arg->op, u, &result);
        if (ret) {
                qcomtee_arg_for_each_input_object(i, u) {
                        qcomtee_user_object_set_notify(u[i].o, false);
                        qcomtee_object_put(u[i].o);
                }

                goto out;
        }

        /* Prase QTEE response and put driver's object copies: */

        if (!result) {
                /* Assume service is UNAVAIL if unable to process the result. */
                if (qcomtee_params_from_args(params, u, arg->num_params, ctx))
                        result = QCOMTEE_MSG_ERROR_UNAVAIL;
        } else {
                /*
                 * qcomtee_params_to_args() gets a copy of IO for the driver to
                 * make sure they do not get released while in the middle of
                 * invocation. On success (!result), qcomtee_params_from_args()
                 * puts them; Otherwise, put them here.
                 */
                qcomtee_arg_for_each_input_object(i, u)
                        qcomtee_object_put(u[i].o);
        }

        arg->ret = result;
out:
        qcomtee_object_put(object);

        return ret;
}

/**
 * qcomtee_supp_recv() - Wait for a request for the supplicant.
 * @ctx: TEE context.
 * @op: requested operation on the object.
 * @num_params: number of elements in the parameter array.
 * @params: parameters for @op.
 *
 * The first parameter is a meta %TEE_IOCTL_PARAM_ATTR_TYPE_VALUE_INOUT.
 * On input, it provides a user buffer. This buffer is used for parameters of
 * type %TEE_IOCTL_PARAM_ATTR_TYPE_UBUF_INPUT in qcomtee_cb_params_from_args().
 * On output, the object ID and request ID are stored in the meta parameter.
 *
 * @num_params is updated to the number of parameters that actually exist
 * in @params on return.
 *
 * Return: On success, returns 0; on failure, returns < 0.
 */
static int qcomtee_supp_recv(struct tee_context *ctx, u32 *op, u32 *num_params,
                             struct tee_param *params)
{
        struct qcomtee_user_object_request_data data;
        void __user *uaddr;
        size_t ubuf_size;
        int i, ret;

        if (!*num_params)
                return -EINVAL;

        /* First parameter should be an INOUT + meta parameter. */
        if (params->attr !=
            (TEE_IOCTL_PARAM_ATTR_TYPE_VALUE_INOUT | TEE_IOCTL_PARAM_ATTR_META))
                return -EINVAL;

        /* Other parameters are none. */
        for (i = 1; i < *num_params; i++)
                if (params[i].attr)
                        return -EINVAL;

        if (!IS_ALIGNED(params->u.value.a, 8))
                return -EINVAL;

        /* User buffer and size from meta parameter. */
        uaddr = u64_to_user_ptr(params->u.value.a);
        ubuf_size = params->u.value.b;
        /* Process TEE parameters. +/-1 to ignore the meta parameter. */
        ret = qcomtee_user_object_select(ctx, params + 1, *num_params - 1,
                                         uaddr, ubuf_size, &data);
        if (ret)
                return ret;

        params->u.value.a = data.object_id;
        params->u.value.b = data.id;
        params->u.value.c = 0;
        *op = data.op;
        *num_params = data.np + 1;

        return 0;
}

/**
 * qcomtee_supp_send() - Submit a response for a request.
 * @ctx: TEE context.
 * @errno: return value for the request.
 * @num_params: number of elements in the parameter array.
 * @params: returned parameters.
 *
 * The first parameter is a meta %TEE_IOCTL_PARAM_ATTR_TYPE_VALUE_OUTPUT.
 * It specifies the request ID this response belongs to.
 *
 * Return: On success, returns 0; on failure, returns < 0.
 */
static int qcomtee_supp_send(struct tee_context *ctx, u32 errno, u32 num_params,
                             struct tee_param *params)
{
        int req_id;

        if (!num_params)
                return -EINVAL;

        /* First parameter should be an OUTPUT + meta parameter. */
        if (params->attr != (TEE_IOCTL_PARAM_ATTR_TYPE_VALUE_OUTPUT |
                             TEE_IOCTL_PARAM_ATTR_META))
                return -EINVAL;

        req_id = params->u.value.a;
        /* Process TEE parameters. +/-1 to ignore the meta parameter. */
        return qcomtee_user_object_submit(ctx, params + 1, num_params - 1,
                                          req_id, errno);
}

static int qcomtee_open(struct tee_context *ctx)
{
        struct qcomtee_context_data *ctxdata __free(kfree) = kzalloc_obj(*ctxdata);
        if (!ctxdata)
                return -ENOMEM;

        /*
         * In the QTEE driver, the same context is used to refcount resources
         * shared by QTEE. For example, teedev_ctx_get() is called for any
         * instance of callback objects (see qcomtee_user_param_to_object()).
         *
         * Maintain a copy of teedev for QTEE as it serves as a direct user of
         * this context. The teedev will be released in the context's release().
         *
         * tee_device_unregister() will remain blocked until all contexts
         * are released. This includes contexts owned by the user, which are
         * closed by teedev_close_context(), as well as those owned by QTEE
         * closed by teedev_ctx_put() in object's release().
         */
        if (!tee_device_get(ctx->teedev))
                return -EINVAL;

        idr_init(&ctxdata->qtee_objects_idr);
        mutex_init(&ctxdata->qtee_lock);
        idr_init(&ctxdata->reqs_idr);
        INIT_LIST_HEAD(&ctxdata->reqs_list);
        mutex_init(&ctxdata->reqs_lock);
        init_completion(&ctxdata->req_c);

        ctx->data = no_free_ptr(ctxdata);

        return 0;
}

/* Gets called when the user closes the device */
static void qcomtee_close_context(struct tee_context *ctx)
{
        struct qcomtee_context_data *ctxdata = ctx->data;
        struct qcomtee_object *object;
        int id;

        /* Process QUEUED or PROCESSING requests. */
        qcomtee_requests_destroy(ctxdata);
        /* Release QTEE objects. */
        idr_for_each_entry(&ctxdata->qtee_objects_idr, object, id)
                qcomtee_object_put(object);
}

/* Gets called when the final reference to the context goes away. */
static void qcomtee_release(struct tee_context *ctx)
{
        struct qcomtee_context_data *ctxdata = ctx->data;

        idr_destroy(&ctxdata->qtee_objects_idr);
        idr_destroy(&ctxdata->reqs_idr);
        kfree(ctxdata);

        /* There is nothing shared in this context with QTEE. */
        tee_device_put(ctx->teedev);
}

static void qcomtee_get_version(struct tee_device *teedev,
                                struct tee_ioctl_version_data *vers)
{
        struct tee_ioctl_version_data v = {
                .impl_id = TEE_IMPL_ID_QTEE,
                .gen_caps = TEE_GEN_CAP_OBJREF,
        };

        *vers = v;
}

/**
 * qcomtee_get_qtee_feature_list() - Query QTEE features versions.
 * @ctx: TEE context.
 * @id: ID of the feature to query.
 * @version: version of the feature.
 *
 * Used to query the verion of features supported by QTEE.
 */
static void qcomtee_get_qtee_feature_list(struct tee_context *ctx, u32 id,
                                          u32 *version)
{
        struct qcomtee_object *client_env, *service;
        struct qcomtee_arg u[3] = { 0 };
        int result;

        struct qcomtee_object_invoke_ctx *oic __free(kfree) =
                qcomtee_object_invoke_ctx_alloc(ctx);
        if (!oic)
                return;

        client_env = qcomtee_object_get_client_env(oic);
        if (client_env == NULL_QCOMTEE_OBJECT)
                return;

        /* Get ''FeatureVersions Service'' object. */
        service = qcomtee_object_get_service(oic, client_env,
                                             QCOMTEE_FEATURE_VER_UID);
        if (service == NULL_QCOMTEE_OBJECT)
                goto out_failed;

        /* IB: Feature to query. */
        u[0].b.addr = &id;
        u[0].b.size = sizeof(id);
        u[0].type = QCOMTEE_ARG_TYPE_IB;

        /* OB: Version returned. */
        u[1].b.addr = version;
        u[1].b.size = sizeof(*version);
        u[1].type = QCOMTEE_ARG_TYPE_OB;

        qcomtee_object_do_invoke(oic, service, QCOMTEE_FEATURE_VER_OP_GET, u,
                                 &result);

out_failed:
        qcomtee_object_put(service);
        qcomtee_object_put(client_env);
}

static const struct tee_driver_ops qcomtee_ops = {
        .get_version = qcomtee_get_version,
        .open = qcomtee_open,
        .close_context = qcomtee_close_context,
        .release = qcomtee_release,
        .object_invoke_func = qcomtee_object_invoke,
        .supp_recv = qcomtee_supp_recv,
        .supp_send = qcomtee_supp_send,
};

static const struct tee_desc qcomtee_desc = {
        .name = "qcomtee",
        .ops = &qcomtee_ops,
        .owner = THIS_MODULE,
};

static int qcomtee_probe(struct platform_device *pdev)
{
        struct workqueue_struct *async_wq;
        struct tee_device *teedev;
        struct tee_shm_pool *pool;
        struct tee_context *ctx;
        struct qcomtee *qcomtee;
        int err;

        qcomtee = kzalloc_obj(*qcomtee);
        if (!qcomtee)
                return -ENOMEM;

        pool = qcomtee_shm_pool_alloc();
        if (IS_ERR(pool)) {
                err = PTR_ERR(pool);

                goto err_free_qcomtee;
        }

        teedev = tee_device_alloc(&qcomtee_desc, NULL, pool, qcomtee);
        if (IS_ERR(teedev)) {
                err = PTR_ERR(teedev);

                goto err_pool_destroy;
        }

        qcomtee->teedev = teedev;
        qcomtee->pool = pool;
        err = tee_device_register(qcomtee->teedev);
        if (err)
                goto err_unreg_teedev;

        platform_set_drvdata(pdev, qcomtee);
        /* Start async wq. */
        async_wq = alloc_ordered_workqueue("qcomtee_wq", 0);
        if (!async_wq) {
                err = -ENOMEM;

                goto err_unreg_teedev;
        }

        qcomtee->wq = async_wq;
        /* Driver context used for async operations of teedev. */
        ctx = teedev_open(qcomtee->teedev);
        if (IS_ERR(ctx)) {
                err = PTR_ERR(ctx);

                goto err_dest_wq;
        }

        qcomtee->ctx = ctx;
        /* Init Object table. */
        qcomtee->xa_last_id = 0;
        xa_init_flags(&qcomtee->xa_local_objects, XA_FLAGS_ALLOC);
        /* Get QTEE verion. */
        qcomtee_get_qtee_feature_list(qcomtee->ctx,
                                      QCOMTEE_FEATURE_VER_OP_GET_QTEE_ID,
                                      &qcomtee->qtee_version);

        pr_info("QTEE version %u.%u.%u\n",
                QTEE_VERSION_GET_MAJOR(qcomtee->qtee_version),
                QTEE_VERSION_GET_MINOR(qcomtee->qtee_version),
                QTEE_VERSION_GET_PATCH(qcomtee->qtee_version));

        return 0;

err_dest_wq:
        destroy_workqueue(qcomtee->wq);
err_unreg_teedev:
        tee_device_unregister(qcomtee->teedev);
err_pool_destroy:
        tee_shm_pool_free(pool);
err_free_qcomtee:
        kfree(qcomtee);

        return err;
}

/**
 * qcomtee_remove() - Device Removal Routine.
 * @pdev: platform device information struct.
 *
 * It is called by the platform subsystem to alert the driver that it should
 * release the device.
 *
 * QTEE does not provide an API to inform it about a callback object going away.
 * However, when releasing QTEE objects, any callback object sent to QTEE
 * previously would be released by QTEE as part of the object release.
 */
static void qcomtee_remove(struct platform_device *pdev)
{
        struct qcomtee *qcomtee = platform_get_drvdata(pdev);

        teedev_close_context(qcomtee->ctx);
        /* Wait for RELEASE operations to be processed for QTEE objects. */
        tee_device_unregister(qcomtee->teedev);
        destroy_workqueue(qcomtee->wq);
        tee_shm_pool_free(qcomtee->pool);
        kfree(qcomtee);
}

static const struct platform_device_id qcomtee_ids[] = { { "qcomtee", 0 }, {} };
MODULE_DEVICE_TABLE(platform, qcomtee_ids);

static struct platform_driver qcomtee_platform_driver = {
        .probe = qcomtee_probe,
        .remove = qcomtee_remove,
        .driver = {
                .name = "qcomtee",
        },
        .id_table = qcomtee_ids,
};

module_platform_driver(qcomtee_platform_driver);

MODULE_AUTHOR("Qualcomm");
MODULE_DESCRIPTION("QTEE driver");
MODULE_VERSION("1.0");
MODULE_LICENSE("GPL");