#include <linux/module.h>
#include <linux/init.h>
#include <linux/sched/task.h>
#include <linux/circ_buf.h>
#include <linux/fs.h>
#include <linux/cdev.h>
#include <linux/completion.h>
#include <linux/slab.h>
#include <linux/syscalls.h>
#include <linux/pid.h>
#include <linux/kprobes.h>
#include <linux/anon_inodes.h>

typedef unsigned long (*kallsyms_lookup_name_t)(const char *name);
kallsyms_lookup_name_t gei_kallsyms_lookup_name;
typedef int (*pidfd_create_t)(struct pid *pid, int flags);
pidfd_create_t pidfd_create;
static struct kprobe kln_kp = {
  .symbol_name = "kallsyms_lookup_name"
};

struct gei_work {
  struct task_struct *execing_task;
  struct completion userspace_done;
  refcount_t usage;
};
struct gei_work dead_sentinel;
#define SUBMIT_QUEUE_SIZE 8
struct gei_work *submit_queue[SUBMIT_QUEUE_SIZE];
unsigned long submit_queue_head = 0;
unsigned long submit_queue_tail = 0;
DECLARE_WAIT_QUEUE_HEAD(submit_read_q);
DECLARE_WAIT_QUEUE_HEAD(submit_write_q);
DEFINE_SPINLOCK(submit_read_lock);
DEFINE_SPINLOCK(submit_write_lock);
DEFINE_MUTEX(clientop_lock);

struct gei_user_msg {
  int gei_fd;
  int pid_fd;
  pid_t pid;
};

int gei_clients = 0;

static inline struct gei_work *get_gei_work(struct gei_work *work) {
  refcount_inc(&work->usage);
  return work;
}
static inline void put_gei_work(struct gei_work *work) {
  if (refcount_dec_and_test(&work->usage)) {
    put_task_struct(work->execing_task);
    kfree(work);
  }
}

static int geifd_release(struct inode *inode, struct file *file) {
  struct gei_work *work = (struct gei_work *)file->private_data;
  complete(&work->userspace_done);
  put_gei_work(work);
  return 0;
}
const struct file_operations geifd_fops = {
  .release = geifd_release,
};

static void gei_ft_cb(unsigned long ip, unsigned long parent_ip,
                      struct ftrace_ops *op, struct pt_regs *regs) {
  struct gei_work* work = kmalloc(sizeof(struct gei_work), GFP_KERNEL);
  work->execing_task = get_task_struct(current);
  init_completion(&work->userspace_done);
  refcount_set(&work->usage, 1); /* this thread and the */
  while (1) {
    spin_lock(&submit_write_lock);
    unsigned long head = submit_queue_head;
    unsigned long tail = READ_ONCE(submit_queue_tail);
    if (CIRC_SPACE(head, tail, SUBMIT_QUEUE_SIZE) >= 1) {
      /* nobody else can see this yet (smp_store_release) - so it's
       * safe to write w/o taking the read spinlock */
      submit_queue[head] = get_gei_work(work);
      smp_store_release(&submit_queue_head,
                        (submit_queue_head + 1) & (SUBMIT_QUEUE_SIZE - 1));
      wake_up(&submit_read_q);
      spin_unlock(&submit_write_lock);
      break;
    } else {
      spin_unlock(&submit_write_lock);
      if (wait_event_killable(submit_write_q,
            CIRC_SPACE(submit_queue_head,
                       READ_ONCE(submit_queue_tail),
                       SUBMIT_QUEUE_SIZE))
          == -ERESTARTSYS) {
        /* TODO: Ensure the process is not allowed to start if it
           receives a non-fatal signal while waiting to enqueue the
           work item */
        put_gei_work(work);
        return;
      }
    }
  }
  if (wait_for_completion_killable(&work->userspace_done) == -ERESTARTSYS) {
    /* TODO: Ensure the process is not allowed to start if it receives
       a non-fatal signal while waiting for userspace to handle the
       work item */
    put_gei_work(work);
    return;
  }
  put_gei_work(work);
}
struct ftrace_ops gei_ft_ops = {
  .func = gei_ft_cb,
};

