From 997c8ce5682bed06885f873dc110cb6ad065e279 Mon Sep 17 00:00:00 2001
From: "Qirui(Keery) Nie" <windmgc@gmail.com>
Date: Tue, 18 Oct 2022 18:06:15 +0800
Subject: [PATCH] feat(zipkin): support adding trace id in response header
 (#9173)

This PR adds a new field in Zipkin plugin schema to support adding trace id in the response header.

Fix FTI-4198, #8124
---
 kong/plugins/zipkin/handler.lua           |  3 ++
 kong/plugins/zipkin/schema.lua            |  3 +-
 spec/03-plugins/34-zipkin/zipkin_spec.lua | 57 +++++++++++++++++++++++
 3 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/kong/plugins/zipkin/handler.lua b/kong/plugins/zipkin/handler.lua
index 9bd0c436ebd8..70ec63287edd 100644
--- a/kong/plugins/zipkin/handler.lua
+++ b/kong/plugins/zipkin/handler.lua
@@ -226,6 +226,9 @@ if subsystem == "http" then
 
     local proxy_span = get_or_add_proxy_span(zipkin, header_filter_start_mu)
     proxy_span:annotate("khs", header_filter_start_mu)
+    if conf.http_response_header_for_traceid then
+      kong.response.add_header(conf.http_response_header_for_traceid, proxy_span.trace_id)
+    end
   end
 
 
diff --git a/kong/plugins/zipkin/schema.lua b/kong/plugins/zipkin/schema.lua
index c3f3ba8e28c3..dba47198356c 100644
--- a/kong/plugins/zipkin/schema.lua
+++ b/kong/plugins/zipkin/schema.lua
@@ -51,7 +51,7 @@ return {
           { sample_ratio = { type = "number",
                              default = 0.001,
                              between = { 0, 1 } } },
-                                        { default_service_name = { type = "string", default = nil } },
+          { default_service_name = { type = "string", default = nil } },
           { include_credential = { type = "boolean", required = true, default = true } },
           { traceid_byte_count = { type = "integer", required = true, default = 16, one_of = { 8, 16 } } },
           { header_type = { type = "string", required = true, default = "preserve",
@@ -65,6 +65,7 @@ return {
           { connect_timeout = typedefs.timeout { default = 2000 } },
           { send_timeout = typedefs.timeout { default = 5000 } },
           { read_timeout = typedefs.timeout { default = 5000 } },
+          { http_response_header_for_traceid = { type = "string", default = nil }},
         },
     }, },
   },
diff --git a/spec/03-plugins/34-zipkin/zipkin_spec.lua b/spec/03-plugins/34-zipkin/zipkin_spec.lua
index b8449abd44d2..10eb9ab98ec8 100644
--- a/spec/03-plugins/34-zipkin/zipkin_spec.lua
+++ b/spec/03-plugins/34-zipkin/zipkin_spec.lua
@@ -423,6 +423,63 @@ for _, strategy in helpers.each_strategy() do
   end)
 end
 
+for _, strategy in helpers.each_strategy() do
+  describe("http_response_header_for_traceid configuration", function()
+    local proxy_client, service
+
+    setup(function()
+      local bp = helpers.get_db_utils(strategy, { "services", "routes", "plugins" })
+
+      service = bp.services:insert {
+        name = string.lower("http-" .. utils.random_string()),
+      }
+
+      -- kong (http) mock upstream
+      bp.routes:insert({
+        name = string.lower("route-" .. utils.random_string()),
+        service = service,
+        hosts = { "http-route" },
+        preserve_host = true,
+      })
+
+      -- enable zipkin plugin globally, with sample_ratio = 1
+      bp.plugins:insert({
+        name = "zipkin",
+        config = {
+          sample_ratio = 1,
+          http_endpoint = fmt("http://%s:%d/api/v2/spans", ZIPKIN_HOST, ZIPKIN_PORT),
+          default_header_type = "b3-single",
+          http_span_name = "method_path",
+          http_response_header_for_traceid = "X-B3-TraceId",
+        }
+      })
+
+      helpers.start_kong({
+        database = strategy,
+        nginx_conf = "spec/fixtures/custom_nginx.template",
+        stream_listen = helpers.get_proxy_ip(false) .. ":19000",
+      })
+
+      proxy_client = helpers.proxy_client()
+    end)
+
+    teardown(function()
+      helpers.stop_kong()
+    end)
+
+    it("custom traceid header included in response headers", function()
+      local r = proxy_client:get("/", {
+        headers = {
+          host  = "http-route",
+        },
+      })
+
+      assert.response(r).has.status(200)
+      assert.response(r).has.header("X-B3-TraceId")
+    end)
+  end)
+end
+
 for _, strategy in helpers.each_strategy() do
   describe("http_span_name configuration", function()
     local proxy_client, zipkin_client, service