// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright (C) 2025 IBM Corporation, Srish Srinivasan <ssrish@linux.ibm.com>
 */

#include <keys/trusted_pkwm.h>
#include <keys/trusted-type.h>
#include <linux/build_bug.h>
#include <linux/key-type.h>
#include <linux/parser.h>
#include <asm/plpks.h>

enum {
        Opt_err,
        Opt_wrap_flags,
};

static const match_table_t key_tokens = {
        {Opt_wrap_flags, "wrap_flags=%s"},
        {Opt_err, NULL}
};

static int getoptions(char *datablob, struct trusted_key_options *opt)
{
        substring_t args[MAX_OPT_ARGS];
        char *p = datablob;
        int token;
        int res;
        u16 wrap_flags;
        unsigned long token_mask = 0;
        struct trusted_pkwm_options *pkwm;

        if (!datablob)
                return 0;

        pkwm = opt->private;

        while ((p = strsep(&datablob, " \t"))) {
                if (*p == '\0' || *p == ' ' || *p == '\t')
                        continue;

                token = match_token(p, key_tokens, args);
                if (test_and_set_bit(token, &token_mask))
                        return -EINVAL;

                switch (token) {
                case Opt_wrap_flags:
                        res = kstrtou16(args[0].from, 16, &wrap_flags);
                        if (res < 0 || wrap_flags > 2)
                                return -EINVAL;
                        pkwm->wrap_flags = wrap_flags;
                        break;
                default:
                        return -EINVAL;
                }
        }
        return 0;
}

static struct trusted_key_options *trusted_options_alloc(void)
{
        struct trusted_key_options *options;
        struct trusted_pkwm_options *pkwm;

        options = kzalloc_obj(*options);

        if (options) {
                pkwm = kzalloc_obj(*pkwm);

                if (!pkwm) {
                        kfree_sensitive(options);
                        options = NULL;
                } else {
                        options->private = pkwm;
                }
        }

        return options;
}

static int trusted_pkwm_seal(struct trusted_key_payload *p, char *datablob)
{
        struct trusted_key_options *options = NULL;
        struct trusted_pkwm_options *pkwm = NULL;
        u8 *input_buf, *output_buf;
        u32 output_len, input_len;
        int rc;

        options = trusted_options_alloc();

        if (!options)
                return -ENOMEM;

        rc = getoptions(datablob, options);
        if (rc < 0)
                goto out;
        dump_options(options);

        input_len = p->key_len;
        input_buf = kmalloc(ALIGN(input_len, 4096), GFP_KERNEL);
        if (!input_buf) {
                pr_err("Input buffer allocation failed. Returning -ENOMEM.");
                rc = -ENOMEM;
                goto out;
        }

        memcpy(input_buf, p->key, p->key_len);

        pkwm = options->private;

        rc = plpks_wrap_object(&input_buf, input_len, pkwm->wrap_flags,
                               &output_buf, &output_len);
        if (!rc) {
                memcpy(p->blob, output_buf, output_len);
                p->blob_len = output_len;
                dump_payload(p);
        } else {
                pr_err("Wrapping of payload key failed: %d\n", rc);
        }

        kfree(input_buf);
        kfree(output_buf);

out:
        kfree_sensitive(options->private);
        kfree_sensitive(options);
        return rc;
}

static int trusted_pkwm_unseal(struct trusted_key_payload *p, char *datablob)
{
        u8 *input_buf, *output_buf;
        u32 input_len, output_len;
        int rc;

        input_len = p->blob_len;
        input_buf = kmalloc(ALIGN(input_len, 4096), GFP_KERNEL);
        if (!input_buf) {
                pr_err("Input buffer allocation failed. Returning -ENOMEM.");
                return -ENOMEM;
        }

        memcpy(input_buf, p->blob, p->blob_len);

        rc = plpks_unwrap_object(&input_buf, input_len, &output_buf,
                                 &output_len);
        if (!rc) {
                memcpy(p->key, output_buf, output_len);
                p->key_len = output_len;
                dump_payload(p);
        } else {
                pr_err("Unwrapping of payload failed: %d\n", rc);
        }

        kfree(input_buf);
        kfree(output_buf);

        return rc;
}

static int trusted_pkwm_init(void)
{
        int ret;

        if (!plpks_wrapping_is_supported()) {
                pr_err("H_PKS_WRAP_OBJECT interface not supported\n");
                return -ENODEV;
        }

        ret = plpks_gen_wrapping_key();
        if (ret) {
                pr_err("Failed to generate default wrapping key\n");
                return -EINVAL;
        }

        return register_key_type(&key_type_trusted);
}

static void trusted_pkwm_exit(void)
{
        unregister_key_type(&key_type_trusted);
}

struct trusted_key_ops pkwm_trusted_key_ops = {
        .migratable = 0, /* non-migratable */
        .init = trusted_pkwm_init,
        .seal = trusted_pkwm_seal,
        .unseal = trusted_pkwm_unseal,
        .exit = trusted_pkwm_exit,
};