From 8b55eb4ab7d71889c60ea7e53e7f882ca85b4550 Mon Sep 17 00:00:00 2001
From: Leandro Carneiro <leandro@carnei.ro>
Date: Fri, 19 Mar 2021 16:47:48 -0300
Subject: [PATCH] feat(span-tag): add header x-request-id to span tag

---
 kong/plugins/zipkin/handler.lua |  7 +++++++
 kong/plugins/zipkin/schema.lua  |  3 +++
 spec/zipkin_spec.lua            | 11 ++++++++++-
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/kong/plugins/zipkin/handler.lua b/kong/plugins/zipkin/handler.lua
index 33c72f5..5190b37 100644
--- a/kong/plugins/zipkin/handler.lua
+++ b/kong/plugins/zipkin/handler.lua
@@ -184,6 +184,13 @@ if subsystem == "http" then
     get_or_add_proxy_span(zipkin, access_start)
 
     tracing_headers.set(conf.header_type, zipkin.header_type, zipkin.proxy_span, conf.default_header_type)
+
+    if conf.include_header_x_request_id then
+      local x_request_id = kong.request.get_header('x-request-id')
+      if x_request_id then
+        zipkin.request_span:set_tag(conf.tag_name_for_x_request_id, x_request_id)
+      end
+    end
   end
 
 
diff --git a/kong/plugins/zipkin/schema.lua b/kong/plugins/zipkin/schema.lua
index 5f91a9b..b227bf7 100644
--- a/kong/plugins/zipkin/schema.lua
+++ b/kong/plugins/zipkin/schema.lua
@@ -52,6 +52,9 @@ return {
                              between = { 0, 1 } } },
                                         { default_service_name = { type = "string", default = nil } },
           { include_credential = { type = "boolean", required = true, default = true } },
+          { include_header_x_request_id = { type = "boolean", required = true, default = false } },
+          { tag_name_for_x_request_id = { type = "string", required = true, default = "guid:x-request-id",
+                            not_one_of = PROTECTED_TAGS } },
           { traceid_byte_count = { type = "integer", required = true, default = 16, one_of = { 8, 16 } } },
           { header_type = { type = "string", required = true, default = "preserve",
                             one_of = { "preserve", "b3", "b3-single", "w3c", "jaeger", "ot" } } },
diff --git a/spec/zipkin_spec.lua b/spec/zipkin_spec.lua
index 6fc9e04..066c670 100644
--- a/spec/zipkin_spec.lua
+++ b/spec/zipkin_spec.lua
@@ -164,6 +164,7 @@ describe("http integration tests with zipkin server [#"
           { name = "static", value = "ok" },
         },
         default_header_type = "b3-single",
+        include_header_x_request_id = true,
       }
     })
 
@@ -221,12 +222,15 @@ describe("http integration tests with zipkin server [#"
 
   it("generates spans, tags and annotations for regular requests", function()
     local start_s = ngx.now()
+    local uuid = require("resty.jit-uuid")
+    local x_request_id = uuid.generate_v4()
 
     local r = proxy_client:get("/", {
       headers = {
         ["x-b3-sampled"] = "1",
         host  = "http-route",
-        ["zipkin-tags"] = "foo=bar; baz=qux"
+        ["zipkin-tags"] = "foo=bar; baz=qux",
+        ["x-request-id"] = x_request_id
       },
     })
     assert.response(r).has.status(200)
@@ -244,6 +248,7 @@ describe("http integration tests with zipkin server [#"
       ["http.method"] = "GET",
       ["http.path"] = "/",
       ["http.status_code"] = "200", -- found (matches server status)
+      ["guid:x-request-id"] = x_request_id,
       lc = "kong",
       static = "ok",
       foo = "bar",
@@ -292,6 +297,8 @@ describe("http integration tests with zipkin server [#"
 
   it("generates spans, tags and annotations for regular requests (#grpc)", function()
     local start_s = ngx.now()
+    local uuid = require("resty.jit-uuid")
+    local x_request_id = uuid.generate_v4()
 
     local ok, resp = proxy_client_grpc({
       service = "hello.HelloService.SayHello",
@@ -300,6 +307,7 @@ describe("http integration tests with zipkin server [#"
       },
       opts = {
         ["-H"] = "'x-b3-sampled: 1'",
+        ["-H"] = "'x-request-id: " .. x_request_id .. "'",
         ["-authority"] = "grpc-route",
       }
     })
@@ -320,6 +328,7 @@ describe("http integration tests with zipkin server [#"
       ["http.method"] = "POST",
       ["http.path"] = "/hello.HelloService/SayHello",
       ["http.status_code"] = "200", -- found (matches server status)
+      ["guid:x-request-id"] = x_request_id,
       lc = "kong",
       static = "ok",
     }, request_tags)