7 files changed, 58 insertions, 27 deletions

diff --git a/src/sp_disabled_functions.c b/src/sp_disabled_functions.c
index 842b15c..379ed75 100644
--- a/src/sp_disabled_functions.c
+++ b/src/sp_disabled_functions.c
@@ -403,8 +403,9 @@ allow:
   return false;
 }
 
-bool should_drop_on_ret_ht(zval* return_value, const char* function_name,
-                           const sp_list_node* config, const HashTable* ht) {
+bool should_drop_on_ret_ht(const zval* return_value, const char* function_name,
+                           const sp_list_node* config, const HashTable* ht,
+                           zend_execute_data* execute_data) {
   const sp_list_node* ht_entry = NULL;
   bool ret = false;
 
@@ -414,17 +415,19 @@ bool should_drop_on_ret_ht(zval* return_value, const char* function_name,
 
   ht_entry = zend_hash_str_find_ptr(ht, function_name, strlen(function_name));
 
-  if (ht_entry && should_drop_on_ret(return_value, ht_entry, function_name)) {
+  if (ht_entry && should_drop_on_ret(return_value, ht_entry, function_name,
+        execute_data)) {
     ret = true;
   } else if (config && config->data) {
-    ret = should_drop_on_ret(return_value, config, function_name);
+    ret = should_drop_on_ret(return_value, config, function_name, execute_data);
   }
 
   return ret;
 }
 
-bool should_drop_on_ret(zval* return_value, const sp_list_node* config,
-                        const char* complete_function_path) {
+bool should_drop_on_ret(const zval* return_value, const sp_list_node* config,
+                        const char* complete_function_path,
+                        zend_execute_data* execute_data) {
   const char* current_filename = zend_get_executed_filename(TSRMLS_C);
   char current_file_hash[SHA256_SIZE * 2 + 1] = {0};
   bool match_type = false, match_value = false;
@@ -436,7 +439,12 @@ bool should_drop_on_ret(zval* return_value, const sp_list_node* config,
 
     assert(config_node->function || config_node->r_function);
 
-    if (config_node->function) {
+    if (config_node->functions_list) {
+      if (false == is_functions_list_matching(execute_data,
+                                              config_node->functions_list)) {
+        goto next;
+      }
+    } else if (config_node->function) {
       if (0 !=
           strcmp(ZSTR_VAL(config_node->function), complete_function_path)) {
         goto next;
@@ -513,7 +521,8 @@ ZEND_FUNCTION(check_disabled_function) {
           return_value, current_function_name,
           SNUFFLEUPAGUS_G(config)
               .config_disabled_functions_reg_ret->disabled_functions,
-          SNUFFLEUPAGUS_G(config).config_disabled_functions_ret_hooked)) {
+          SNUFFLEUPAGUS_G(config).config_disabled_functions_ret_hooked,
+          execute_data)) {
     sp_terminate();
   }
 }
diff --git a/src/sp_disabled_functions.h b/src/sp_disabled_functions.h
index 258e2e6..ae1500a 100644
--- a/src/sp_disabled_functions.h
+++ b/src/sp_disabled_functions.h
@@ -9,9 +9,10 @@ bool should_disable(zend_execute_data *, const char *, const zend_string *,
                     const char *, const sp_list_node *, const zend_string *);
 bool should_disable_ht(zend_execute_data *, const char *, const zend_string *,
                        const char *, const sp_list_node *, const HashTable *);
-bool should_drop_on_ret_ht(zval *, const char *, const sp_list_node *config,
-                           const HashTable *);
-bool should_drop_on_ret(zval *, const sp_list_node *config, const char *);
+bool should_drop_on_ret_ht(const zval *, const char *, const sp_list_node *config,
+                           const HashTable *, zend_execute_data *);
+bool should_drop_on_ret(const zval *, const sp_list_node *config, const char *,
+                        zend_execute_data *);
 char *get_complete_function_path(zend_execute_data const *const);
 
 #endif /* __SP_DISABLE_FUNCTIONS_H */
diff --git a/src/sp_execute.c b/src/sp_execute.c
index 844647e..5447ea1 100644
--- a/src/sp_execute.c
+++ b/src/sp_execute.c
@@ -151,6 +151,7 @@ static void sp_execute_ex(zend_execute_data *execute_data) {
 
   if (SNUFFLEUPAGUS_G(config).hook_execute) {
     char *function_name = get_complete_function_path(execute_data);
+    zval ret_val;
 
     if (!function_name) {
       orig_execute_ex(execute_data);
@@ -185,23 +186,34 @@ static void sp_execute_ex(zend_execute_data *execute_data) {
       }
     }
 
+    // When a function's return value isn't used, php doesn't store it in the execute_data,
+    // so we need to use a local variable to be able to match on it later.
+    if (EX(return_value) == NULL) {
+      memset(&ret_val, 0, sizeof(ret_val));
+      EX(return_value) = &ret_val;
+    }
+
     orig_execute_ex(execute_data);
 
-    if (EX(return_value) != NULL) {
-      if (UNEXPECTED(
-            true ==
-            should_drop_on_ret_ht(
-              EX(return_value), function_name,
-              SNUFFLEUPAGUS_G(config)
-              .config_disabled_functions_reg_ret->disabled_functions,
-              SNUFFLEUPAGUS_G(config).config_disabled_functions_ret))) {
-        sp_terminate();
-      }
+    if (UNEXPECTED(
+          true ==
+          should_drop_on_ret_ht(
+            EX(return_value), function_name,
+            SNUFFLEUPAGUS_G(config)
+            .config_disabled_functions_reg_ret->disabled_functions,
+            SNUFFLEUPAGUS_G(config).config_disabled_functions_ret,
+            execute_data))) {
+      sp_terminate();
     }
     efree(function_name);
+
+    if (EX(return_value) == &ret_val) {
+      EX(return_value) = NULL;
+    }
   } else {
     orig_execute_ex(execute_data);
   }
+
 }
 
 static void sp_zend_execute_internal(INTERNAL_FUNCTION_PARAMETERS) {
diff --git a/src/sp_utils.c b/src/sp_utils.c
index 41d817e..e872abd 100644
--- a/src/sp_utils.c
+++ b/src/sp_utils.c
@@ -151,7 +151,7 @@ static char* zend_string_to_char(const zend_string* zs) {
   return copy;
 }
 
-const zend_string* sp_zval_to_zend_string(zval* zv) {
+const zend_string* sp_zval_to_zend_string(const zval* zv) {
   switch (Z_TYPE_P(zv)) {
     case IS_LONG: {
       char* msg;
diff --git a/src/sp_utils.h b/src/sp_utils.h
index 97c97fd..c094fac 100644
--- a/src/sp_utils.h
+++ b/src/sp_utils.h
@@ -46,7 +46,7 @@
 
 void sp_log_msg(char const *feature, int type, const char *fmt, ...);
 int compute_hash(const char *const filename, char *file_hash);
-const zend_string *sp_zval_to_zend_string(zval *);
+const zend_string *sp_zval_to_zend_string(const zval *);
 bool sp_match_value(const zend_string *, const zend_string *, const sp_pcre *);
 bool sp_match_array_key(const zval *, const zend_string *, const sp_pcre *);
 bool sp_match_array_value(const zval *, const zend_string *, const sp_pcre *);
diff --git a/src/tests/disabled_functions_chain_call_user_func_ret.phpt b/src/tests/disabled_functions_chain_call_user_func_ret.phpt
index 44c9b24..d0ee389 100644
--- a/src/tests/disabled_functions_chain_call_user_func_ret.phpt
+++ b/src/tests/disabled_functions_chain_call_user_func_ret.phpt
@@ -21,6 +21,15 @@ echo one('matching') . "\n";
 echo one('still not matching') . "\n";
 
 ?>
---XFAIL--
 --EXPECTF--
-Match on ret is broken :/
+one
+two
+not matching_one
+one
+two
+
+Warning: [snuffleupagus][disabled_function] Aborted execution on return of the function 'two', because the function returned 'matching_two', which matched a rule in %a/tests/disabled_functions_chain_call_user_func_ret.php on line %d
+matching_one
+one
+two
+still not matching_one
diff --git a/src/tests/disabled_functions_ret_user.phpt b/src/tests/disabled_functions_ret_user.phpt
index 597a6b8..89c1bce 100644
--- a/src/tests/disabled_functions_ret_user.phpt
+++ b/src/tests/disabled_functions_ret_user.phpt
@@ -12,5 +12,5 @@ function qwe() {
 qwe();
 echo 1;
 ?>
---EXPECT--
-1
+--EXPECTF--
+Fatal error: [snuffleupagus][disabled_function] Aborted execution on return of the function 'qwe', because the function returned 'asd', which matched a rule in %a/tests/disabled_functions_ret_user.php on line %d