Fix setting multiple cookies bug

2 files changed, 47 insertions, 6 deletions

diff --git a/lib/resty/cookie.lua b/lib/resty/cookie.lua
index be12e5f..a96da60 100644
--- a/lib/resty/cookie.lua
+++ b/lib/resty/cookie.lua
@@ -8,6 +8,7 @@ local sub           = string.sub
 local format        = string.format
 local log           = ngx.log
 local ERR           = ngx.ERR
+local ngx_header    = ngx.header
 
 local EQUAL         = byte("=")
 local SEMICOLON     = byte(";")
@@ -20,6 +21,11 @@ if not ok then
     new_tab = function (narr, nrec) return {} end
 end
 
+local ok, clear_tab = pcall(require, "table.clear")
+if not ok then
+    clear_tab = function(tab) for k, _ in pairs(tab) do tab[k] = nil end end
+end
+
 local _M = new_tab(0, 2)
 
 _M._VERSION = '0.01'
@@ -98,7 +104,8 @@ function _M.new(self)
     --if not _cookie then
         --return nil, "no cookie found in current request"
     --end
-    return setmetatable({ _cookie = _cookie }, mt)
+    return setmetatable({ _cookie = _cookie, set_cookie_table = new_tab(4, 0) },
+        mt)
 end
 
 function _M.get(self, key)
@@ -150,8 +157,32 @@ function _M.set(self, cookie)
     if not cookie_str then
         return nil, err
     end
-    print(cookie_str)
-    ngx.header['Set-Cookie'] = cookie_str
+
+    local set_cookie = ngx_header['Set-Cookie']
+    local set_cookie_type = type(set_cookie)
+    local t = self.set_cookie_table
+    clear_tab(t)
+
+    if set_cookie_type == "string" then
+        -- only one cookie has been setted
+        t[1] = set_cookie
+        t[2] = cookie_str
+        ngx_header['Set-Cookie'] = t
+    elseif set_cookie_type == "table" then
+        -- more than one cookies has been setted
+        local size = #set_cookie
+
+        -- we can not set cookie like ngx.header['Set-Cookie'][3] = val
+        -- so create a new table, copy all the values, and then set it back
+        for i=1, size do
+            t[i] = ngx_header['Set-Cookie'][i]
+        end
+        t[size + 1] = cookie_str
+        ngx_header['Set-Cookie'] = t
+    else
+        -- no cookie has been setted
+        ngx_header['Set-Cookie'] = cookie_str
+    end
     return true
 end
 
diff --git a/t/sanity.t b/t/sanity.t
index 0899a82..1385ebb 100644
--- a/t/sanity.t
+++ b/t/sanity.t
@@ -235,6 +235,15 @@ Set cookie
                 ngx.log(ngx.ERR, err)
                 return
             end
+
+            local ok, err = cookie:set({
+                key = "ID", value = "0xf7898",
+                expires = "Wed, 09 Jun 2021 10:18:14 GMT"
+            })
+            if not ok then
+                ngx.log(ngx.ERR, err)
+                return
+            end
             ngx.say("Set cookie")
         ';
     }
@@ -242,8 +251,9 @@ Set cookie
 GET /t
 --- no_error_log
 [error]
---- response_headers
-Set-Cookie: Name=Bob; Path=/
-Set-Cookie: Age=20
+--- comment
+because "--- response_headers" does not work with multiple headers with the same
+key, so use "--- raw_response_headers_like" instead
+--- raw_response_headers_like: Set-Cookie: Name=Bob; Path=/\r\nSet-Cookie: Age=20\r\nSet-Cookie: ID=0xf7898; Expires=Wed, 09 Jun 2021 10:18:14 GMT
 --- response_body
 Set cookie