Merge branch 'fix-cookies' into 'master'

Fix cookies

See merge request honeyryderchuck/httpx!36

lib/httpx/plugins/cookies.rb

@@ -3,20 +3,82 @@
 module HTTPX
   module Plugins
     module Cookies
+      using URIExtensions
+
+      class Store
+        def initialize(cookies = nil)
+          @store = Hash.new { |hash, origin| hash[origin] = HTTP::CookieJar.new }
+          return unless cookies
+
+          cookies = cookies.split(/ *; */) if cookies.is_a?(String)
+          @default_cookies = cookies.map do |cookie, v|
+            if cookie.is_a?(HTTP::Cookie)
+              cookie
+            else
+              HTTP::Cookie.new(cookie.to_s, v.to_s)
+            end
+          end
+        end
+
+        def set(origin, cookies)
+          return unless cookies
+
+          @store[origin].parse(cookies, origin)
+        end
+
+        def [](uri)
+          store = @store[uri.origin]
+          @default_cookies.each do |cookie|
+            c = cookie.dup
+            c.domain ||= uri.authority
+            c.path ||= uri.path
+            store.add(c)
+          end if @default_cookies
+          store
+        end
+      end
+
       def self.load_dependencies(*)
         require "http/cookie"
       end
 
       module InstanceMethods
-        def cookies(cookies)
-          branch(default_options.with_cookies(cookies))
-        end
-      end
+        attr_reader :cookies_store
 
-      module RequestMethods
         def initialize(*)
           super
-          @headers.cookies(@options.cookies, self)
+          @cookies_store = @options.cookies || Store.new
+        end
+
+        def with_cookies(cookies)
+          branch(default_options.with_cookies(cookies))
+        end
+
+        def wrap
+          return unless block_given?
+
+          super do |client|
+            old_cookies_store = @cookies_store
+            @cookies_store = old_cookies_store.dup
+            begin
+              yield client
+            ensure
+              @cookies_store = old_cookies_store
+            end
+          end
+        end
+
+        private
+
+        def on_response(request, response)
+          @cookies_store.set(request.origin, response.headers["set-cookie"])
+          super
+        end
+
+        def __build_req(*)
+          request = super
+          request.headers.cookies(@cookies_store[request.uri], request)
+          request
         end
       end
 
@@ -24,38 +86,20 @@ module HTTPX
         def cookies(jar, request)
           return unless jar
 
-          unless jar.is_a?(HTTP::CookieJar)
-            jar = jar.each_with_object(HTTP::CookieJar.new) do |(k, v), j|
-              cookie = k.is_a?(HTTP::Cookie) ? v : HTTP::Cookie.new(k.to_s, v.to_s)
-              cookie.domain = request.authority
-              cookie.path = request.path
-              j.add(cookie)
-            end
-          end
-          self["cookie"] = HTTP::Cookie.cookie_value(jar.cookies)
-        end
-      end
+          cookie_value = HTTP::Cookie.cookie_value(jar.cookies(request.uri))
+          return if cookie_value.empty?
 
-      module ResponseMethods
-        def cookie_jar
-          return @cookie_jar if defined?(@cookie_jar)
-          return nil unless headers.key?("set-cookie")
-
-          @cookie_jar ||= begin
-            jar = HTTP::CookieJar.new
-            jar.parse(headers["set-cookie"], @request.uri)
-            jar
-          end
+          add("cookie", cookie_value)
         end
-        alias_method :cookies, :cookie_jar
       end
 
       module OptionsMethods
         def self.included(klass)
           super
           klass.def_option(:cookies) do |cookies|
-            cookies.split(/ *; */) if cookies.is_a?(String)
-            cookies
+            return cookies if cookies.is_a?(Store)
+
+            Store.new(cookies)
           end
         end
       end

test/support/requests/plugins/cookies.rb

@@ -5,14 +5,12 @@ module Requests
     module Cookies
       def test_plugin_cookies_get
         client = HTTPX.plugin(:cookies)
