From 5fdb0c4bedc57356dea242a4486a1e121a38628f Mon Sep 17 00:00:00 2001 From: Ken Collins Date: Wed, 3 Mar 2021 23:18:13 -0500 Subject: [PATCH] Support Multiple Set-Cookie Headers --- CHANGELOG.md | 6 +++ Gemfile.lock | 2 +- lib/lamby/handler.rb | 8 ++++ lib/lamby/rack_alb.rb | 5 +- lib/lamby/rack_http.rb | 12 +++++ lib/lamby/rack_rest.rb | 5 ++ lib/lamby/version.rb | 2 +- .../app/controllers/application_controller.rb | 6 +++ test/dummy_app/config/routes.rb | 1 + test/handler_test.rb | 47 ++++++++++++++++++- 10 files changed, 90 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1a97611..58a5acd 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/Gemfile.lock b/Gemfile.lock index 0b05df7..051999a 100644 --- a/Gemfile.lock +++ b/Gemfile.lock @@ -1,7 +1,7 @@ PATH remote: . specs: - lamby (2.5.3) + lamby (2.6.0) rack GEM diff --git a/lib/lamby/handler.rb b/lib/lamby/handler.rb index 0defe3d..199f4f2 100644 --- a/lib/lamby/handler.rb +++ b/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 diff --git a/lib/lamby/rack_alb.rb b/lib/lamby/rack_alb.rb index 4954108..6d2b0da 100644 --- a/lib/lamby/rack_alb.rb +++ b/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 diff --git a/lib/lamby/rack_http.rb b/lib/lamby/rack_http.rb index 7295760..74741de 100644 --- a/lib/lamby/rack_http.rb +++ b/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 diff --git a/lib/lamby/rack_rest.rb b/lib/lamby/rack_rest.rb index f8c8d23..c012d9f 100644 --- a/lib/lamby/rack_rest.rb +++ b/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 diff --git a/lib/lamby/version.rb b/lib/lamby/version.rb index d559e61..29c400c 100644 --- a/lib/lamby/version.rb +++ b/lib/lamby/version.rb @@ -1,3 +1,3 @@ module Lamby - VERSION = '2.5.3' + VERSION = '2.6.0' end diff --git a/test/dummy_app/app/controllers/application_controller.rb b/test/dummy_app/app/controllers/application_controller.rb index 7e12558..aa6af90 100644 --- a/test/dummy_app/app/controllers/application_controller.rb +++ b/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? diff --git a/test/dummy_app/config/routes.rb b/test/dummy_app/config/routes.rb index 173cfc0..cc2d389 100644 --- a/test/dummy_app/config/routes.rb +++ b/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 diff --git a/test/handler_test.rb b/test/handler_test.rb index ccdc093..ca0e8fa 100644 --- a/test/handler_test.rb +++ b/test/handler_test.rb @@ -15,6 +15,18 @@ class HandlerTest < LambySpec expect(result[:body]).must_match %r{
false
} 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{

Hello Lamby

} + 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{
false
} 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{

Hello Lamby

} + 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{
false
} 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{

Hello Lamby

} + 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{
false
} 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{

Hello Lamby

} + 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