2 files changed, 15 insertions, 3 deletions
diff --git a/user/init.c b/user/init.c
index 5df6deb..13764ca 100644
--- a/user/init.c
+++ b/user/init.c
@@ -31,8 +31,20 @@ main(void)
       printf("init: exec sh failed\n");
       exit(1);
     }
-    while((wpid=wait(0)) >= 0 && wpid != pid){
-      //printf("zombie!\n");
+
+    for(;;){
+      // this call to wait() returns if the shell exits,
+      // or if a parentless process exits.
+      wpid = wait((int *) 0);
+      if(wpid == pid){
+        // the shell exited; restart it.
+        break;
+      } else if(wpid < 0){
+        printf("init: wait returned an error\n");
+        exit(1);
+      } else {
+        // it was a parentless process; do nothing.
+      }
     }
   }
 }
diff --git a/user/initcode.S b/user/initcode.S
index ca76972..e8f7a91 100644
--- a/user/initcode.S
+++ b/user/initcode.S
@@ -1,4 +1,4 @@
-# Initial process execs /init.
+# Initial process that execs /init.
 # This code runs in user space.
 
 #include "syscall.h"