-        assert client.respond_to?(:cookies), "client should be cookie-enabled"
         response = client.get(cookies_uri)
-        assert response.respond_to?(:cookies), "response should have cookies"
         body = json_body(response)
         assert body.key?("cookies")
         assert body["cookies"].empty?
 
-        session_response = client.cookies("abc" => "def").get(cookies_uri)
+        session_response = client.with_cookies("abc" => "def").get(cookies_uri)
         body = json_body(session_response)
         assert body.key?("cookies")
         assert body["cookies"]["abc"] == "def", "abc wasn't properly set"
@@ -22,21 +20,39 @@ module Requests
         client = HTTPX.plugin(:cookies)
         session_cookies = { "a" => "b", "c" => "d" }
         session_uri = cookies_set_uri(session_cookies)
-        session_response = client.get(cookies_set_uri(session_cookies))
-        assert session_response.status == 302, "response should redirect"
+        session_response = client.get(session_uri)
+        verify_status(session_response, 302)
+        verify_cookies(client.cookies_store[URI(session_uri)], session_cookies)
 
-        assert !session_response.cookies.nil?, "there should be cookies in the response"
-        response_cookies = session_response.cookie_jar
-        assert !response_cookies.empty?
-        response_cookies.cookies(session_uri).each do |cookie|
-          assert(session_cookies.one? { |k, v| k == cookie.name && v == cookie.value })
-        end
-
-        response = client.cookies(response_cookies).get(cookies_uri)
+        # first request sets the session
+        response = client.get(cookies_uri)
         body = json_body(response)
         assert body.key?("cookies")
-        assert body["cookies"]["a"] == "b"
-        assert body["cookies"]["c"] == "d"
+        verify_cookies(body["cookies"], session_cookies)
+
+        # second request reuses the session
+        extra_cookie_response = client.with_cookies("e" => "f").get(cookies_uri)
+        body = json_body(extra_cookie_response)
+        assert body.key?("cookies")
+        verify_cookies(body["cookies"], session_cookies.merge("e" => "f"))
+
+        # redirect to a different origin only uses the option cookies
+        other_origin_response = client.with_cookies("e" => "f").get(redirect_uri(origin("google.com")))
+        verify_status(other_origin_response, 302)
+        assert !other_origin_response.headers.key?("set-cookie"), "cookies should not transition to next origin"
+      end
+
+      def test_plugin_cookies_follow
+        client = HTTPX.plugins(:follow_redirects, :cookies)
+        session_cookies = { "a" => "b", "c" => "d" }
+        session_uri = cookies_set_uri(session_cookies)
+
+        response = client.get(session_uri)
+        verify_status(response, 200)
+        assert response.uri.to_s == cookies_uri
+        body = json_body(response)
+        assert body.key?("cookies")
+        verify_cookies(body["cookies"], session_cookies)
       end
 
       private
@@ -48,6 +64,19 @@ module Requests
       def cookies_set_uri(cookies)
         build_uri("/cookies/set?" + URI.encode_www_form(cookies))
       end
+
+      def verify_cookies(jar, cookies)
+        assert !jar.nil? && !jar.empty?, "there should be cookies in the response"
+        assert jar.all? { |cookie|
+          case cookie
+          when HTTP::Cookie
+            cookies.one? { |k, v| k == cookie.name && v == cookie.value }
+          else
+            cookie_name, cookie_value = cookie
+            cookies.one? { |k, v| k == cookie_name && v == cookie_value }
+          end
+        }, "jar should contain all expected cookies"
+      end
     end
   end
 end

test/support/requests/plugins/follow_redirects.rb

@@ -52,8 +52,8 @@ module Requests
 
       private
 
-      def redirect_uri
-        build_uri("/redirect-to?url=" + redirect_location)
+      def redirect_uri(redirect_uri = redirect_location)
+        build_uri("/redirect-to?url=" + redirect_uri)
       end
 
       def max_redirect_uri(n)