// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright (C) 2025 - Google Inc
 * Author: Mostafa Saleh <smostafa@google.com>
 * IOMMU API debug page alloc sanitizer
 */
#include <linux/atomic.h>
#include <linux/iommu.h>
#include <linux/iommu-debug-pagealloc.h>
#include <linux/kernel.h>
#include <linux/page_ext.h>
#include <linux/page_owner.h>

#include "iommu-priv.h"

static bool needed;
DEFINE_STATIC_KEY_FALSE(iommu_debug_initialized);

struct iommu_debug_metadata {
	atomic_t ref;
};

static __init bool need_iommu_debug(void)
{
	return needed;
}

struct page_ext_operations page_iommu_debug_ops = {
	.size = sizeof(struct iommu_debug_metadata),
	.need = need_iommu_debug,
};

static struct iommu_debug_metadata *get_iommu_data(struct page_ext *page_ext)
{
	return page_ext_data(page_ext, &page_iommu_debug_ops);
}

static void iommu_debug_inc_page(phys_addr_t phys)
{
	struct page_ext *page_ext = page_ext_from_phys(phys);
	struct iommu_debug_metadata *d;

	if (!page_ext)
		return;

	d = get_iommu_data(page_ext);
	WARN_ON(atomic_inc_return_relaxed(&d->ref) <= 0);
	page_ext_put(page_ext);
}

static void iommu_debug_dec_page(phys_addr_t phys)
{
	struct page_ext *page_ext = page_ext_from_phys(phys);
	struct iommu_debug_metadata *d;

	if (!page_ext)
		return;

	d = get_iommu_data(page_ext);
	WARN_ON(atomic_dec_return_relaxed(&d->ref) < 0);
	page_ext_put(page_ext);
}

/*
 * IOMMU page size doesn't have to match the CPU page size. So, we use
 * the smallest IOMMU page size to refcount the pages in the vmemmap.
 * That is important as both map and unmap has to use the same page size
 * to update the refcount to avoid double counting the same page.
 * And as we can't know from iommu_unmap() what was the original page size
 * used for map, we just use the minimum supported one for both.
 */
static size_t iommu_debug_page_size(struct iommu_domain *domain)
{
	return 1UL << __ffs(domain->pgsize_bitmap);
}

static bool iommu_debug_page_count(const struct page *page)
{
	unsigned int ref;
	struct page_ext *page_ext = page_ext_get(page);
	struct iommu_debug_metadata *d = get_iommu_data(page_ext);

	ref = atomic_read(&d->ref);
	page_ext_put(page_ext);
	return ref != 0;
}

void __iommu_debug_check_unmapped(const struct page *page, int numpages)
{
	while (numpages--) {
		if (WARN_ON(iommu_debug_page_count(page))) {
			pr_warn("iommu: Detected page leak!\n");
			dump_page_owner(page);
		}
		page++;
	}
}

void __iommu_debug_map(struct iommu_domain *domain, phys_addr_t phys, size_t size)
{
	size_t off, end;
	size_t page_size = iommu_debug_page_size(domain);

	if (WARN_ON(!phys || check_add_overflow(phys, size, &end)))
		return;

	for (off = 0 ; off < size ; off += page_size)
		iommu_debug_inc_page(phys + off);
}

static void __iommu_debug_update_iova(struct iommu_domain *domain,
				      unsigned long iova, size_t size, bool inc)
{
	size_t off, end;
	size_t page_size = iommu_debug_page_size(domain);

	if (WARN_ON(check_add_overflow(iova, size, &end)))
		return;

	for (off = 0 ; off < size ; off += page_size) {
		phys_addr_t phys = iommu_iova_to_phys(domain, iova + off);

		if (!phys)
			continue;

		if (inc)
			iommu_debug_inc_page(phys);
		else
			iommu_debug_dec_page(phys);
	}
}

void __iommu_debug_unmap_begin(struct iommu_domain *domain,
			       unsigned long iova, size_t size)
{
	__iommu_debug_update_iova(domain, iova, size, false);
}

void __iommu_debug_unmap_end(struct iommu_domain *domain,
			     unsigned long iova, size_t size,
			     size_t unmapped)
{
	if ((unmapped == size) || WARN_ON_ONCE(unmapped > size))
		return;

	/* If unmap failed, re-increment the refcount. */
	__iommu_debug_update_iova(domain, iova + unmapped,
				  size - unmapped, true);
}

void iommu_debug_init(void)
{
	if (!needed)
		return;

	pr_info("iommu: Debugging page allocations, expect overhead or disable iommu.debug_pagealloc");
	static_branch_enable(&iommu_debug_initialized);
}

static int __init iommu_debug_pagealloc(char *str)
{
	return kstrtobool(str, &needed);
}
early_param("iommu.debug_pagealloc", iommu_debug_pagealloc);