// SPDX-License-Identifier: GPL-2.0
/*
 * Miscellaneous cgroup controller
 *
 * Copyright 2020 Google LLC
 * Author: Vipin Sharma <vipinsh@google.com>
 */

#include <linux/limits.h>
#include <linux/cgroup.h>
#include <linux/errno.h>
#include <linux/atomic.h>
#include <linux/slab.h>
#include <linux/misc_cgroup.h>

#define MAX_STR "max"
#define MAX_NUM U64_MAX

/* Miscellaneous res name, keep it in sync with enum misc_res_type */
static const char *const misc_res_name[] = {
#ifdef CONFIG_KVM_AMD_SEV
        /* AMD SEV ASIDs resource */
        "sev",
        /* AMD SEV-ES ASIDs resource */
        "sev_es",
#endif
#ifdef CONFIG_INTEL_TDX_HOST
        /* Intel TDX HKIDs resource */
        "tdx",
#endif
};

/* Root misc cgroup */
static struct misc_cg root_cg;

/*
 * Miscellaneous resources capacity for the entire machine. 0 capacity means
 * resource is not initialized or not present in the host.
 *
 * root_cg.max and capacity are independent of each other. root_cg.max can be
 * more than the actual capacity. We are using Limits resource distribution
 * model of cgroup for miscellaneous controller.
 */
static u64 misc_res_capacity[MISC_CG_RES_TYPES];

/**
 * parent_misc() - Get the parent of the passed misc cgroup.
 * @cgroup: cgroup whose parent needs to be fetched.
 *
 * Context: Any context.
 * Return:
 * * struct misc_cg* - Parent of the @cgroup.
 * * %NULL - If @cgroup is null or the passed cgroup does not have a parent.
 */
static struct misc_cg *parent_misc(struct misc_cg *cgroup)
{
        return cgroup ? css_misc(cgroup->css.parent) : NULL;
}

/**
 * valid_type() - Check if @type is valid or not.
 * @type: misc res type.
 *
 * Context: Any context.
 * Return:
 * * true - If valid type.
 * * false - If not valid type.
 */
static inline bool valid_type(enum misc_res_type type)
{
        return type >= 0 && type < MISC_CG_RES_TYPES;
}

/**
 * misc_cg_set_capacity() - Set the capacity of the misc cgroup res.
 * @type: Type of the misc res.
 * @capacity: Supported capacity of the misc res on the host.
 *
 * If capacity is 0 then the charging a misc cgroup fails for that type.
 *
 * Context: Any context.
 * Return:
 * * %0 - Successfully registered the capacity.
 * * %-EINVAL - If @type is invalid.
 */
int misc_cg_set_capacity(enum misc_res_type type, u64 capacity)
{
        if (!valid_type(type))
                return -EINVAL;

        WRITE_ONCE(misc_res_capacity[type], capacity);
        return 0;
}
EXPORT_SYMBOL_GPL(misc_cg_set_capacity);

/**
 * misc_cg_cancel_charge() - Cancel the charge from the misc cgroup.
 * @type: Misc res type in misc cg to cancel the charge from.
 * @cg: Misc cgroup to cancel charge from.
 * @amount: Amount to cancel.
 *
 * Context: Any context.
 */
static void misc_cg_cancel_charge(enum misc_res_type type, struct misc_cg *cg,
                                  u64 amount)
{
        WARN_ONCE(atomic64_add_negative(-amount, &cg->res[type].usage),
                  "misc cgroup resource %s became less than 0",
                  misc_res_name[type]);
}

static void misc_cg_update_watermark(struct misc_res *res, u64 new_usage)
{
        u64 old;

        while (true) {
                old = atomic64_read(&res->watermark);
                if (new_usage <= old)
                        break;
                if (atomic64_cmpxchg(&res->watermark, old, new_usage) == old)
                        break;
        }
}

static void misc_cg_event(enum misc_res_type type, struct misc_cg *cg)
{
        atomic64_inc(&cg->res[type].events_local);
        cgroup_file_notify(&cg->events_local_file);

        for (; parent_misc(cg); cg = parent_misc(cg)) {
                atomic64_inc(&cg->res[type].events);
                cgroup_file_notify(&cg->events_file);
        }
}

/**
 * misc_cg_try_charge() - Try charging the misc cgroup.
 * @type: Misc res type to charge.
 * @cg: Misc cgroup which will be charged.
 * @amount: Amount to charge.
 *
 * Charge @amount to the misc cgroup. Caller must use the same cgroup during
 * the uncharge call.
 *
 * Context: Any context.
 * Return:
 * * %0 - If successfully charged.
 * * -EINVAL - If @type is invalid or misc res has 0 capacity.
 * * -EBUSY - If max limit will be crossed or total usage will be more than the
 *            capacity.
 */
