Kernel: More out of memory condition checks

10 files changed, 76 insertions, 56 deletions

diff --git a/kernel/arch/i386/mmu.c b/kernel/arch/i386/mmu.c
index de4c3c4..a505639 100644
--- a/kernel/arch/i386/mmu.c
+++ b/kernel/arch/i386/mmu.c
@@ -58,7 +58,7 @@ void *ksbrk(size_t s) {
   }
   if (!mmu_allocate_shared_kernel_region((void *)rc,
                                          (data_end - (uintptr_t)rc))) {
-    return NULL;
+    return (void *)-1;
   }
   assert(((uintptr_t)rc % PAGE_SIZE) == 0);
   return (void *)rc;
@@ -372,7 +372,6 @@ void *mmu_map_user_frames(void *const ptr, size_t s) {
     p->rw = rw;
     p->user = !is_kernel;
     p->frame = (uintptr_t)(ptr + i * 0x1000) / 0x1000;
-    kprintf("mapped user frame: %x\n", p->frame);
     (void)write_to_frame((uintptr_t)ptr + i * 0x1000, 1);
   }
   return r;
@@ -383,7 +382,10 @@ void *mmu_map_frames(void *const ptr, size_t s) {
   size_t num_pages = s / 0x1000;
   for (size_t i = 0; i <= num_pages; i++) {
     Page *p = get_page((void *)(r + i * 0x1000), NULL, PAGE_ALLOCATE, 0);
-    assert(p);
+    if (!p) {
+      mmu_free_address_range(r, i * 0x1000, NULL);
+      return NULL;
+    }
     int rw = 1;
     int is_kernel = 1;
     p->present = 1;
@@ -463,36 +465,24 @@ void mmu_free_address_range(void *ptr, size_t length, PageDirectory *pd) {
   }
 }
 
-void mmu_map_directories(void *dst, PageDirectory *d, void *src,
-                         PageDirectory *s, size_t length) {
-  d = (!d) ? get_active_pagedirectory() : d;
-  s = (!s) ? get_active_pagedirectory() : s;
-  size_t num_pages = (u32)align_page((void *)length) / 0x1000;
-  for (size_t i = 0; i < num_pages; i++, dst += 0x1000, src += 0x1000) {
-    Page *p = get_page(dst, d, PAGE_ALLOCATE, 1);
-    assert(p);
-    p->present = 1;
-    p->rw = 1;
-    p->user = 1;
-    void *physical = virtual_to_physical(src, s);
-    p->frame = (u32)physical / PAGE_SIZE;
-  }
-  flush_tlb();
-}
-
-void mmu_map_physical(void *dst, PageDirectory *d, void *physical,
-                      size_t length) {
+int mmu_map_physical(void *dst, PageDirectory *d, void *physical,
+                     size_t length) {
+  void *const dst_orig = dst;
   d = (!d) ? get_active_pagedirectory() : d;
   size_t num_pages = (u32)align_page((void *)length) / 0x1000;
   for (size_t i = 0; i < num_pages; i++, dst += 0x1000, physical += 0x1000) {
     Page *p = get_page(dst, d, PAGE_ALLOCATE, 1);
-    assert(p);
+    if (!p) {
+      mmu_free_address_range(dst_orig, i * 0x1000, d);
+      return 0;
+    }
     p->present = 1;
     p->rw = 1;
     p->user = 1;
     p->frame = (uintptr_t)physical / PAGE_SIZE;
     (void)write_to_frame((uintptr_t)physical, 1);
   }
+  return 1;
 }
 
 struct PhysVirtMap {
@@ -542,6 +532,9 @@ void *virtual_to_physical(void *address, PageDirectory *directory) {
     directory = get_active_pagedirectory();
   }
   Page *p = get_page((void *)address, directory, PAGE_NO_ALLOCATE, 0);
+  if (!p) {
+    return NULL;
+  }
   return (void *)((u32)p->frame * 0x1000) + (((uintptr_t)address) & 0xFFF);
 }
 
diff --git a/kernel/drivers/ahci.c b/kernel/drivers/ahci.c
index f655709..1fad0a3 100644
--- a/kernel/drivers/ahci.c
+++ b/kernel/drivers/ahci.c
@@ -496,10 +496,10 @@ void add_devfs_drive_file(u8 port) {
   inode->inode_num = port;
 }
 
-void ahci_init(void) {
+int ahci_init(void) {
   struct PCI_DEVICE device;
   if (!pci_devices_by_id(0x01, 0x06, &device)) {
-    return;
+    return 0;
   }
   kprintf("vendor: %x\n", device.vendor);
   kprintf("device: %x\n", device.device);
@@ -509,6 +509,9 @@ void ahci_init(void) {
   pci_get_bar(&device, 5, &bar);
 
   u8 *HBA_base = mmu_map_frames((void *)bar.address, bar.size);
+  if (!HBA_base) {
+    return 0;
+  }
   hba = (volatile struct HBA_MEM *)(HBA_base);
   for (u8 i = 0; i < 32; i++) {
     if (!((hba->pi >> i) & 1)) {
@@ -535,4 +538,5 @@ void ahci_init(void) {
       break;
     }
   }
+  return 1;
 }
diff --git a/kernel/drivers/ahci.h b/kernel/drivers/ahci.h
index d075f07..6ec9a97 100644
--- a/kernel/drivers/ahci.h
+++ b/kernel/drivers/ahci.h
@@ -1 +1 @@
-void ahci_init(void);
+int ahci_init(void);
diff --git a/kernel/drivers/vbe.c b/kernel/drivers/vbe.c
index 9c6fa4a..e67a3d9 100644
--- a/kernel/drivers/vbe.c
+++ b/kernel/drivers/vbe.c
@@ -24,7 +24,7 @@ struct DISPLAY_INFO {
 
 struct DISPLAY_INFO vbe_info;
 
-void display_driver_init(multiboot_info_t *mbi) {
+int display_driver_init(multiboot_info_t *mbi) {
   assert(CHECK_FLAG(mbi->flags, 12));
   framebuffer_width = mbi->framebuffer_width;
   framebuffer_height = mbi->framebuffer_height;
@@ -37,10 +37,14 @@ void display_driver_init(multiboot_info_t *mbi) {
   framebuffer_physical = mbi->framebuffer_addr;
   framebuffer =
       mmu_map_frames((void *)(u32)mbi->framebuffer_addr, framebuffer_size);
+  if(!framebuffer) {
+    return 0;
+  }
 
   vbe_info.width = framebuffer_width;
   vbe_info.height = framebuffer_height;
   vbe_info.bpp = mbi->framebuffer_bpp;
+  return 1;
 }
 
 vfs_vm_object_t *vbe_get_vm_object(u64 length, u64 offset, vfs_fd_t *fd) {
diff --git a/kernel/drivers/vbe.h b/kernel/drivers/vbe.h
index 4df2221..369f384 100644
--- a/kernel/drivers/vbe.h
+++ b/kernel/drivers/vbe.h
@@ -1,7 +1,7 @@
 #ifndef VBE_H
 #define VBE_H
 #include <multiboot.h>
-void display_driver_init(multiboot_info_t *mb);
+int display_driver_init(multiboot_info_t *mb);
 void display_driver_cross(multiboot_info_t *mbi);
 void add_vbe_device(void);
 #endif // VBE_H
diff --git a/kernel/fs/ext2.c b/kernel/fs/ext2.c
index b11bc9c..046408c 100644
--- a/kernel/fs/ext2.c
+++ b/kernel/fs/ext2.c
@@ -800,8 +800,15 @@ int ext2_create_file(const char *path, int mode) {
 }
 
 vfs_inode_t *ext2_mount(void) {
-  int fd = vfs_open("/dev/sda", O_RDWR, 0);
   cache = kcalloc(3000, sizeof(struct BLOCK_CACHE));
+  if (!cache) {
+    return NULL;
+  }
+  int fd = vfs_open("/dev/sda", O_RDWR, 0);
+  if(0 > fd) {
+    kfree(cache);
+    return NULL;
+  }
   // TODO: Can this be done better? Maybe create a seperate function in
   // the VFS?
   mount_fd = get_vfs_fd(fd, NULL);
diff --git a/kernel/includes/mmu.h b/kernel/includes/mmu.h
index a248b8a..319169f 100644
--- a/kernel/includes/mmu.h
+++ b/kernel/includes/mmu.h
@@ -45,11 +45,9 @@ int mmu_allocate_shared_kernel_region(void *rc, size_t n);
 void *mmu_find_unallocated_virtual_range(void *addr, size_t length);
 void mmu_remove_virtual_physical_address_mapping(void *ptr, size_t length);
 void mmu_free_address_range(void *ptr, size_t length, PageDirectory *pd);
-void mmu_map_directories(void *dst, PageDirectory *d, void *src,
-                         PageDirectory *s, size_t length);
 u32 mmu_get_number_of_allocated_frames(void);
 void *mmu_map_frames(void *ptr, size_t s);
-void mmu_map_physical(void *dst, PageDirectory *d, void *physical,
+int mmu_map_physical(void *dst, PageDirectory *d, void *physical,
                       size_t length);
 void mmu_free_pages(void *a, u32 n);
 void *mmu_map_user_frames(void *const ptr, size_t s);
diff --git a/kernel/init/kernel.c b/kernel/init/kernel.c
index 642f4b8..2b96359 100644
--- a/kernel/init/kernel.c
+++ b/kernel/init/kernel.c
@@ -58,6 +58,7 @@ void kernel_main(u32 kernel_end, unsigned long magic, unsigned long addr,
   paging_init(memsize_kb, mb);
   klog("Paging Initalized", LOG_SUCCESS);
   mb = mmu_map_frames((multiboot_info_t *)addr, sizeof(multiboot_info_t));
+  assert(mb);
 
   gdt_init();
   klog("GDT Initalized", LOG_SUCCESS);
@@ -85,8 +86,10 @@ void kernel_main(u32 kernel_end, unsigned long magic, unsigned long addr,
   klog("PS2 Keyboard driver installed", LOG_SUCCESS);
 
   vfs_mount("/dev", devfs_mount());
-  ahci_init();
-  vfs_mount("/", ext2_mount());
+  assert(ahci_init());
+  vfs_inode_t *ext2_mount_point = ext2_mount();
+  assert(ext2_mount_point);
+  vfs_mount("/", ext2_mount_point);
 
   add_stdout();
   add_serial();
@@ -108,7 +111,7 @@ void kernel_main(u32 kernel_end, unsigned long magic, unsigned long addr,
 
   gen_ipv4(&ip_address, 10, 0, 2, 15);
 
-  display_driver_init(mb);
+  assert(display_driver_init(mb));
   add_vbe_device();
   int pid;
   if (0 == (pid = fork())) {
diff --git a/kernel/kmalloc.c b/kernel/kmalloc.c
index 087609e..12e8232 100644
--- a/kernel/kmalloc.c
+++ b/kernel/kmalloc.c
@@ -12,7 +12,11 @@
 void *kmalloc_align(size_t s, void **physical) {
   // TODO: It should reuse virtual regions so that it does not run out
   // of address space.
-  return ksbrk_physical(s, physical);
+  void *rc;
+  if ((void *)-1 == (rc = ksbrk_physical(s, physical))) {
+    return NULL;
+  }
+  return rc;
 }
 
 void kmalloc_align_free(void *p, size_t s) {
diff --git a/kernel/sched/scheduler.c b/kernel/sched/scheduler.c
index 57ab33b..7d5740b 100644
--- a/kernel/sched/scheduler.c
+++ b/kernel/sched/scheduler.c
@@ -257,9 +257,13 @@ void exit(int status) {
   switch_task();
 }
 
-u32 setup_stack(u32 stack_pointer, int argc, char **argv) {
-  mmu_allocate_region(STACK_LOCATION - STACK_SIZE, STACK_SIZE, MMU_FLAG_RW,
-                      NULL);
+u32 setup_stack(u32 stack_pointer, int argc, char **argv, int *err) {
+  *err = 0;
+  if (!mmu_allocate_region(STACK_LOCATION - STACK_SIZE, STACK_SIZE, MMU_FLAG_RW,
+                           NULL)) {
+    *err = 1;
+    return 0;
+  }
   flush_tlb();
 
   u32 ptr = stack_pointer;
@@ -316,7 +320,11 @@ int exec(const char *filename, char **argv, int dealloc_argv,
 
   current_task->data_segment_end = align_page((void *)end_of_code);
 
-  u32 ptr = setup_stack(0x90000000, argc, argv);
+  int err;
+  u32 ptr = setup_stack(0x90000000, argc, argv, &err);
+  if (err) {
+    return 0;
+  }
 
   if (dealloc_argv) {
     for (int i = 0; i < argc; i++) {
@@ -557,32 +565,27 @@ void *allocate_virtual_user_memory(size_t length, int prot, int flags) {
   (void)prot;
   (void)flags;
   void *rc = get_free_virtual_memory(length);
-  if ((void *)-1 == rc) {
-    return (void *)-1;
+  if (!rc) {
+    return NULL;
   }
 
-  mmu_allocate_region(rc, length, MMU_FLAG_RW, NULL);
-  return rc;
-}
-
-void *user_kernel_mapping(void *kernel_addr, size_t length) {
-  void *rc = get_free_virtual_memory(length);
-  if ((void *)-1 == rc) {
-    return (void *)-1;
+  if (!mmu_allocate_region(rc, length, MMU_FLAG_RW, NULL)) {
+    return NULL;
   }
-
-  mmu_map_directories(rc, NULL, kernel_addr, NULL, length);
   return rc;
 }
 
 void *create_physical_mapping(void **physical_addresses, size_t length) {
   void *rc = get_free_virtual_memory(length);
-  if ((void *)-1 == rc) {
-    return (void *)-1;
+  if (!rc) {
+    return NULL;
   }
   int n = (uintptr_t)align_page((void *)length) / 0x1000;
   for (int i = 0; i < n; i++) {
-    mmu_map_physical(rc + (i * 0x1000), NULL, physical_addresses[i], 0x1000);
+    if (!mmu_map_physical(rc + (i * 0x1000), NULL, physical_addresses[i],
+                          0x1000)) {
+      return NULL;
+    }
   }
   return rc;
 }
@@ -616,7 +619,7 @@ void *mmap(void *addr, size_t length, int prot, int flags, int fd,
 
   if (-1 == fd) {
     void *rc = allocate_virtual_user_memory(length, prot, flags);
-    if ((void *)-1 == rc) {
+    if (!rc) {
       kprintf("ENOMEM\n");
       return (void *)-ENOMEM;
     }
@@ -643,6 +646,10 @@ void *mmap(void *addr, size_t length, int prot, int flags, int fd,
     length = vmobject->size;
   }
   void *rc = create_physical_mapping(vmobject->object, length);
+  if (!rc) {
+      kprintf("ENOMEM\n");
+      return (void *)-ENOMEM;
+  }
   free_map->u_address = rc;
   free_map->k_address = NULL;
   free_map->length = length;