lab cow: finishcow

8 files changed, 107 insertions, 9 deletions

diff --git a/Makefile b/Makefile
index ecb0f9e..217311b 100644
--- a/Makefile
+++ b/Makefile
@@ -13,6 +13,7 @@ OBJS = \
   $K/kalloc.o \
   $K/string.o \
   $K/main.o \
+  $K/cow.o \
   $K/vm.o \
   $K/proc.o \
   $K/swtch.o \
diff --git a/kernel/cow.c b/kernel/cow.c
new file mode 100644
index 0000000..b3634fc
--- /dev/null
+++ b/kernel/cow.c
@@ -0,0 +1,30 @@
+// COW pagefault handler
+#include "types.h"
+#include "riscv.h"
+#include "defs.h"
+
+int
+cow_handler(pagetable_t pagetable, uint64 va)
+{
+  // you can't really write to rediculous pointers
+  if(va >= MAXVA || PGROUNDDOWN(va) == 0)
+    return -1;
+  pte_t *pte = walk(pagetable, va, 0);
+  if(pte == 0 || (*pte & PTE_V) == 0 || (*pte & PTE_U) == 0)
+    return -1;
+  if(*pte & PTE_C){
+    uint64 pa_orig = PTE2PA(*pte);
+    uint64 pa_new = (uint64)kalloc();
+    if(pa_new == 0){
+      printf("cow pagefault: kalloc failed\n");
+      return -1;
+    }
+    // copy the page and add write permission
+    memmove((void*)pa_new, (void*)pa_orig, PGSIZE);
+    uint64 flags = (PTE_FLAGS(*pte) | PTE_W) & ~PTE_C;
+    *pte = PA2PTE(pa_new) | flags;
+    kfree((void*)pa_orig);
+  } else if ((*pte & PTE_W) == 0)
+    return -1;
+  return 0;
+}
diff --git a/kernel/defs.h b/kernel/defs.h
index 04a0276..870c984 100644
--- a/kernel/defs.h
+++ b/kernel/defs.h
@@ -23,6 +23,9 @@ void            consoleinit(void);
 void            consoleintr(int);
 void            consputc(int);
 
+// cow.c
+int             cow_handler(pagetable_t, uint64);
+
 // exec.c
 int             exec(char*, char**);
 
@@ -61,6 +64,8 @@ void            ramdiskintr(void);
 void            ramdiskrw(struct buf*);
 
 // kalloc.c
+int             refcnt_inc(uint64);
+int             refcnt_dec(uint64);
 void*           kalloc(void);
 void            kfree(void *);
 void            kinit(void);
diff --git a/kernel/kalloc.c b/kernel/kalloc.c
index c2fdb86..581f0f6 100644
--- a/kernel/kalloc.c
+++ b/kernel/kalloc.c
@@ -23,10 +23,41 @@ struct {
   struct run *freelist;
 } kmem;
 
+int phypg_refcnt[PHYSTOP/PGSIZE];
+
+// Increase the refcnt
+int
+refcnt_inc(uint64 pa)
+{
+  acquire(&kmem.lock);
+  int *prefcnt = &phypg_refcnt[pa/PGSIZE];
+  if(pa > PHYSTOP || *prefcnt < 1)
+    panic("increase refcnt");
+  (*prefcnt)++;
+  release(&kmem.lock);
+  return *prefcnt;
+}
+
+// Decrease the refcnt
+int
+refcnt_dec(uint64 pa)
+{
+  acquire(&kmem.lock);
+  int *prefcnt = &phypg_refcnt[pa/PGSIZE];
+  if(pa > PHYSTOP || *prefcnt < 1)
+    panic("decrease refcnt");
+  (*prefcnt)--;
+  release(&kmem.lock);
+  return *prefcnt;
+}
+
 void
 kinit()
 {
   initlock(&kmem.lock, "kmem");
+  // init all refcnt to 1, which would later be freed to 0 by kfree()
+  for(uint64 p = PGROUNDUP((uint64)end); p + PGSIZE <= PHYSTOP; p += PGSIZE)
+    phypg_refcnt[p/PGSIZE] = 1;
   freerange(end, (void*)PHYSTOP);
 }
 
@@ -51,6 +82,12 @@ kfree(void *pa)
   if(((uint64)pa % PGSIZE) != 0 || (char*)pa < end || (uint64)pa >= PHYSTOP)
     panic("kfree");
 
+  refcnt_dec((uint64)pa);
+
+  if(phypg_refcnt[(uint64)pa/PGSIZE] > 0)
+    // We still have refs to this phy page, do not actually free it
+    return;
+
   // Fill with junk to catch dangling refs.
   memset(pa, 1, PGSIZE);
 