int misc_cg_try_charge(enum misc_res_type type, struct misc_cg *cg, u64 amount)
{
        struct misc_cg *i, *j;
        int ret;
        struct misc_res *res;
        u64 new_usage;

        if (!(valid_type(type) && cg && READ_ONCE(misc_res_capacity[type])))
                return -EINVAL;

        if (!amount)
                return 0;

        for (i = cg; i; i = parent_misc(i)) {
                res = &i->res[type];

                new_usage = atomic64_add_return(amount, &res->usage);
                if (new_usage > READ_ONCE(res->max) ||
                    new_usage > READ_ONCE(misc_res_capacity[type])) {
                        ret = -EBUSY;
                        goto err_charge;
                }
                misc_cg_update_watermark(res, new_usage);
        }
        return 0;

err_charge:
        misc_cg_event(type, i);

        for (j = cg; j != i; j = parent_misc(j))
                misc_cg_cancel_charge(type, j, amount);
        misc_cg_cancel_charge(type, i, amount);
        return ret;
}
EXPORT_SYMBOL_GPL(misc_cg_try_charge);

/**
 * misc_cg_uncharge() - Uncharge the misc cgroup.
 * @type: Misc res type which was charged.
 * @cg: Misc cgroup which will be uncharged.
 * @amount: Charged amount.
 *
 * Context: Any context.
 */
void misc_cg_uncharge(enum misc_res_type type, struct misc_cg *cg, u64 amount)
{
        struct misc_cg *i;

        if (!(amount && valid_type(type) && cg))
                return;

        for (i = cg; i; i = parent_misc(i))
                misc_cg_cancel_charge(type, i, amount);
}
EXPORT_SYMBOL_GPL(misc_cg_uncharge);

/**
 * misc_cg_max_show() - Show the misc cgroup max limit.
 * @sf: Interface file
 * @v: Arguments passed
 *
 * Context: Any context.
 * Return: 0 to denote successful print.
 */
static int misc_cg_max_show(struct seq_file *sf, void *v)
{
        int i;
        struct misc_cg *cg = css_misc(seq_css(sf));
        u64 max;

        for (i = 0; i < MISC_CG_RES_TYPES; i++) {
                if (READ_ONCE(misc_res_capacity[i])) {
                        max = READ_ONCE(cg->res[i].max);
                        if (max == MAX_NUM)
                                seq_printf(sf, "%s max\n", misc_res_name[i]);
                        else
                                seq_printf(sf, "%s %llu\n", misc_res_name[i],
                                           max);
                }
        }

        return 0;
}

/**
 * misc_cg_max_write() - Update the maximum limit of the cgroup.
 * @of: Handler for the file.
 * @buf: Data from the user. It should be either "max", 0, or a positive
 *       integer.
 * @nbytes: Number of bytes of the data.
 * @off: Offset in the file.
 *
 * User can pass data like:
 * echo sev 23 > misc.max, OR
 * echo sev max > misc.max
 *
 * Context: Any context.
 * Return:
 * * >= 0 - Number of bytes processed in the input.
 * * -EINVAL - If buf is not valid.
 * * -ERANGE - If number is bigger than the u64 capacity.
 */
static ssize_t misc_cg_max_write(struct kernfs_open_file *of, char *buf,
                                 size_t nbytes, loff_t off)
{
        struct misc_cg *cg;
        u64 max;
        int ret = 0, i;
        enum misc_res_type type = MISC_CG_RES_TYPES;
        char *token;

        buf = strstrip(buf);
        token = strsep(&buf, " ");

        if (!token || !buf)
                return -EINVAL;

        for (i = 0; i < MISC_CG_RES_TYPES; i++) {
                if (!strcmp(misc_res_name[i], token)) {
                        type = i;
                        break;
                }
        }

        if (type == MISC_CG_RES_TYPES)
                return -EINVAL;

        if (!strcmp(MAX_STR, buf)) {
                max = MAX_NUM;
        } else {
                ret = kstrtou64(buf, 0, &max);
                if (ret)
                        return ret;
        }

        cg = css_misc(of_css(of));

        if (READ_ONCE(misc_res_capacity[type]))
                WRITE_ONCE(cg->res[type].max, max);
        else
                ret = -EINVAL;

        return ret ? ret : nbytes;
}

/**
 * misc_cg_current_show() - Show the current usage of the misc cgroup.
 * @sf: Interface file
 * @v: Arguments passed
 *
 * Context: Any context.
 * Return: 0 to denote successful print.
 */
static int misc_cg_current_show(struct seq_file *sf, void *v)
{
        int i;
        u64 usage;
        struct misc_cg *cg = css_misc(seq_css(sf));

        for (i = 0; i < MISC_CG_RES_TYPES; i++) {
                usage = atomic64_read(&cg->res[i].usage);
                if (READ_ONCE(misc_res_capacity[i]) || usage)
                        seq_printf(sf, "%s %llu\n", misc_res_name[i], usage);
        }

        return 0;
}

