From abfe9999f447c15d904b3c11f32d4a22a45470a0 Mon Sep 17 00:00:00 2001
From: Robert Morris <rtm@csail.mit.edu>
Date: Mon, 1 Jul 2019 17:46:06 -0400
Subject: [PATCH] have fork() fail, not panic, if not enough phys mem

---
 Makefile      |  1 +
 kernel/defs.h |  2 +-
 kernel/proc.c | 42 ++++++++++++++++++++++++++-------------
 kernel/vm.c   | 11 +++++++++--
 user/cow.c    | 54 +++++++++++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 94 insertions(+), 16 deletions(-)
 create mode 100644 user/cow.c

diff --git a/Makefile b/Makefile
index 88130e1..03befd7 100644
--- a/Makefile
+++ b/Makefile
@@ -128,6 +128,7 @@ UPROGS=\
 	$U/_usertests\
 	$U/_wc\
 	$U/_zombie\
+	$U/_cow\
 
 fs.img: mkfs/mkfs README $(UPROGS)
 	mkfs/mkfs fs.img README $(UPROGS)
diff --git a/kernel/defs.h b/kernel/defs.h
index 1b397fe..1c0b39a 100644
--- a/kernel/defs.h
+++ b/kernel/defs.h
@@ -185,7 +185,7 @@ pagetable_t     uvmcreate(void);
 void            uvminit(pagetable_t, uchar *, uint);
 uint64          uvmalloc(pagetable_t, uint64, uint64);
 uint64          uvmdealloc(pagetable_t, uint64, uint64);
-void            uvmcopy(pagetable_t, pagetable_t, uint64);
+int            uvmcopy(pagetable_t, pagetable_t, uint64);
 void            uvmfree(pagetable_t, uint64);
 void            mappages(pagetable_t, uint64, uint64, uint64, int);
 void            unmappages(pagetable_t, uint64, uint64, int);
diff --git a/kernel/proc.c b/kernel/proc.c
index 4ae34c8..e0ee965 100644
--- a/kernel/proc.c
+++ b/kernel/proc.c
@@ -109,6 +109,28 @@ found:
   return p;
 }
 
+// free a proc structure and the data hanging from it,
+// including user pages.
+// the proc lock must be held.
+static void
+freeproc(struct proc *p)
+{
+  if(p->kstack)
+    kfree(p->kstack);
+  p->kstack = 0;
+  if(p->tf)
+    kfree((void*)p->tf);
+  p->tf = 0;
+  if(p->pagetable)
+    proc_freepagetable(p->pagetable, p->sz);
+  p->pagetable = 0;
+  p->pid = 0;
+  p->parent = 0;
+  p->name[0] = 0;
+  p->killed = 0;
+  p->state = UNUSED;
+}
+
 // Create a page table for a given process,
 // with no users pages, but with trampoline pages.
 // Called both when creating a process, and
@@ -145,7 +167,8 @@ proc_freepagetable(pagetable_t pagetable, uint64 sz)
 {
   unmappages(pagetable, TRAMPOLINE, PGSIZE, 0);
   unmappages(pagetable, TRAMPOLINE-PGSIZE, PGSIZE, 0);
-  uvmfree(pagetable, sz);
+  if(sz > 0)
+    uvmfree(pagetable, sz);
 }
 
 // a user program that calls exec("/init")
@@ -223,7 +246,10 @@ fork(void)
   }
 
   // Copy user memory from parent to child.
-  uvmcopy(p->pagetable, np->pagetable, p->sz);
+  if(uvmcopy(p->pagetable, np->pagetable, p->sz) < 0){
+    freeproc(np);
+    return -1;
+  }
   np->sz = p->sz;
 
   np->parent = p;
@@ -319,17 +345,7 @@ wait(void)
       if(np->state == ZOMBIE){
         // Found one.
         pid = np->pid;
-        kfree(np->kstack);
-        np->kstack = 0;
-        kfree((void*)np->tf);
-        np->tf = 0;
-        proc_freepagetable(np->pagetable, np->sz);
-        np->pagetable = 0;
-        np->pid = 0;
-        np->parent = 0;
-        np->name[0] = 0;
-        np->killed = 0;
-        np->state = UNUSED;
+        freeproc(np);
         release(&ptable.lock);
         return pid;
       }
diff --git a/kernel/vm.c b/kernel/vm.c
index 580669f..7d67464 100644
--- a/kernel/vm.c
+++ b/kernel/vm.c
@@ -273,7 +273,9 @@ uvmfree(pagetable_t pagetable, uint64 sz)
 // its memory into a child's page table.
 // Copies both the page table and the
 // physical memory.
-void
+// returns 0 on success, -1 on failure.
+// frees any allocated pages on failure.
+int
 uvmcopy(pagetable_t old, pagetable_t new, uint64 sz)
 {
   pte_t *pte;
@@ -289,10 +291,15 @@ uvmcopy(pagetable_t old, pagetable_t new, uint64 sz)
     pa = PTE2PA(*pte);
     flags = PTE_FLAGS(*pte);
     if((mem = kalloc()) == 0)
-      panic("uvmcopy: kalloc failed");
+      goto err;
     memmove(mem, (char*)pa, PGSIZE);
     mappages(new, i, PGSIZE, (uint64)mem, flags);
   }
+  return 0;
+
+ err:
+  unmappages(new, 0, i, 1);
+  return -1;
 }
 
 // Copy from kernel to user.
diff --git a/user/cow.c b/user/cow.c
new file mode 100644
index 0000000..f770cc9
--- /dev/null
+++ b/user/cow.c
@@ -0,0 +1,54 @@
+//
+// tests for copy-on-write fork() assignment.
+//
+
+#include "kernel/types.h"
+#include "kernel/memlayout.h"
+#include "user/user.h"
+
+// allocate more than half of physical memory,
+// then fork. this will fail in the default
+// kernel, which does not support copy-on-write.
+void
+simpletest()
+{
+  uint64 phys_size = PHYSTOP - KERNBASE;
+  int sz = (phys_size / 3) * 2;
+
+  printf(1, "simple: ");
+  
+  char *p = sbrk(sz);
+  if(p == (char*)0xffffffffffffffffL){
+    printf(1, "sbrk(%d) failed\n", sz);
+    exit();
+  }
+
+  for(char *q = p; q < p + sz; q += 4096){
+    *(int*)q = getpid();
+  }
+
+  int pid = fork();
+  if(pid < 0){
+    printf(1, "fork() failed\n");
+    exit();
+  }
+
+  if(pid == 0)
+    exit();
+
+  wait();
+
+  if(sbrk(-sz) == (char*)0xffffffffffffffffL){
+    printf(1, "sbrk(-%d) failed\n", sz);
+    exit();
+  }
+
+  printf(1, "simple ok\n");
+}
+
+int
+main(int argc, char *argv[])
+{
+  simpletest();
+  exit();
+}