diff --git a/lib/httpx/plugins/cookies.rb b/lib/httpx/plugins/cookies.rb index 8caa3e67..fdd4a6e1 100644 --- a/lib/httpx/plugins/cookies.rb +++ b/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 diff --git a/test/support/requests/plugins/cookies.rb b/test/support/requests/plugins/cookies.rb index 5554388d..229b7d96 100644 --- a/test/support/requests/plugins/cookies.rb +++ b/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 diff --git a/test/support/requests/plugins/follow_redirects.rb b/test/support/requests/plugins/follow_redirects.rb index 1f728edc..c0132060 100644 --- a/test/support/requests/plugins/follow_redirects.rb +++ b/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)