From 5edfe47f0a3fd0fd4db72522bb47ac9ab5564cdf Mon Sep 17 00:00:00 2001
From: "Qirui(Keery) Nie" <windmgc@gmail.com>
Date: Mon, 25 Sep 2023 20:50:17 +0800
Subject: [PATCH] fix(vault): vault try should avoid using semaphore in
 non-yieldable phases (#11617)

* tests(vault): add test to start kong with vault referenced postgres config

* feat(*): add check_phase_yieldable

* fix(*): with_coroutine_mutex bypass any non-yieldable phase

* fix(*): rename function to in_yieldable_phase
---
 kong/concurrency.lua                          |  6 +--
 kong/tools/utils.lua                          | 53 ++++++++++++++++---
 .../02-integration/13-vaults/03-mock_spec.lua | 22 ++++++++
 3 files changed, 71 insertions(+), 10 deletions(-)

diff --git a/kong/concurrency.lua b/kong/concurrency.lua
index 82bc5ad75fb4..58077d0aeed5 100644
--- a/kong/concurrency.lua
+++ b/kong/concurrency.lua
@@ -1,5 +1,6 @@
 local resty_lock = require "resty.lock"
 local ngx_semaphore = require "ngx.semaphore"
+local in_yieldable_phase = require("kong.tools.utils").in_yieldable_phase
 
 
 local type  = type
@@ -7,9 +8,6 @@ local error = error
 local pcall = pcall
 
 
-local get_phase = ngx.get_phase
-
-
 local concurrency = {}
 
 
@@ -91,7 +89,7 @@ function concurrency.with_coroutine_mutex(opts, fn)
     error("invalid value for opts.on_timeout", 2)
   end
 
-  if get_phase() == "init_worker" then
+  if not in_yieldable_phase() then
     return fn()
   end
 
diff --git a/kong/tools/utils.lua b/kong/tools/utils.lua
index 060aed457f93..d95a5adcbc85 100644
--- a/kong/tools/utils.lua
+++ b/kong/tools/utils.lua
@@ -1648,19 +1648,45 @@ function _M.sort_by_handler_priority(a, b)
   return prio_a > prio_b
 end
 
-do
+---
+-- Check if the phase is yieldable.
+-- @tparam string phase the phase to check, if not specified then
+-- the default value will be the current phase
+-- @treturn boolean true if the phase is yieldable, false otherwise
+local in_yieldable_phase do
   local get_phase = ngx.get_phase
+
+  -- https://github.com/openresty/lua-nginx-module/blob/c89469e920713d17d703a5f3736c9335edac22bf/src/ngx_http_lua_util.h#L35C10-L35C10
+  local LUA_CONTEXT_YIELDABLE_PHASE = {
+    rewrite = true,
+    server_rewrite = true,
+    access = true,
+    content = true,
+    timer = true,
+    ssl_client_hello = true,
+    ssl_certificate = true,
+    ssl_session_fetch = true,
+    preread = true,
+  }
+
+  in_yieldable_phase = function(phase)
+    if LUA_CONTEXT_YIELDABLE_PHASE[phase or get_phase()] == nil then
+      return false
+    end
+    return true
+  end
+end
+
+_M.in_yieldable_phase = in_yieldable_phase
+
+do
   local ngx_sleep = _G.native_ngx_sleep or ngx.sleep
 
   local YIELD_ITERATIONS = 1000
   local counter = YIELD_ITERATIONS
 
   function _M.yield(in_loop, phase)
-    if ngx.IS_CLI then
-      return
-    end
-    phase = phase or get_phase()
-    if phase == "init" or phase == "init_worker" then
+    if ngx.IS_CLI or not in_yieldable_phase(phase) then
       return
     end
     if in_loop then
@@ -1785,15 +1811,30 @@ _M.sha256_hex       = sha256_hex
 _M.sha256_base64    = sha256_base64
 _M.sha256_base64url = sha256_base64url
 
+
+local get_now_ms
 local get_updated_now_ms
+local get_start_time_ms
 do
   local now           = ngx.now
   local update_time   = ngx.update_time
+  local start_time    = ngx.req.start_time
+
+  function get_now_ms()
+    return now() * 1000 -- time is kept in seconds with millisecond resolution.
+  end
+
   function get_updated_now_ms()
     update_time()
     return now() * 1000 -- time is kept in seconds with millisecond resolution.
   end
+
+  function get_start_time_ms()
+    return start_time() * 1000 -- time is kept in seconds with millisecond resolution.
+  end
 end
+_M.get_now_ms         = get_now_ms
 _M.get_updated_now_ms = get_updated_now_ms
+_M.get_start_time_ms  = get_start_time_ms
 
 return _M
diff --git a/spec/02-integration/13-vaults/03-mock_spec.lua b/spec/02-integration/13-vaults/03-mock_spec.lua
index d8e2c896bfa6..77508422afea 100644
--- a/spec/02-integration/13-vaults/03-mock_spec.lua
+++ b/spec/02-integration/13-vaults/03-mock_spec.lua
@@ -145,4 +145,26 @@ for _, strategy in helpers.each_strategy() do
       end)
     end)
   end)
+
+  if strategy == "postgres" then
+    describe("ENV Vault #" .. strategy, function ()
+      describe("Kong Start", function ()
+        it("can resolve reference in init_phase", function ()
+          helpers.setenv("TEST_ENV_VAULT_LOGLEVEL", "debug")
+
+          assert(helpers.start_kong {
+            database = strategy,
+            prefix = helpers.test_conf.prefix,
+            nginx_conf = "spec/fixtures/custom_nginx.template",
+            vaults = "env",
+            log_level = "{vault://env/TEST_ENV_VAULT_LOGLEVEL}"
+          })
+
+          finally(function ()
+            assert(helpers.stop_kong())
+          end)
+        end)
+      end)
+    end)
+  end
 end