// 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
};

/* 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_res_total_usage() - Get the current total usage of the resource.
 * @type: misc res type.
 *
 * Context: Any context.
 * Return: Current total usage of the resource.
 */
u64 misc_cg_res_total_usage(enum misc_res_type type)
{
	if (valid_type(type))
		return atomic64_read(&root_cg.res[type].usage);

	return 0;
}
EXPORT_SYMBOL_GPL(misc_cg_res_total_usage);

/**
 * 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(sizeof(*cg), GFP_KERNEL);
		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,
};