From 4e12036b71b006122753ec7e16f3ec17ca27a206 Mon Sep 17 00:00:00 2001
From: "Juan M. Uys" <opyate@gmail.com>
Date: Wed, 23 Mar 2016 07:12:51 +0000
Subject: [PATCH] make 'echo downstream' configurable and false default - also
 revert busted:match, since it's not for regex

---
 kong/plugins/correlation-id/handler.lua     | 10 ++++---
 kong/plugins/correlation-id/schema.lua      |  4 +++
 spec/plugins/correlation-id/access_spec.lua | 29 ++++++++++++++-------
 3 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/kong/plugins/correlation-id/handler.lua b/kong/plugins/correlation-id/handler.lua
index 4ff072bc7907..e4ac667ae893 100644
--- a/kong/plugins/correlation-id/handler.lua
+++ b/kong/plugins/correlation-id/handler.lua
@@ -44,13 +44,17 @@ function CorrelationIdHandler:access(conf)
     req_set_header(conf.header_name, header_value)
   end
 
-  -- For later use, to echo it back downstream
-  ngx.ctx.correlationid_header_value = header_value
+  if conf.echo_downstream then
+    -- For later use, to echo it back downstream
+    ngx.ctx.correlationid_header_value = header_value
+  end
 end
 
 function CorrelationIdHandler:header_filter(conf)
   CorrelationIdHandler.super.header_filter(self)
-  ngx.header[conf.header_name] = ngx.ctx.correlationid_header_value
+  if conf.echo_downstream then
+    ngx.header[conf.header_name] = ngx.ctx.correlationid_header_value
+  end
 end
 
 return CorrelationIdHandler
diff --git a/kong/plugins/correlation-id/schema.lua b/kong/plugins/correlation-id/schema.lua
index 838cd78cbb35..6f9782f6a753 100644
--- a/kong/plugins/correlation-id/schema.lua
+++ b/kong/plugins/correlation-id/schema.lua
@@ -8,6 +8,10 @@ return {
       type = "string",
       default = "uuid#counter",
       enum = {"uuid", "uuid#counter"}
+    },
+    echo_downstream = {
+      type = "boolean",
+      default = false
     }
   }
 }
diff --git a/spec/plugins/correlation-id/access_spec.lua b/spec/plugins/correlation-id/access_spec.lua
index 943012cf4a74..99e411ce4a08 100644
--- a/spec/plugins/correlation-id/access_spec.lua
+++ b/spec/plugins/correlation-id/access_spec.lua
@@ -15,12 +15,14 @@ describe("Correlation ID Plugin", function()
       api = {
         {request_host = "correlation1.com", upstream_url = "http://mockbin.com"},
         {request_host = "correlation2.com", upstream_url = "http://mockbin.com"},
-        {request_host = "correlation3.com", upstream_url = "http://mockbin.com"}
+        {request_host = "correlation3.com", upstream_url = "http://mockbin.com"},
+        {request_host = "correlation4.com", upstream_url = "http://mockbin.com"}
       },
       plugin = {
-        {name = "correlation-id", config = {}, __api = 1},
-        {name = "correlation-id", config = {header_name = "Foo-Bar-Id"}, __api = 2},
-        {name = "correlation-id", config = {generator = "uuid"}, __api = 3}
+        {name = "correlation-id", config = {echo_downstream = true}, __api = 1},
+        {name = "correlation-id", config = {header_name = "Foo-Bar-Id", echo_downstream = true}, __api = 2},
+        {name = "correlation-id", config = {generator = "uuid", echo_downstream = true}, __api = 3},
+        {name = "correlation-id", config = {}, __api = 4},
       }
     }
     spec_helper.start_kong()
@@ -31,15 +33,15 @@ describe("Correlation ID Plugin", function()
   end)
 
   local function test_with(host, header, pattern)
-    local response1, status1 = http_client.get(STUB_GET_URL, nil, {host = host})
+    local _, status1, headers1 = http_client.get(STUB_GET_URL, nil, {host = host})
     assert.equal(200, status1)
-    local correlation_id1 = json.decode(response1).headers[header:lower()]
-    assert.match(correlation_id1, pattern)
+    local correlation_id1 = headers1[header:lower()]
+    assert.truthy(correlation_id1:match(pattern))
 
-    local response2, status2 = http_client.get(STUB_GET_URL, nil, {host = host})
+    local _, status2, headers2 = http_client.get(STUB_GET_URL, nil, {host = host})
     assert.equal(200, status2)
-    local correlation_id2 = json.decode(response2).headers[header:lower()]
-    assert.match(correlation_id2, pattern)
+    local correlation_id2 = headers2[header:lower()]
+    assert.truthy(correlation_id2:match(pattern))
 
     assert.are_not_equals(correlation_id1, correlation_id2)
 
@@ -79,4 +81,11 @@ describe("Correlation ID Plugin", function()
     local correlation_id = json.decode(response).headers[DEFAULT_HEADER_NAME:lower()]
     assert.equals(existing_correlation_id, correlation_id)
   end)
+
+  it("should not echo back the correlation header", function()
+    local _, status, headers = http_client.get(STUB_GET_URL, nil, {host = "correlation4.com"})
+    assert.equal(200, status)
+    local correlation_id = headers[DEFAULT_HEADER_NAME:lower()]
+    assert.falsy(correlation_id)
+  end)
 end)