From b3c5edc03de4b09be055dffe83d80617cdde5038 Mon Sep 17 00:00:00 2001
From: Matthieu Paret <matthieu.paret@flex.ai>
Date: Fri, 12 Jul 2024 14:52:01 +0200
Subject: [PATCH] add feature to jwt plugin:allow set headers from claim

---
 kong/plugins/jwt/handler.lua | 14 ++++++++++++++
 kong/plugins/jwt/schema.lua  |  5 +++++
 2 files changed, 19 insertions(+)

diff --git a/kong/plugins/jwt/handler.lua b/kong/plugins/jwt/handler.lua
index 45af8fd64c85..0bd3ab12a13e 100644
--- a/kong/plugins/jwt/handler.lua
+++ b/kong/plugins/jwt/handler.lua
@@ -12,6 +12,9 @@ local pairs = pairs
 local tostring = tostring
 local re_gmatch = ngx.re.gmatch
 
+local consts = {
+  JWT_CLAIM_HEADER_PREFIX = "X-Jwt-Claim"
+}
 
 local JwtHandler = {
   VERSION = kong_meta.version,
@@ -150,6 +153,15 @@ local function unauthorized(message, www_auth_content, errors)
   return { status = 401, message = message, headers = { ["WWW-Authenticate"] = www_auth_content }, errors = errors }
 end
 
+-- set header keys from claims
+local function set_headers(conf, claims)
+  local set_header = kong.service.request.set_header
+  for _, v in ipairs(conf.headers_to_set) do
+    if claims[v] then
+      set_header(consts.JWT_CLAIM_HEADER_PREFIX.."-"..v, claims[v])
+    end
+  end
+end
 
 local function do_authentication(conf)
   local token, err = retrieve_tokens(conf)
@@ -254,6 +266,8 @@ local function do_authentication(conf)
 
   set_consumer(consumer, jwt_secret, token)
 
+  set_headers(conf, claims)
+
   return true
 end
 
diff --git a/kong/plugins/jwt/schema.lua b/kong/plugins/jwt/schema.lua
index 0bfaef6e1354..4361a994853a 100644
--- a/kong/plugins/jwt/schema.lua
+++ b/kong/plugins/jwt/schema.lua
@@ -30,6 +30,11 @@ return {
                 type = "string",
                 one_of = { "exp", "nbf" },
           }, }, },
+          { headers_to_set = {
+            type = "set",
+            elements = {
+              type = "string"
+          }, }, },
           { anonymous = { description = "An optional string (consumer UUID or username) value to use as an “anonymous” consumer if authentication fails.", type = "string" }, },
           { run_on_preflight = { description = "A boolean value that indicates whether the plugin should run (and try to authenticate) on OPTIONS preflight requests. If set to false, then OPTIONS requests will always be allowed.", type = "boolean", required = true, default = true }, },
           { maximum_expiration = {