/**
 * misc_cg_peak_show() - Show the peak usage of the misc cgroup.
 * @sf: Interface file
 * @v: Arguments passed
 *
 * Context: Any context.
 * Return: 0 to denote successful print.
 */
static int misc_cg_peak_show(struct seq_file *sf, void *v)
{
        int i;
        u64 watermark;
        struct misc_cg *cg = css_misc(seq_css(sf));

        for (i = 0; i < MISC_CG_RES_TYPES; i++) {
                watermark = atomic64_read(&cg->res[i].watermark);
                if (READ_ONCE(misc_res_capacity[i]) || watermark)
                        seq_printf(sf, "%s %llu\n", misc_res_name[i], watermark);
        }

        return 0;
}

/**
 * misc_cg_capacity_show() - Show the total capacity of misc res on the host.
 * @sf: Interface file
 * @v: Arguments passed
 *
 * Only present in the root cgroup directory.
 *
 * Context: Any context.
 * Return: 0 to denote successful print.
 */
static int misc_cg_capacity_show(struct seq_file *sf, void *v)
{
        int i;
        u64 cap;

        for (i = 0; i < MISC_CG_RES_TYPES; i++) {
                cap = READ_ONCE(misc_res_capacity[i]);
                if (cap)
                        seq_printf(sf, "%s %llu\n", misc_res_name[i], cap);
        }

        return 0;
}

static int __misc_events_show(struct seq_file *sf, bool local)
{
        struct misc_cg *cg = css_misc(seq_css(sf));
        u64 events;
        int i;

        for (i = 0; i < MISC_CG_RES_TYPES; i++) {
                if (local)
                        events = atomic64_read(&cg->res[i].events_local);
                else
                        events = atomic64_read(&cg->res[i].events);
                if (READ_ONCE(misc_res_capacity[i]) || events)
                        seq_printf(sf, "%s.max %llu\n", misc_res_name[i], events);
        }
        return 0;
}

static int misc_events_show(struct seq_file *sf, void *v)
{
        return __misc_events_show(sf, false);
}

static int misc_events_local_show(struct seq_file *sf, void *v)
{
        return __misc_events_show(sf, true);
}

/* Misc cgroup interface files */
static struct cftype misc_cg_files[] = {
        {
                .name = "max",
                .write = misc_cg_max_write,
                .seq_show = misc_cg_max_show,
                .flags = CFTYPE_NOT_ON_ROOT,
        },
        {
                .name = "current",
                .seq_show = misc_cg_current_show,
        },
        {
                .name = "peak",
                .seq_show = misc_cg_peak_show,
        },
        {
                .name = "capacity",
                .seq_show = misc_cg_capacity_show,
                .flags = CFTYPE_ONLY_ON_ROOT,
        },
        {
                .name = "events",
                .flags = CFTYPE_NOT_ON_ROOT,
                .file_offset = offsetof(struct misc_cg, events_file),
                .seq_show = misc_events_show,
        },
        {
                .name = "events.local",
                .flags = CFTYPE_NOT_ON_ROOT,
                .file_offset = offsetof(struct misc_cg, events_local_file),
                .seq_show = misc_events_local_show,
        },
        {}
};

/**
 * misc_cg_alloc() - Allocate misc cgroup.
 * @parent_css: Parent cgroup.
 *
 * Context: Process context.
 * Return:
 * * struct cgroup_subsys_state* - css of the allocated cgroup.
 * * ERR_PTR(-ENOMEM) - No memory available to allocate.
 */
static struct cgroup_subsys_state *
misc_cg_alloc(struct cgroup_subsys_state *parent_css)
{
        enum misc_res_type i;
        struct misc_cg *cg;

        if (!parent_css) {
                cg = &root_cg;
        } else {
                cg = kzalloc_obj(*cg);
                if (!cg)
                        return ERR_PTR(-ENOMEM);
        }

        for (i = 0; i < MISC_CG_RES_TYPES; i++) {
                WRITE_ONCE(cg->res[i].max, MAX_NUM);
                atomic64_set(&cg->res[i].usage, 0);
        }

        return &cg->css;
}

/**
 * misc_cg_free() - Free the misc cgroup.
 * @css: cgroup subsys object.
 *
 * Context: Any context.
 */
static void misc_cg_free(struct cgroup_subsys_state *css)
{
        kfree(css_misc(css));
}

/* Cgroup controller callbacks */
struct cgroup_subsys misc_cgrp_subsys = {
        .css_alloc = misc_cg_alloc,
        .css_free = misc_cg_free,
        .legacy_cftypes = misc_cg_files,
        .dfl_cftypes = misc_cg_files,
};