Fix job_set_critical

zircon-object/src/task/job.rs

@@ -56,7 +56,6 @@ struct JobInner {
     policy: JobPolicy,
     children: Vec<Arc<Job>>,
     processes: Vec<Arc<Process>>,
-    critical_proc: Option<(KoID, bool)>,
     timer_policy: TimerSlack,
 }
 
@@ -96,6 +95,11 @@ impl Job {
         self.inner.lock().policy.merge(&self.parent_policy)
     }
 
+    /// Get the parent job.
+    pub fn parent(&self) -> Option<Arc<Self>> {
+        self.parent.clone()
+    }
+
     /// Sets one or more security and/or resource policies to an empty job.
     ///
     /// The job's effective policies is the combination of the parent's
@@ -131,44 +135,11 @@ impl Job {
         Ok(())
     }
 
-    /// Set a process as critical to the job.
-    ///
-    /// When process terminates, job will be terminated as if `task_kill()` was
-    /// called on it. The return code used will be `ZX_TASK_RETCODE_CRITICAL_PROCESS_KILL`.
-    ///
-    /// The job specified must be the parent of process, or an ancestor.
-    ///
-    /// If `retcode_nonzero` is true, then job will only be terminated if process
-    /// has a non-zero return code.
-    pub fn set_critical(&self, proc: &Arc<Process>, retcode_nonzero: bool) -> ZxResult {
-        let mut inner = self.inner.lock();
-        if let Some((pid, _)) = inner.critical_proc {
-            if proc.id() == pid {
-                return Err(ZxError::ALREADY_BOUND);
-            }
-        }
-        if !inner.processes.iter().any(|p| proc.id() == p.id()) {
-            return Err(ZxError::INVALID_ARGS);
-        }
-        inner.critical_proc = Some((proc.id(), retcode_nonzero));
-        Ok(())
-    }
-
     /// Add a process to the job.
     pub(super) fn add_process(&self, process: Arc<Process>) {
         self.inner.lock().processes.push(process);
     }
 
-    pub(super) fn process_exit(&self, id: KoID, retcode: i64) {
-        let mut inner = self.inner.lock();
-        inner.processes.retain(|proc| proc.id() != id);
-        if let Some((pid, retcode_nonzero)) = inner.critical_proc {
-            if pid == id && !(retcode_nonzero && retcode == 0) {
-                unimplemented!("kill the job")
-            }
-        }
-    }
-
     pub fn get_info(&self) -> JobInfo {
         JobInfo::default()
     }

zircon-object/src/task/process.rs

@@ -82,6 +82,7 @@ struct ProcessInner {
     // special info
     debug_addr: usize,
     dyn_break_on_load: usize,
+    critical_to_job: Option<(Arc<Job>, bool)>,
 }
 
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
@@ -175,7 +176,13 @@ impl Process {
         }
         inner.threads.clear();
         inner.handles.clear();
-        self.job.process_exit(self.base.id, retcode);
+
+        // If we are critical to a job, we need to take action.
+        if let Some((_job, retcode_nonzero)) = &inner.critical_to_job {
+            if !retcode_nonzero || retcode != 0 {
+                unimplemented!("kill the job")
+            }
+        }
     }
 
     /// Check whether `condition` is allowed in the parent job's policy.
@@ -191,6 +198,36 @@ impl Process {
         }
     }
 
+    /// Set a process as critical to the job.
+    ///
+    /// When process terminates, job will be terminated as if `task_kill()` was
+    /// called on it. The return code used will be `ZX_TASK_RETCODE_CRITICAL_PROCESS_KILL`.
+    ///
+    /// The job specified must be the parent of process, or an ancestor.
+    ///
+    /// If `retcode_nonzero` is true, then job will only be terminated if process
+    /// has a non-zero return code.
+    pub fn set_critical_at_job(&self, critical_to_job: &Arc<Job>, retcode_nonzero: bool) -> ZxResult {
+        let mut inner = self.inner.lock();
+        if inner.critical_to_job.is_some() {
+            return Err(ZxError::ALREADY_BOUND);
+        }
+
+        let mut job = self.job.clone();
+        loop {
+            if job.id() == critical_to_job.id() {
+                inner.critical_to_job = Some((job, retcode_nonzero));
+                return Ok(());
+            }
+            if let Some(p) = job.parent() {
+                job = p;
+            } else {
+                break;
+            }
+        }
+        Err(ZxError::INVALID_ARGS)
+    }
+
     /// Get process status.
     pub fn status(&self) -> Status {
         self.inner.lock().status

zircon-syscall/src/task.rs

@@ -132,7 +132,7 @@ impl Syscall<'_> {
         let proc = self.thread.proc();
         let job = proc.get_object_with_rights::<Job>(job_handle, Rights::DESTROY)?;
         let process = proc.get_object_with_rights::<Process>(process_handle, Rights::WAIT)?;
-        job.set_critical(&process, retcode_nonzero)?;
+        process.set_critical_at_job(&job, retcode_nonzero)?;
         Ok(())
     }