static int gei_open(struct inode *ino, struct file *filp) {
  if (mutex_lock_interruptible(&clientop_lock)) {
    return -ERESTARTSYS;
  }
  printk("gei: found daemon\n");
  if (gei_clients++ == 0) {
    printk("gei: starting");
    if (register_ftrace_function(&gei_ft_ops) < 0) {
      return -EBUSY;
    }
  }
  printk("gei: total clients %d\n", gei_clients);
  mutex_unlock(&clientop_lock);
  return 0;
}
static int gei_release(struct inode *ino, struct file *filp) {
  if (mutex_lock_interruptible(&clientop_lock)) {
    return -ERESTARTSYS;
  }
  printk("gei: daemon left\n");
  if (--gei_clients == 0) {
    printk("gei: stopping");
    unregister_ftrace_function(&gei_ft_ops);
  }
  printk("gei: total clients %d\n", gei_clients);
  mutex_unlock(&clientop_lock);
  return 0;
}

static ssize_t gei_read(struct file *filp,
                        char __user *buff, size_t count,
                        loff_t *offp) {
  struct gei_work *work = NULL;

  if (count != sizeof(struct gei_user_msg)) { return -EINVAL; }
  while (1) {
    spin_lock(&submit_read_lock);
    unsigned long head = smp_load_acquire(&submit_queue_head);
    unsigned long tail = submit_queue_tail;
    if (CIRC_CNT(head, tail, SUBMIT_QUEUE_SIZE) >= 1) {
      work = submit_queue[tail];
      smp_store_release(&submit_queue_tail,
                        (tail + 1) & (SUBMIT_QUEUE_SIZE - 1));
      spin_unlock(&submit_read_lock);
      break;
    } else {
      spin_unlock(&submit_read_lock);
      if (wait_event_killable(submit_read_q,
            CIRC_CNT(smp_load_acquire(&submit_queue_head),
                     submit_queue_tail, SUBMIT_QUEUE_SIZE) >= 1)
          == -ERESTARTSYS) {
        return -ERESTARTSYS;
      }
    }
  }
  struct pid *pid = get_task_pid(work->execing_task, PIDTYPE_PID);
  struct gei_user_msg msg = {
    .gei_fd = anon_inode_getfd("[gei]", &geifd_fops, work, O_CLOEXEC),
    .pid_fd = pidfd_create(pid, 0),
    .pid = pid_vnr(pid),
  };
  put_pid(pid);
  ssize_t ret = copy_to_user(buff, &msg, sizeof(struct gei_user_msg));
  return sizeof(struct gei_user_msg) - ret;
}
static const struct file_operations gei_fops = {
  .owner = THIS_MODULE,
  .open = gei_open,
  .read = gei_read,
  .release = gei_release,
};

static dev_t dev;
static struct class *gei_class;
static struct device *gei_device;
static struct cdev cdev;

static int __init gei_init(void) {
  int ret;

  pr_debug("untroubled exec? no more\n");

  if ((ret = alloc_chrdev_region(&dev, 0, 1, "global_execve_interceptor"))) {
    return ret;
  }
  cdev_init(&cdev, &gei_fops);
  if ((ret = cdev_add(&cdev, dev, 1) < 0)) {
    unregister_chrdev_region(dev, 1);
    return ret;
  }
  gei_class = class_create(THIS_MODULE, "global_execve_interceptor");
  gei_device = device_create(gei_class, NULL,
                             dev, NULL,
                             "global_execve_interceptor");
  
  register_kprobe(&kln_kp);
  gei_kallsyms_lookup_name = (kallsyms_lookup_name_t)kln_kp.addr;
  unregister_kprobe(&kln_kp);
  pidfd_create = (pidfd_create_t)gei_kallsyms_lookup_name("pidfd_create");

 if (ftrace_set_filter(&gei_ft_ops, "finalize_exec",
                       strlen("finalize_exec"), 0) < 0) {
    goto err;
 }

  return 0;
err:
  device_destroy(gei_class, dev);
  class_destroy(gei_class);
  unregister_chrdev_region(dev, 1);
  return -EINVAL;
}
static void __exit gei_exit(void) {
  pr_debug("we kindly return your exec\n");
  if (gei_clients) { unregister_ftrace_function(&gei_ft_ops); }
  device_destroy(gei_class, dev);
  class_destroy(gei_class);
  unregister_chrdev_region(dev, 1);
}
module_init(gei_init);
module_exit(gei_exit);
MODULE_LICENSE("GPL");
MODULE_AUTHOR("semiotics");