From b615511f5b64d65e397ed6c978ef4714a35cd205 Mon Sep 17 00:00:00 2001
From: Jun Ouyang <ouyangjun1999@gmail.com>
Date: Sat, 12 Oct 2024 16:24:33 +0800
Subject: [PATCH] feat(ai-proxy): fix code

---
 kong/init.lua                            |  2 ++
 kong/llm/proxy/handler.lua               |  2 +-
 kong/pdk/client/tls.lua                  |  1 +
 kong/pdk/private/phases.lua              |  1 +
 t/01-pdk/14-client-tls/00-phase_checks.t | 22 ++++++++++++++++++++++
 5 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/kong/init.lua b/kong/init.lua
index a758dc6f653f..44597dc9bdba 100644
--- a/kong/init.lua
+++ b/kong/init.lua
@@ -1033,6 +1033,8 @@ function Kong.ssl_certificate()
 end
 
 function Kong.ssl_client_hello()
+  local ctx = get_ctx_table(fetch_table(CTX_NS, CTX_NARR, CTX_NREC))
+  ctx.KONG_PHASE = PHASES.client_hello
 end
 
 function Kong.preread()
diff --git a/kong/llm/proxy/handler.lua b/kong/llm/proxy/handler.lua
index 39ce3ef333a9..6147208e0908 100644
--- a/kong/llm/proxy/handler.lua
+++ b/kong/llm/proxy/handler.lua
@@ -495,7 +495,7 @@ function _M:init_worker_for_plugin(plugin_name)
 
     local ctx = ngx.ctx
     -- ensure phases are set
-    ctx.KONG_PHASE = PHASES.certificate
+    ctx.KONG_PHASE = PHASES.client_hello
 
     kong_global.set_namespaced_log(kong, plugin_name)
     local snis_set, err = kong.cache:get(sni_cache_key, TTL_FOREVER,
diff --git a/kong/pdk/client/tls.lua b/kong/pdk/client/tls.lua
index a0ac55c6015a..a678976e1bf2 100644
--- a/kong/pdk/client/tls.lua
+++ b/kong/pdk/client/tls.lua
@@ -181,6 +181,7 @@ local function new()
   --   -- do something with err
   -- end
   function _TLS.disable_http2_alpn()
+    check_phase(PHASES.client_hello)
     return kong_tls.disable_http2_alpn()
   end
 
diff --git a/kong/pdk/private/phases.lua b/kong/pdk/private/phases.lua
index d3a2bca57179..12fcf5de10b3 100644
--- a/kong/pdk/private/phases.lua
+++ b/kong/pdk/private/phases.lua
@@ -10,6 +10,7 @@ local PHASES = {
   --init            = 0x00000001,
   init_worker       = 0x00000001,
   certificate       = 0x00000002,
+  client_hello      = 0x00000008,
   --set             = 0x00000004,
   rewrite           = 0x00000010,
   access            = 0x00000020,
diff --git a/t/01-pdk/14-client-tls/00-phase_checks.t b/t/01-pdk/14-client-tls/00-phase_checks.t
index d43f6519129b..8be8ca746431 100644
--- a/t/01-pdk/14-client-tls/00-phase_checks.t
+++ b/t/01-pdk/14-client-tls/00-phase_checks.t
@@ -26,6 +26,11 @@ qq{
             phase_check_functions(phases.certificate)
         }
 
+
+        ssl_client_hello_by_lua_block {
+            phase_check_functions(phases.client_hello)
+        }
+
         location / {
             set \$upstream_uri '/t';
             set \$upstream_scheme 'https';
@@ -66,6 +71,7 @@ qq{
                 args          = {},
                 init_worker   = "forced false",
                 certificate   = true,
+                client_hello  = "forced false",
                 rewrite       = "forced false",
                 access        = "forced false",
                 header_filter = "forced false",
@@ -78,6 +84,7 @@ qq{
                 args          = {},
                 init_worker   = false,
                 certificate   = true,
+                client_hello  = false,
                 rewrite       = false,
                 access        = false,
                 header_filter = false,
@@ -90,6 +97,7 @@ qq{
                 args          = {},
                 init_worker   = false,
                 certificate   = false,
+                client_hello  = false,
                 rewrite       = true,
                 access        = true,
                 response      = true,
@@ -101,6 +109,7 @@ qq{
                 method        = "set_client_verify",
                 args          = { "SUCCESS", },
                 init_worker   = "forced false",
+                client_hello  = "forced false",
                 certificate   = "forced false",
                 rewrite       = nil,
                 access        = nil,
@@ -109,6 +118,19 @@ qq{
                 body_filter   = "forced false",
                 log           = "forced false",
                 admin_api     = false,
+            }, {
+                method        = "disable_http2_alpn",
+                args          = {},
+                init_worker   = false,
+                client_hello  = true,
+                certificate   = false,
+                rewrite       = false,
+                access        = false,
+                header_filter = false,
+                response      = false,
+                body_filter   = false,
+                log           = false,
+                admin_api     = false,
             },
         }