@@ -72,8 +109,12 @@ kalloc(void)
 
   acquire(&kmem.lock);
   r = kmem.freelist;
-  if(r)
+  if(r){
+    if(phypg_refcnt[(uint64)r/PGSIZE])
+      panic("kalloc: invalid refcnt");
+    phypg_refcnt[(uint64)r/PGSIZE] = 1;
     kmem.freelist = r->next;
+  }
   release(&kmem.lock);
 
   if(r)
diff --git a/kernel/riscv.h b/kernel/riscv.h
index 5ede50a..7d6eb6e 100644
--- a/kernel/riscv.h
+++ b/kernel/riscv.h
@@ -353,6 +353,7 @@ typedef uint64 *pagetable_t; // 512 PTEs
 #define PTE_X (1L << 3)
 #define PTE_U (1L << 4) // user can access
 #define PTE_A (1L << 6) // riscv access bit                      
+#define PTE_C (1L << 8) // RSW low bit, use it to mark whether a page is COW
 
 // shift a physical address to the right place for a PTE.
 #define PA2PTE(pa) ((((uint64)pa) >> 12) << 10)
diff --git a/kernel/sysproc.c b/kernel/sysproc.c
index fac1d83..715a511 100644
--- a/kernel/sysproc.c
+++ b/kernel/sysproc.c
@@ -66,7 +66,7 @@ sys_sleep(void)
     sleep(&ticks, &tickslock);
   }
 
-  backtrace();
+  // backtrace();
 
   release(&tickslock);
   return 0;
diff --git a/kernel/trap.c b/kernel/trap.c
index eb8009f..bd3b100 100644
--- a/kernel/trap.c
+++ b/kernel/trap.c
@@ -73,9 +73,16 @@ usertrap(void)
     syscall();
   } else if((which_dev = devintr()) != 0){
     // ok
+  } else if(r_scause() == 13 || r_scause() == 15){
+    // deal with page fault
+    uint64 va = r_stval();
+    if(cow_handler(p->pagetable, va) < 0)
+      goto err;
   } else {
     printf("usertrap(): unexpected scause %p pid=%d\n", r_scause(), p->pid);
     printf("            sepc=%p stval=%p\n", r_sepc(), r_stval());
+  err:
+    printf("killing the process...\n");
     setkilled(p);
   }
 
diff --git a/kernel/vm.c b/kernel/vm.c
index 9c17fe7..11f9e0a 100644
--- a/kernel/vm.c
+++ b/kernel/vm.c
@@ -315,20 +315,26 @@ uvmcopy(pagetable_t old, pagetable_t new, uint64 sz)
   pte_t *pte;
   uint64 pa, i;
   uint flags;
-  char *mem;
+  // char *mem;
 
   for(i = 0; i < sz; i += PGSIZE){
     if((pte = walk(old, i, 0)) == 0)
       panic("uvmcopy: pte should exist");
     if((*pte & PTE_V) == 0)
       panic("uvmcopy: page not present");
-    pa = PTE2PA(*pte);
-    flags = PTE_FLAGS(*pte);
+    // do not do the actual copy, just increase the refcnt and mark pages readonly COW
+    /*
     if((mem = kalloc()) == 0)
       goto err;
     memmove(mem, (char*)pa, PGSIZE);
-    if(mappages(new, i, PGSIZE, (uint64)mem, flags) != 0){
-      kfree(mem);
+    */
+    *pte &= ~PTE_W;
+    *pte |= PTE_C;
+    pa = PTE2PA(*pte);
+    refcnt_inc(pa);
+    flags = PTE_FLAGS(*pte);
+    if(mappages(new, i, PGSIZE, (uint64)pa, flags) != 0){
+      // kfree(mem);
       goto err;
     }
   }
@@ -359,17 +365,24 @@ int
 copyout(pagetable_t pagetable, uint64 dstva, char *src, uint64 len)
 {
   uint64 n, va0, pa0;
-  pte_t *pte;
+  // pte_t *pte;
 
   while(len > 0){
     va0 = PGROUNDDOWN(dstva);
-    if(va0 >= MAXVA)
+
+    if(cow_handler(pagetable, va0) < 0)
       return -1;
+    
+    /*
     pte = walk(pagetable, va0, 0);
     if(pte == 0 || (*pte & PTE_V) == 0 || (*pte & PTE_U) == 0 ||
        (*pte & PTE_W) == 0)
       return -1;
     pa0 = PTE2PA(*pte);
+    */
+    pa0 = walkaddr(pagetable, va0);
+    if(pa0 == 0)
+      return -1;
     n = PGSIZE - (dstva - va0);
     if(n > len)
       n = len;