add headers_sent

VKCOM · Dec 11, 2024 · 1f3fdab · 1f3fdab
1 parent 2398687
commit 1f3fdab
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 4 deletions.
diff --git a/builtin-functions/kphp-full/_functions.txt b/builtin-functions/kphp-full/_functions.txt
@@ -141,6 +141,7 @@ function ob_get_flush () ::: string | false;
 function ob_get_length () ::: int | false;
 function ob_get_level () ::: int;
 
+function headers_sent () ::: bool;
 function header ($str ::: string, $replace ::: bool = true, $http_response_code ::: int = 0) ::: void;
 function headers_list () ::: string[];
 function send_http_103_early_hints($headers ::: string[]) ::: void;

diff --git a/runtime/interface.cpp b/runtime/interface.cpp
@@ -344,6 +344,10 @@ static void header(const char *str, int str_len, bool replace = true, int http_r
   }
 }
 
+bool f$headers_sent() {
+  return headers_sent;
+}
+
 void f$header(const string &str, bool replace, int64_t http_response_code) {
   header(str.c_str(), (int)str.size(), replace, static_cast<int32_t>(http_response_code));
 }
@@ -568,9 +572,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) {
+  if (!headers_sent && query_type == QUERY_TYPE_HTTP) {
     headers_sent = true;
-    headers_custom_handler_function();
+    if (headers_custom_handler_function) {
+      headers_custom_handler_function();
+    }
   }
   string_buffer const * http_body = compress_http_query_body(&oub[ob_system_level]);
   string_buffer const * http_headers = nullptr;
@@ -586,9 +592,11 @@ void f$flush() {
 
 void f$fastcgi_finish_request(int64_t exit_code) {
   // Run custom headers handler before body processing
-  if (headers_custom_handler_function && !headers_sent && query_type == QUERY_TYPE_HTTP) {
+  if (!headers_sent && query_type == QUERY_TYPE_HTTP) {
     headers_sent = true;
-    headers_custom_handler_function();
+    if (headers_custom_handler_function) {
+      headers_custom_handler_function();
+    }
   }
   int ob_total_buffer = ob_merge_buffers();
   if (php_worker.has_value() && php_worker->flushed_http_connection) {

diff --git a/runtime/interface.h b/runtime/interface.h
@@ -47,6 +47,8 @@ int64_t f$ob_get_level();
 
 void f$flush();
 
+bool f$headers_sent();
+
 void f$header(const string &str, bool replace = true, int64_t http_response_code = 0);
 
 array<string> f$headers_list();

diff --git a/tests/python/tests/http_server/php/index.php b/tests/python/tests/http_server/php/index.php
@@ -178,6 +178,10 @@ public function work(string $output) {
 
         break;
     }
+} else if ($_SERVER["PHP_SELF"] === "/test_headers_sent") {
+    echo (int)headers_sent();
+    flush();
+    echo (int)headers_sent();
 } else if ($_SERVER["PHP_SELF"] === "/test_ignore_user_abort") {
     register_shutdown_function('shutdown_function');
     /** @var I */

diff --git a/tests/python/tests/http_server/test_headers_sent.py b/tests/python/tests/http_server/test_headers_sent.py
@@ -0,0 +1,13 @@
+import socket
+from python.lib.testcase import KphpServerAutoTestCase
+from python.lib.http_client import RawResponse
+
+class TestHeadersSent(KphpServerAutoTestCase):
+    def test(self):
+        request = b"GET /test_headers_sent HTTP/1.1\r\nHost:localhost\r\n\r\n"
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.connect(('127.0.0.1', self.kphp_server.http_port))
+            s.send(request)
+
+            self.assertEqual(b"0", RawResponse(s.recv(4096)).content)
+            self.assertEqual(b"1", s.recv(20))