Fix bug in main script, add Ctrl+C handler when compiling ropchains

cli-tool/ropium

@@ -156,6 +156,7 @@ def load(args):
     seen_filename = False
     arch = None
     filenames = []
+    compiler_was_none = False
 
     # Parse arguments
     if not args:
@@ -187,6 +188,13 @@ def load(args):
     if arch and (arch not in str_to_arch):
         raise LoadException(f"Unsupported architecture: {arch}")
 
+    # Instanciate compiler if not already
+    if compiler is None:
+        compiler = ROPium(str_to_arch[arch])
+        compiler_was_none = True
+    elif compiler and arch and (str_to_arch[arch] != compiler.arch):
+        raise LoadException(f"Already working on a different architecture than '{arch}'")
+
     loaded_at_least_one = False
     print('') # So it's moar pretty
     for f in filenames:
@@ -198,11 +206,9 @@ def load(args):
             info(f"Loaded: {f}")
             loaded_at_least_one = True
 
-    # Instanciate compiler if not already
-    if compiler is None and loaded_at_least_one:
-        compiler = ROPium(str_to_arch[arch])
-    elif compiler and arch and (str_to_arch[arch] != compiler.arch):
-        raise LoadException(f"Already working on a different architecture than '{arch}'")
+
+    if compiler_was_none and not loaded_at_least_one:
+        compiler = None
 
 
 # Find command

libropium/compiler/compiler.cpp

@@ -26,7 +26,17 @@ ROPChain* CompilerTask::compile(Arch* arch, GadgetDB* db, Constraint* constraint
     ROPChain* res = nullptr;
     nb_tries = 3000;
 
+    // Set sigint handler to catch Ctrl+C
+    set_sigint_handler();
+
     while( nb_tries-- > 0 && !pending_strategies.empty() && !res){
+        // Check if user entered Ctrl+C
+        if( is_pending_sigint() ){
+            notify_sigint_handled();
+            unset_signint_handler();
+            return nullptr;
+        }
+
         graph = pending_strategies.back();
         pending_strategies.pop_back();
         if( graph->select_gadgets(*db, constraint, arch) ){
@@ -37,6 +47,10 @@ ROPChain* CompilerTask::compile(Arch* arch, GadgetDB* db, Constraint* constraint
         }
         delete graph; graph = nullptr;
     }
+
+    // Restore original sigint handler
+    unset_signint_handler();
+
     return res;
 }
 

libropium/compiler/strategy_graph.cpp

@@ -866,6 +866,11 @@ bool StrategyGraph::select_gadgets(GadgetDB& db, Constraint* constraint, Arch* a
         throw runtime_exception("StrategyGraph::select_gadget(): should NEVER be called with a non-NULL constraint and a NULL arch");
     }
 
+    // Check if SIGINT
+    if( is_pending_sigint()){
+        return false;
+    }
+
     // Otherwise do proper gadget selection : 
 
     // If root call

libropium/include/utils.hpp

@@ -147,4 +147,12 @@ string value_to_hex_str(int octets, addr_t addr);
 
 void disable_colors();
 void enable_colors();
+
+/* ========= Catching ctrl+C ============= */
+void set_sigint_handler();
+void unset_signint_handler();
+bool is_pending_sigint();
+void notify_sigint_handled();
+
+
 #endif

libropium/utils/utils.cpp

@@ -9,6 +9,7 @@
 #include <array>
 #include <vector>
 #include <exception>
+#include <signal.h>
 
 using std::ifstream;
 using std::ofstream;
@@ -180,3 +181,34 @@ void enable_colors(){
     g_EXPLOIT_DESCRIPTION_ANSI = DEFAULT_EXPLOIT_DESCRIPTION_ANSI;
     g_END_COLOR_ANSI = DEFAULT_END_COLOR_ANSI ;    
 }
+
+
+
+/* ========= Catching ctrl+C ============= */
+struct sigaction g_ropium_sigint_handler;
+struct sigaction g_ropium_prev_sigint_handler;
+bool g_ropium_sigint_flag = false;
+
+void ropium_sigint_handler(int s){
+    g_ropium_sigint_flag = true;
+}
+
+void set_sigint_handler(){
+    g_ropium_sigint_handler.sa_handler = ropium_sigint_handler;
+    sigemptyset(&g_ropium_sigint_handler.sa_mask);
+    g_ropium_sigint_handler.sa_flags = 0;
+
+    sigaction(SIGINT, &g_ropium_sigint_handler, &g_ropium_prev_sigint_handler);
+}
+
+void unset_signint_handler(){
+    sigaction(SIGINT, &g_ropium_prev_sigint_handler, nullptr);
+}
+
+bool is_pending_sigint(){
+    return g_ropium_sigint_flag;
+}
+
+void notify_sigint_handled(){
+    g_ropium_sigint_flag = false;
+}