Add header_register_callback builtin support (#1101)

* Add header_register_callback builtin support

Signed-off-by: Petr Shumilov <[email protected]>

builtin-functions/kphp-full/_functions.txt

@@ -149,6 +149,7 @@ function setrawcookie ($name ::: string, $value ::: string, $expire ::: int = 0,
 function register_shutdown_function (callable():void $function) ::: void;
 function ignore_user_abort ($enable ::: ?bool = null) ::: int;
 function flush() ::: void;
+function header_register_callback (callable():void $callback) ::: bool;
 /* // removed because it's not working now.
   function fastcgi_finish_request() ::: void;
 */

compiler/pipes/final-check.cpp

@@ -319,20 +319,28 @@ static void check_kphp_tracing_func_enter_branch_call(FunctionPtr current_functi
              "kphp_tracing_func_enter_branch() is available only inside functions with @kphp-tracing aggregate");
 }
 
-void check_register_shutdown_functions(VertexAdaptor<op_func_call> call) {
-  auto callback = call->args()[0].as<op_callback_of_builtin>();
+void raise_error_if_throwable(const std::string &where, const VertexAdaptor<op_callback_of_builtin> &callback) {
   if (!callback->func_id->can_throw()) {
     return;
   }
   std::vector<std::string> throws;
   for (const auto &e : callback->func_id->exceptions_thrown) {
     throws.emplace_back(e->name);
   }
-  kphp_error(false,
-             fmt_format("register_shutdown_callback should not throw exceptions\n"
-                        "But it may throw {}\n"
-                        "Throw chain: {}",
-                        vk::join(throws, ", "), callback->func_id->get_throws_call_chain()));
+  kphp_error(false, fmt_format("{} should not throw exceptions\n"
+                               "But it throws {}\n"
+                               "Throw chain: {}",
+                               where.c_str(), vk::join(throws, ", "), callback->func_id->get_throws_call_chain()));
+}
+
+void check_register_shutdown_functions(VertexAdaptor<op_func_call> call) {
+  auto callback = call->args()[0].as<op_callback_of_builtin>();
+  raise_error_if_throwable(call->func_id->name, callback);
+}
+
+void check_header_register_callback(VertexAdaptor<op_func_call> call) {
+  auto callback = call->args()[0].as<op_callback_of_builtin>();
+  raise_error_if_throwable(call->func_id->name, callback);
 }
 
 void mark_global_vars_for_memory_stats(const std::vector<VarPtr> &vars_list) {
@@ -871,6 +879,8 @@ void FinalCheckPass::check_op_func_call(VertexAdaptor<op_func_call> call) {
       kphp_error(arg_type->can_store_null(), fmt_format("is_null() will be always false for {}", arg_type->as_human_readable()));
     } else if (function_name == "register_shutdown_function") {
       check_register_shutdown_functions(call);
+    } else if (function_name == "header_register_callback") {
+      check_header_register_callback(call);
     } else if (function_name == "to_mixed") {
       check_to_mixed_call(call);
     } else if (vk::string_view{function_name}.starts_with("rpc_tl_query")) {

runtime/interface.cpp

@@ -206,6 +206,8 @@ static string http_status_line;
 static char headers_storage[sizeof(array<string>)];
 static array<string> *headers = reinterpret_cast <array<string> *> (headers_storage);
 static long long header_last_query_num = -1;
+static bool headers_sent = false;
+static headers_custom_handler_function_type headers_custom_handler_function;
 
 static bool check_status_line_int(const char *str, int str_len, int *pos) {
   if (*pos != str_len && str[*pos] == '0') {
@@ -562,7 +564,11 @@ static int ob_merge_buffers() {
 
 void f$flush() {
   php_assert(ob_cur_buffer >= 0 && php_worker.has_value());
-
+  // Run custom headers handler before body processing
+  if (headers_custom_handler_function && !headers_sent && query_type == QUERY_TYPE_HTTP) {
+    headers_sent = true;
+    headers_custom_handler_function();
+  }
   string_buffer const * http_body = compress_http_query_body(&oub[ob_system_level]);
   string_buffer const * http_headers = nullptr;
   if (!php_worker->flushed_http_connection) {
@@ -576,7 +582,12 @@ void f$flush() {
 }
 
 void f$fastcgi_finish_request(int64_t exit_code) {
-  int const ob_total_buffer = ob_merge_buffers();
+  // Run custom headers handler before body processing
+  if (headers_custom_handler_function && !headers_sent && query_type == QUERY_TYPE_HTTP) {
+    headers_sent = true;
+    headers_custom_handler_function();
+  }
+  int ob_total_buffer = ob_merge_buffers();
   if (php_worker.has_value() && php_worker->flushed_http_connection) {
     string const raw_response = oub[ob_total_buffer].str();
     http_set_result(nullptr, 0, raw_response.c_str(), raw_response.size(), static_cast<int32_t>(exit_code));
@@ -690,6 +701,14 @@ void register_shutdown_function_impl(shutdown_function_type &&f) {
   new(&shutdown_functions[shutdown_functions_count++]) shutdown_function_type{std::move(f)};
 }
 
+void register_header_handler_impl(headers_custom_handler_function_type &&f) {
+  dl::CriticalSectionGuard critical_section;
+  // Move assignment leads to lhs object invalidation and fires memory releasing mechanism
+  // But memory is already released by destructor after previous run
+  // Therefore we need to use placement new
+  new(&headers_custom_handler_function) headers_custom_handler_function_type{std::move(f)};
+}
+
 void finish(int64_t exit_code, bool from_exit) {
   check_script_timeout();
   if (!finished) {
@@ -2344,9 +2363,17 @@ static void free_shutdown_functions() {
   shutdown_functions_count = 0;
 }
 
+static void free_header_handler_function() {
+  headers_custom_handler_function.~headers_custom_handler_function_type();
+  new(&headers_custom_handler_function) headers_custom_handler_function_type{};
+  headers_sent = false;
+}
+
+
 static void free_interface_lib() {
   dl::enter_critical_section();//OK
   free_shutdown_functions();
+  free_header_handler_function();
   if (dl::query_num == uploaded_files_last_query_num) {
     const array<bool> *const_uploaded_files = uploaded_files;
     for (auto p = const_uploaded_files->begin(); p != const_uploaded_files->end(); ++p) {

runtime/interface.h

@@ -17,6 +17,7 @@
 
 extern string_buffer *coub;//TODO static
 using shutdown_function_type = std::function<void()>;
+using headers_custom_handler_function_type = std::function<void()>;
 
 enum class shutdown_functions_status {
   not_executed,
@@ -80,6 +81,14 @@ void f$register_shutdown_function(F &&f) {
   register_shutdown_function_impl(shutdown_function_type{std::forward<F>(f)});
 }
 
+void register_header_handler_impl(headers_custom_handler_function_type &&f);
+
+template <typename F>
+bool f$header_register_callback(F &&f) {
+  register_header_handler_impl(headers_custom_handler_function_type{std::forward<F>(f)});
+  return true;
+}
+
 void f$fastcgi_finish_request(int64_t exit_code = 0);
 
 __attribute__((noreturn))

tests/phpt/header_register_callback/throw_error_1.php

@@ -0,0 +1,10 @@
+@kphp_should_fail k2_skip
+/header_register_callback should not throw exceptions/
+<?php
+
+/** @kphp-required */
+function throwing_func() {
+  throw new Exception('BAD');
+}
+
+header_register_callback('throwing_func');

tests/phpt/header_register_callback/throw_error_2.php

@@ -0,0 +1,7 @@
+@kphp_should_fail k2_skip
+/header_register_callback should not throw exceptions/
+<?php
+
+header_register_callback(function () {
+  throw new Exception('BAD');
+});

tests/phpt/header_register_callback/throw_error_3.php

@@ -0,0 +1,13 @@
+@kphp_should_fail k2_skip
+/header_register_callback should not throw exceptions/
+<?php
+
+function may_throw(bool $cond) {
+  if ($cond) {
+    throw new Exception('BAD');
+  }
+}
+
+header_register_callback(function () {
+  may_throw(false);
+});

tests/phpt/pk/030_header_register_callback_unreachable_in_cli.phpt

@@ -0,0 +1,8 @@
+@ok
+<?php
+
+echo "Zero";
+header_register_callback(function () {
+    echo "In CLI is unreachable";
+});
+echo "One";

tests/phpt/pk/031_header_register_callback_no_invoked_by_flush.phpt

@@ -0,0 +1,9 @@
+@ok
+<?php
+
+echo "Zero";
+header_register_callback(function () {
+    echo "In CLI is unreachable, even via flush";
+});
+flush();
+echo "One";

tests/phpt/shutdown_functions/throw_error1.php

@@ -1,5 +1,5 @@
 @kphp_should_fail k2_skip
-/register_shutdown_callback should not throw exceptions/
+/register_shutdown_function should not throw exceptions/
 <?php
 
 /** @kphp-required */

tests/phpt/shutdown_functions/throw_error2.php

@@ -1,5 +1,5 @@
 @kphp_should_fail k2_skip
-/register_shutdown_callback should not throw exceptions/
+/register_shutdown_function should not throw exceptions/
 <?php
 
 register_shutdown_function(function () {

tests/phpt/shutdown_functions/throw_error3.php

@@ -1,5 +1,5 @@
 @kphp_should_fail k2_skip
-/register_shutdown_callback should not throw exceptions/
+/register_shutdown_function should not throw exceptions/
 <?php
 
 function may_throw(bool $cond) {

tests/python/tests/http_server/php/index.php

@@ -14,11 +14,11 @@ function assert($flag) {
 }
 
 interface I {
-    public function work();
+    public function work(string $output);
 }
 
 class ResumableWorker implements I {
-    public function work() {
+    public function work(string $output) {
         $job = function() {
                 sched_yield_sleep(2);
                 return true;
@@ -31,7 +31,7 @@ public function work() {
         foreach ($responses as $resp) {
             assert($resp);
         }
-        fwrite(STDERR, "test_ignore_user_abort/finish_resumable_work_" . $_GET["level"] . "\n");
+        fwrite(STDERR, $output);
     }
 }
 
@@ -42,13 +42,13 @@ public function __construct(int $port) {
         $this->port = $port;
     }
 
-    public function work() {
+    public function work(string $output) {
         $conn = new_rpc_connection('localhost', $this->port, 0, 5);
         $req_id = rpc_tl_query_one($conn, ["_" => "engine.sleep",
                                                     "time_ms" => 120]);
         $resp = rpc_tl_query_result_one($req_id);
         assert($resp['result']);
-        fwrite(STDERR, "test_ignore_user_abort/finish_rpc_work_" . $_GET["level"] . "\n");
+        fwrite(STDERR, $output);
    }
 }
 
@@ -119,6 +119,37 @@ public function work() {
         flush();
         sleep(2);
         throw new Exception('Exception');
+     case "flush_and_header_register_callback_flush_inside_callback":
+        echo "Zero ";
+        header_register_callback(function () {
+            echo "Two ";
+            flush();
+            sleep(2);
+            echo "Three ";
+        });
+        echo "One ";
+        return;
+     case 'flush_and_header_register_callback_invoked_after_flush':
+        header_register_callback(function () {
+            echo "One ";
+        });
+        echo "Zero ";
+        flush();
+        sleep(2);
+        echo "Two ";
+        return;
+     case 'flush_and_header_register_callback_no_double_invoked_after_flush':
+        header_register_callback(function () {
+            echo "One ";
+        });
+        echo "Zero ";
+        flush();
+        sleep(2);
+        echo "Two ";
+        flush();
+        sleep(2);
+        echo "Three ";
+        return;
     }
 
     echo "OK";
@@ -151,37 +182,40 @@ public function work() {
     register_shutdown_function('shutdown_function');
     /** @var I */
     $worker = null;
+    $msg = "";
     switch($_GET["type"]) {
      case "rpc":
         $worker = new RpcWorker(intval($_GET["port"]));
+        $msg = "test_ignore_user_abort/finish_rpc_work_" . $_GET["level"] . "\n";
         break;
      case "resumable":
         $worker = new ResumableWorker;
+        $msg = "test_ignore_user_abort/finish_resumable_work_" . $_GET["level"] . "\n";
         break;
      default:
         echo "ERROR"; return;
     }
     switch($_GET["level"]) {
      case "no_ignore":
-        $worker->work();
+        $worker->work($msg);
         break;
      case "ignore":
         ignore_user_abort(true);
-        $worker->work();
+        $worker->work($msg);
         fwrite(STDERR, "test_ignore_user_abort/finish_ignore_" . $_GET["type"] . "\n");
         ignore_user_abort(false);
         break;
      case "multi_ignore":
         ignore_user_abort(true);
-        $worker->work();
-        $worker->work();
+        $worker->work($msg);
+        $worker->work($msg);
         fwrite(STDERR, "test_ignore_user_abort/finish_multi_ignore_" . $_GET["type"] . "\n");
         ignore_user_abort(false);
         break;
      case "nested_ignore":
         ignore_user_abort(true);
         ignore_user_abort(true);
-        $worker->work();
+        $worker->work($msg);
         ignore_user_abort(false);
         fwrite(STDERR, "test_ignore_user_abort/finish_nested_ignore_" . $_GET["type"] . "\n");
         ignore_user_abort(false);
@@ -217,9 +251,25 @@ public function work() {
 } else if ($_SERVER["PHP_SELF"] === "/pid") {
     echo "pid=" . posix_getpid();
 } else if ($_SERVER["PHP_SELF"] === "/test_script_errors") {
-  critical_error("Test error");
+    critical_error("Test error");
 } else if ($_SERVER["PHP_SELF"] === "/test_oom_handler") {
-  require_once "test_oom_handler.php";
+    require_once "test_oom_handler.php";
+} else if ($_SERVER["PHP_SELF"] === "/test_header_register_callback") {
+    header_register_callback(function () {
+        global $_GET;
+        switch($_GET["act_in_callback"]) {
+         case "rpc":
+            $msg = "test_header_register_callback/rpc_in_callback\n";
+            (new RpcWorker(intval($_GET["port"])))->work($msg);
+            break;
+         case "exit":
+            $msg = "test_header_register_callback/exit_in_callback";
+            exit($msg);
+         default:
+            echo "ERROR";
+            return;
+        }
+    });
 } else {
     if ($_GET["hints"] === "yes") {
         send_http_103_early_hints(["Content-Type: text/plain or application/json", "Link: </script.js>; rel=preload; as=script"]);