// SPDX-License-Identifier: GPL-2.0
/*
 * KUnit userspace memory allocation resource management.
 */
#include <kunit/resource.h>
#include <kunit/test.h>
#include <linux/kthread.h>
#include <linux/mm.h>

struct kunit_vm_mmap_resource {
 unsigned long addr;
 size_t size;
};

/* vm_mmap() arguments */
struct kunit_vm_mmap_params {
 struct file *file;
 unsigned long addr;
 unsigned long len;
 unsigned long prot;
 unsigned long flag;
 unsigned long offset;
};

int kunit_attach_mm(void)
{
 struct mm_struct *mm;

 if (current->mm)
  return 0;

 /* arch_pick_mmap_layout() is only sane with MMU systems. */
 if (!IS_ENABLED(CONFIG_MMU))
  return -EINVAL;

 mm = mm_alloc();
 if (!mm)
  return -ENOMEM;

 /* Define the task size. */
 mm->task_size = TASK_SIZE;

 /* Make sure we can allocate new VMAs. */
 arch_pick_mmap_layout(mm, ¤t->signal->rlim[RLIMIT_STACK]);

 /* Attach the mm. It will be cleaned up when the process dies. */
 kthread_use_mm(mm);

 return 0;
}
EXPORT_SYMBOL_GPL(kunit_attach_mm);

static int kunit_vm_mmap_init(struct kunit_resource *res, void *context)
{
 struct kunit_vm_mmap_params *p = context;
 struct kunit_vm_mmap_resource vres;
 int ret;

 ret = kunit_attach_mm();
 if (ret)
  return ret;

 vres.size = p->len;
 vres.addr = vm_mmap(p->file, p->addr, p->len, p->prot, p->flag, p->offset);
 if (!vres.addr)
  return -ENOMEM;
 res->data = kmemdup(&vres, sizeof(vres), GFP_KERNEL);
 if (!res->data) {
  vm_munmap(vres.addr, vres.size);
  return -ENOMEM;
 }

 return 0;
}

static void kunit_vm_mmap_free(struct kunit_resource *res)
{
 struct kunit_vm_mmap_resource *vres = res->data;

 /*
 * Since this is executed from the test monitoring process,
 * the test's mm has already been torn down. We don't need
 * to run vm_munmap(vres->addr, vres->size), only clean up
 * the vres.
 */

 kfree(vres);
 res->data = NULL;
}

unsigned long kunit_vm_mmap(struct kunit *test, struct file *file,
       unsigned long addr, unsigned long len,
       unsigned long prot, unsigned long flag,
       unsigned long offset)
{
 struct kunit_vm_mmap_params params = {
  .file = file,
  .addr = addr,
  .len = len,
  .prot = prot,
  .flag = flag,
  .offset = offset,
 };
 struct kunit_vm_mmap_resource *vres;

 vres = kunit_alloc_resource(test,
        kunit_vm_mmap_init,
        kunit_vm_mmap_free,
        GFP_KERNEL,
        ¶ms);
 if (vres)
  return vres->addr;
 return 0;
}
EXPORT_SYMBOL_GPL(kunit_vm_mmap);

MODULE_IMPORT_NS("EXPORTED_FOR_KUNIT_TESTING");