Matrix: prevent new programs from being launched after a task is killed

1 files changed, 48 insertions, 12 deletions

diff --git a/src/Program.cpp b/src/Program.cpp
index 736532e..dc92f16 100644
--- a/src/Program.cpp
+++ b/src/Program.cpp
@@ -14,37 +14,73 @@
 #define READ_END 0
 #define WRITE_END 1
 
+struct ThreadProgram {
+    ReadProgram read_program;
+    bool killed;
+};
+
+static std::unordered_map<std::thread::id, ThreadProgram> thread_current_program;
+static std::mutex thread_current_program_mutex;
+
 class CurrentThreadProgram {
 public:
+    CurrentThreadProgram() {
+        std::lock_guard<std::mutex> lock(thread_current_program_mutex);
+        ThreadProgram thread_program;
+        thread_program.read_program.pid = -1;
+        thread_program.read_program.read_fd = -1;
+        thread_program.killed = false;
+        thread_current_program[std::this_thread::get_id()] = std::move(thread_program);
+    }
+
+    ~CurrentThreadProgram() {
+        thread_current_program.erase(std::this_thread::get_id());
+    }
+
     void set(ReadProgram read_program) {
-        std::lock_guard<std::mutex> lock(mutex);
-        thread_current_program[std::this_thread::get_id()] = read_program;
+        std::lock_guard<std::mutex> lock(thread_current_program_mutex);
+        auto it = thread_current_program.find(std::this_thread::get_id());
+        if(it != thread_current_program.end())
+            it->second.read_program = std::move(read_program);
     }
 
     void clear() {
-        std::lock_guard<std::mutex> lock(mutex);
-        thread_current_program.erase(std::this_thread::get_id());
+        std::lock_guard<std::mutex> lock(thread_current_program_mutex);
+        auto it = thread_current_program.find(std::this_thread::get_id());
+        if(it != thread_current_program.end()) {
+            it->second.read_program.pid = -1;
+            it->second.read_program.read_fd = -1;
+        }
     }
 
     void kill_in_thread(const std::thread::id &thread_id) {
-        std::lock_guard<std::mutex> lock(mutex);
+        std::lock_guard<std::mutex> lock(thread_current_program_mutex);
         auto it = thread_current_program.find(thread_id);
-        if(it != thread_current_program.end()) {
-            close(it->second.read_fd);
-            kill(it->second.pid, SIGTERM);
+        if(it != thread_current_program.end() && it->second.read_program.pid != -1 && it->second.read_program.read_fd != -1) {
+            close(it->second.read_program.read_fd);
+            kill(it->second.read_program.pid, SIGTERM);
+            it->second.killed = true;
         }
     }
-private:
-    std::unordered_map<std::thread::id, ReadProgram> thread_current_program;
-    std::mutex mutex;
+
+    bool is_killed() {
+        std::lock_guard<std::mutex> lock(thread_current_program_mutex);
+        auto it = thread_current_program.find(std::this_thread::get_id());
+        if(it != thread_current_program.end())
+            return it->second.killed;
+        return false;
+    }
 };
 
-static CurrentThreadProgram current_thread_program;
+thread_local CurrentThreadProgram current_thread_program;
 
 int exec_program_pipe(const char **args, ReadProgram *read_program) {
     /* 1 arguments */
     if(args[0] == NULL)
         return -1;
+
+    if(current_thread_program.is_killed())
+        return -1;
     
     int fd[2];
     if(pipe(fd) == -1) {