mop3/kernel/proc/procgroup.c

#include <fs/vfs.h>
#include <id/id_alloc.h>
#include <libk/rbtree.h>
#include <libk/std.h>
#include <libk/string.h>
#include <limine/requests.h>
#include <mm/malloc.h>
#include <mm/pmm.h>
#include <proc/env.h>
#include <proc/proc.h>
#include <proc/procgroup.h>
#include <proc/reschedule.h>
#include <sync/spin_lock.h>
#include <sys/debug.h>
#include <sys/mm.h>

#define PGIDS_MAX 1024

static struct rb_node_link* procgroup_tree = NULL;

static struct id_alloc pgid_alloc;

void procgroup_pgid_alloc_init(void) { id_alloc_init(&pgid_alloc, PGIDS_MAX); }

struct procgroup* procgroup_find(int pgid) {
  struct procgroup* procgroup = NULL;

  rbtree_find(struct procgroup, &procgroup_tree, pgid, procgroup, procgroup_tree_link, pgid);

  return procgroup;
}

uintptr_t procgroup_map(struct procgroup* procgroup, uintptr_t vaddr, size_t pages, uint32_t flags,
                        uintptr_t* out_paddr) {
  vaddr = (vaddr == 0) ? procgroup->map_base : vaddr;

  struct proc_mapping* mapping = malloc(sizeof(*mapping));

  if (mapping == NULL) {
    return 0;
  }

  uintptr_t paddr = pmm_alloc(pages);

  if (paddr == PMM_ALLOC_ERR) {
    free(mapping);
    return 0;
  }

  if (out_paddr != NULL)
    *out_paddr = paddr;

  mapping->paddr = paddr;
  mapping->vaddr = vaddr;
  mapping->size = pages * PAGE_SIZE;

  procgroup->map_base += pages * PAGE_SIZE;

  list_append(procgroup->mappings, &mapping->proc_mappings_link);

  for (uintptr_t vpage = vaddr, ppage = paddr; vpage < vaddr + pages * PAGE_SIZE;
       vpage += PAGE_SIZE, ppage += PAGE_SIZE) {
    mm_map_page(&procgroup->pd, ppage, vpage, flags);
  }

  return vaddr;
}

bool procgroup_unmap(struct procgroup* procgroup, uintptr_t start_vaddr, size_t pages) {
  size_t unmap_size = pages * PAGE_SIZE;
  uintptr_t end_vaddr = start_vaddr + unmap_size;

  struct list_node_link *mapping_link, *mapping_link_tmp;

  bool used_tail_mapping = false;

  struct proc_mapping* tail_mapping = malloc(sizeof(*tail_mapping));
  if (tail_mapping == NULL)
    return false;

  list_foreach(procgroup->mappings, mapping_link, mapping_link_tmp) {
    struct proc_mapping* mapping =
        list_entry(mapping_link, struct proc_mapping, proc_mappings_link);

    uintptr_t m_start = mapping->vaddr;
    uintptr_t m_end = mapping->vaddr + mapping->size;

    /* check overlap */
    if ((start_vaddr < m_end) && (end_vaddr > mapping->vaddr)) {
      uintptr_t free_vstart = (start_vaddr > m_start) ? start_vaddr : m_start;
      uintptr_t free_vend = (end_vaddr < m_end) ? end_vaddr : m_end;
      size_t free_size = free_vend - free_vstart;

      uintptr_t ppage_to_free = mapping->paddr + (free_vstart - m_start);
      pmm_free(ppage_to_free, free_size / PAGE_SIZE);

      /* split in the middle */
      if ((start_vaddr > m_start) && (end_vaddr < m_end)) {
        tail_mapping->vaddr = end_vaddr;
        tail_mapping->paddr = mapping->paddr + (end_vaddr - m_start);
        tail_mapping->size = m_end - end_vaddr;

        mapping->size = start_vaddr - m_start;

        list_insert_after(procgroup->mappings, &mapping->proc_mappings_link,
                          &tail_mapping->proc_mappings_link);

        used_tail_mapping = true;

        break;
      } else if ((start_vaddr <= m_start) && (end_vaddr < m_end)) { /* shrink left */
        size_t diff = end_vaddr - m_start;
        mapping->vaddr += diff;
        mapping->paddr += diff;
        mapping->size -= diff;
      } else if ((start_vaddr > m_start) && (end_vaddr >= m_end)) { /* shrink right */
        mapping->size = start_vaddr - m_start;
      } else { /* full overlap */
        list_remove(procgroup->mappings, &mapping->proc_mappings_link);
        free(mapping);
      }
    }
  }

  if (!used_tail_mapping)
    free(tail_mapping);

  for (uintptr_t vpage = start_vaddr; vpage < end_vaddr; vpage += PAGE_SIZE) {
    mm_unmap_page(&procgroup->pd, vpage);
  }

  return true;
}

