Support Multiple Set-Cookie Headers

CHANGELOG.md

@@ -3,6 +3,12 @@
 
 See this http://keepachangelog.com link for information on how we want this documented formatted.
 
+## v2.6.0
+
+#### Fixed
+
+* Support multiple Set-Cookie headers for all rest types.
+
 ## v2.5.3
 
 #### Fixed

Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    lamby (2.5.3)
+    lamby (2.6.0)
       rack
 
 GEM

lib/lamby/handler.rb

@@ -31,6 +31,13 @@ def headers
       @headers
     end
 
+    def set_cookies
+      return @set_cookies if defined?(@set_cookies)
+      @set_cookies = if @headers && @headers['Set-Cookie']
+        @headers.delete('Set-Cookie').split("\n")
+      end
+    end
+
     def body
       @rbody ||= ''.tap do |rbody|
         @body.each { |part| rbody << part }
@@ -40,6 +47,7 @@ def body
     def call
       return self if @called
       @status, @headers, @body = call_app
+      set_cookies
       @called = true
       self
     end

lib/lamby/rack_alb.rb

@@ -11,7 +11,10 @@ def multi_value?
 
     def response(handler)
       hhdrs = handler.headers
-      multivalue_headers = hhdrs.transform_values { |v| Array[v].compact.flatten } if multi_value?
+      if multi_value?
+        multivalue_headers = hhdrs.transform_values { |v| Array[v].compact.flatten }
+        multivalue_headers['Set-Cookie'] = handler.set_cookies if handler.set_cookies
+      end
       status_description = "#{handler.status} #{::Rack::Utils::HTTP_STATUS_CODES[handler.status]}"
       base64_encode = hhdrs['Content-Transfer-Encoding'] == 'binary' || hhdrs['X-Lamby-Base64'] == '1'
       body = Base64.strict_encode64(handler.body) if base64_encode

lib/lamby/rack_http.rb

@@ -6,6 +6,15 @@ def response(handler)
         { isBase64Encoded: true, body: handler.body64 }
       else
         super
+      end.tap do |r|
+        if cookies = handler.set_cookies
+          if payload_version_one?
+            r[:multiValueHeaders] ||= {}
+            r[:multiValueHeaders]['Set-Cookie'] = cookies
+          else
+            r[:cookies] = cookies
+          end
+        end
       end
     end
 
@@ -79,6 +88,9 @@ def server_protocol
         'HTTP/1.1'
     end
 
+    def payload_version_one?
+      event['version'] == '1.0'
+    end
 
   end
 end

lib/lamby/rack_rest.rb

@@ -6,6 +6,11 @@ def response(handler)
         { isBase64Encoded: true, body: handler.body64 }
       else
         super
+      end.tap do |r|
+        if cookies = handler.set_cookies
+          r[:multiValueHeaders] ||= {}
+          r[:multiValueHeaders]['Set-Cookie'] = cookies
+        end
       end
     end
 

lib/lamby/version.rb

@@ -1,3 +1,3 @@
 module Lamby
-  VERSION = '2.5.3'
+  VERSION = '2.6.0'
 end

test/dummy_app/app/controllers/application_controller.rb

@@ -26,6 +26,12 @@ def exception
     raise 'hell'
   end
 
+  def cooks
+    cookies['1'] = '1'
+    cookies['2'] = '2'
+    render :index
+  end
+
   private
 
   def logged_in?

test/dummy_app/config/routes.rb

@@ -4,4 +4,5 @@
   post 'login', to: 'application#login'
   delete 'logout', to: 'application#logout'
   get 'exception', to: 'application#exception'
+  get 'cooks', to: 'application#cooks'
 end

test/handler_test.rb

@@ -15,6 +15,18 @@ class HandlerTest < LambySpec
       expect(result[:body]).must_match %r{<div id="logged_in">false</div>}
     end
 
+    it 'get - multiple cookies' do
+      event = TestHelpers::Events::HttpV2.create(
+        'rawPath' => '/production/cooks',
+        'requestContext' => { 'http' => {'path' => '/production/cooks'} }
+      )
+      result = Lamby.handler app, event, context, rack: :http
+      expect(result[:statusCode]).must_equal 200
+      expect(result[:headers]['Set-Cookie']).must_be_nil
+      expect(result[:cookies]).must_equal ["1=1; path=/", "2=2; path=/"]
+      expect(result[:body]).must_match %r{<h1>Hello Lamby</h1>}
+    end
+
     it 'get - image' do
       event = TestHelpers::Events::HttpV2.create(
         'rawPath' => '/production/image',
@@ -77,6 +89,18 @@ class HandlerTest < LambySpec
       expect(result[:body]).must_match %r{<div id="logged_in">false</div>}
     end
 
+    it 'get - multiple cookies' do
+      event = TestHelpers::Events::HttpV1.create(
+        'path' => '/production/cooks',
+        'requestContext' => { 'path' => '/production/cooks'}
+      )
+      result = Lamby.handler app, event, context, rack: :http
+      expect(result[:statusCode]).must_equal 200
+      expect(result[:headers]['Set-Cookie']).must_be_nil
+      expect(result[:multiValueHeaders]['Set-Cookie']).must_equal ["1=1; path=/", "2=2; path=/"]
+      expect(result[:body]).must_match %r{<h1>Hello Lamby</h1>}
+    end
+
     it 'get - image' do
       event = TestHelpers::Events::HttpV1.create(
         'path' => '/production/image',
@@ -139,6 +163,18 @@ class HandlerTest < LambySpec
       expect(result[:body]).must_match %r{<div id="logged_in">false</div>}
     end
 
+    it 'get - multiple cookies' do
+      event = TestHelpers::Events::Rest.create(
+        'path' => '/cooks',
+        'requestContext' => { 'path' => '/cooks'}
+      )
+      result = Lamby.handler app, event, context, rack: :rest
+      expect(result[:statusCode]).must_equal 200
+      expect(result[:headers]['Set-Cookie']).must_be_nil
+      expect(result[:multiValueHeaders]['Set-Cookie']).must_equal ["1=1; path=/", "2=2; path=/"]
+      expect(result[:body]).must_match %r{<h1>Hello Lamby</h1>}
+    end
+
     it 'get - image' do
       event = TestHelpers::Events::Rest.create(
         'path' => '/image',
@@ -201,6 +237,15 @@ class HandlerTest < LambySpec
       expect(result[:body]).must_match %r{<div id="logged_in">false</div>}
     end
 
+    it 'get - multiple cookies' do
+      event = TestHelpers::Events::Alb.create 'path' => '/cooks'
+      result = Lamby.handler app, event, context, rack: :rest
+      expect(result[:statusCode]).must_equal 200
+      expect(result[:headers]['Set-Cookie']).must_be_nil
+      expect(result[:multiValueHeaders]['Set-Cookie']).must_equal ["1=1; path=/", "2=2; path=/"]
+      expect(result[:body]).must_match %r{<h1>Hello Lamby</h1>}
+    end
+
     it 'get - image' do
       event = TestHelpers::Events::Alb.create 'path' => '/image'
       result = Lamby.handler app, event, context, rack: :alb
@@ -233,7 +278,7 @@ class HandlerTest < LambySpec
   private
 
   def session_cookie(result)
-    cookies = result[:headers]['Set-Cookie']
+    cookies = (result[:cookies] || result[:multiValueHeaders]['Set-Cookie'])[0]
     cookies.split('; ').detect { |x| x =~ /session=/ }
   end