struct procgroup* procgroup_create(void) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  struct procgroup* procgroup = malloc(sizeof(*procgroup));
  if (procgroup == NULL) {
    return NULL;
  }

  memset(procgroup, 0, sizeof(*procgroup));

  if (!id_alloc_init(&procgroup->rid_alloc, PROCGROUP_RESOURCES_MAX)) {
    free(procgroup);
    return NULL;
  }

  procgroup->pgid = id_alloc(&pgid_alloc);

  if (procgroup->pgid < 0) {
    id_alloc_fini(&procgroup->rid_alloc);
    free(procgroup);
    return NULL;
  }

  procgroup->memb_proc_tree = NULL;
  procgroup->pd.cr3_paddr = mm_alloc_user_pd_phys();
  procgroup->map_base = PROC_MAP_BASE;

  procgroup->uvaddr_cmdline = procgroup_map(procgroup, PROC_CMDLINE_BASE, 1,
                                            MM_PG_PRESENT | MM_PG_USER, &procgroup->paddr_cmdline);
  memset((void*)((uintptr_t)hhdm->offset + procgroup->paddr_cmdline), 0, PROCGROUP_CMDLINE_MAX);

  if (proc_create_resource_mail(procgroup) == NULL) {
    id_alloc_fini(&procgroup->rid_alloc);
    free(procgroup);
    return NULL;
  }

  if (proc_create_resource_stream(procgroup) == NULL) {
    id_alloc_fini(&procgroup->rid_alloc);
    free(procgroup);
    return NULL;
  }

  if (proc_create_resource_stream(procgroup) == NULL) {
    id_alloc_fini(&procgroup->rid_alloc);
    free(procgroup);
    return NULL;
  }

  rbtree_insert(struct procgroup, &procgroup_tree, &procgroup->procgroup_tree_link,
                procgroup_tree_link, pgid);

  return procgroup;
}

void procgroup_attach(struct procgroup* procgroup, struct proc* proc) {
  rbtree_insert(struct proc, &procgroup->memb_proc_tree, &proc->procgroup_memb_tree_link,
                procgroup_memb_tree_link, pid);
}

static void procgroup_delete(struct procgroup* procgroup, struct reschedule_ctx* rctx) {
  rbtree_delete(&procgroup_tree, &procgroup->procgroup_tree_link);

  /* delete resources */
  struct rb_node_link* rnode;
  rbtree_first(&procgroup->resource_tree, rnode);
  while (rnode) {
    struct rb_node_link* next;
    rbtree_next(rnode, next);

    struct proc_resource* resource = rbtree_entry(rnode, struct proc_resource, resource_tree_link);

    rnode = next;

    proc_delete_resource(procgroup, resource, rctx);
  }

  /* delete mappings */
  struct list_node_link *mapping_link, *mapping_link_tmp;
  list_foreach(procgroup->mappings, mapping_link, mapping_link_tmp) {
    struct proc_mapping* mapping =
        list_entry(mapping_link, struct proc_mapping, proc_mappings_link);

    pmm_free(mapping->paddr, mapping->size / PAGE_SIZE);
    free(mapping);
  }

  proc_env_cleanup(procgroup);

  pmm_free(procgroup->paddr_cmdline, 1);

  pmm_free(procgroup->pd.cr3_paddr, 1);

  free(procgroup->tls.tls_tmpl);

  id_alloc_fini(&procgroup->rid_alloc);
  id_free(&pgid_alloc, procgroup->pgid);

  free(procgroup);
}

void procgroup_detach(struct procgroup* procgroup, struct proc* proc, struct reschedule_ctx* rctx) {
  rbtree_delete(&procgroup->memb_proc_tree, &proc->procgroup_memb_tree_link);
  struct rb_node_link* memb_tree = procgroup->memb_proc_tree;

  if (memb_tree == NULL) {
    procgroup_delete(procgroup, rctx);
  }
}