From 71b0bc6c0f03b714bbd4cc9ae73a94a878431971 Mon Sep 17 00:00:00 2001
From: evazion <noizave@gmail.com>
Date: Sat, 20 Jun 2020 16:48:00 -0500
Subject: [PATCH] danbooru::http: support tracking cookies between requests.

Allow cookies to be saved and sent back when making several requests in
a row. Usage:

    http = Danbooru::Http.use(:session)

    # saves the foo=42 cookie sent by the response.
    http.get("https://httpbin.org/cookies/set/foo/42")

    # sends back the foo=42 cookie from the previous request.
    http.get("https://httpbin.org/cookies")
---
 app/logical/danbooru/http.rb         | 12 ++++++++-
 app/logical/danbooru/http/session.rb | 37 ++++++++++++++++++++++++++++
 test/unit/danbooru_http_test.rb      | 12 +++++++++
 3 files changed, 60 insertions(+), 1 deletion(-)
 create mode 100644 app/logical/danbooru/http/session.rb

diff --git a/app/logical/danbooru/http.rb b/app/logical/danbooru/http.rb
index c57f85f0b..8fe06a0d1 100644
--- a/app/logical/danbooru/http.rb
+++ b/app/logical/danbooru/http.rb
@@ -1,6 +1,7 @@
 require "danbooru/http/html_adapter"
 require "danbooru/http/xml_adapter"
 require "danbooru/http/retriable"
+require "danbooru/http/session"
 
 module Danbooru
   class Http
@@ -13,7 +14,7 @@ module Danbooru
     attr_writer :cache, :max_size, :http
 
     class << self
-      delegate :get, :head, :put, :post, :delete, :cache, :follow, :max_size, :timeout, :auth, :basic_auth, :headers, :public_only, :download_media, to: :new
+      delegate :get, :head, :put, :post, :delete, :cache, :follow, :max_size, :timeout, :auth, :basic_auth, :headers, :cookies, :use, :public_only, :download_media, to: :new
     end
 
     def get(url, **options)
@@ -64,6 +65,14 @@ module Danbooru
       dup.tap { |o| o.http = o.http.headers(*args) }
     end
 
+    def cookies(*args)
+      dup.tap { |o| o.http = o.http.cookies(*args) }
+    end
+
+    def use(*args)
+      dup.tap { |o| o.http = o.http.use(*args) }
+    end
+
     # allow requests only to public IPs, not to local or private networks.
     def public_only
       dup.tap do |o|
@@ -143,6 +152,7 @@ module Danbooru
       @http ||=
         ::Danbooru::Http::ApplicationClient.new
         .follow(strict: false, max_hops: MAX_REDIRECTS)
+        .use(:session)
         .timeout(DEFAULT_TIMEOUT)
         .use(:auto_inflate)
         .headers(Danbooru.config.http_headers)
diff --git a/app/logical/danbooru/http/session.rb b/app/logical/danbooru/http/session.rb
new file mode 100644
index 000000000..3bb532928
--- /dev/null
+++ b/app/logical/danbooru/http/session.rb
@@ -0,0 +1,37 @@
+module Danbooru
+  class Http
+    class Session < HTTP::Feature
+      HTTP::Options.register_feature :session, self
+
+      attr_reader :cookie_jar
+
+      def initialize(cookie_jar: HTTP::CookieJar.new)
+        @cookie_jar = cookie_jar
+      end
+
+      def perform(request)
+        add_cookies(request)
+        response = yield request
+        save_cookies(response)
+        response
+      end
+
+      def add_cookies(request)
+        cookies = cookies_for_request(request)
+        request.headers["Cookie"] = cookies if cookies.present?
+      end
+
+      def cookies_for_request(request)
+        saved_cookies = cookie_jar.each(request.uri).map { |c| [c.name, c.value] }.to_h
+        request_cookies = HTTP::Cookie.cookie_value_to_hash(request.headers["Cookie"].to_s)
+        saved_cookies.merge(request_cookies).map { |name, value| "#{name}=#{value}" }.join("; ")
+      end
+
+      def save_cookies(response)
+        response.cookies.each do |cookie|
+          cookie_jar.add(cookie)
+        end
+      end
+    end
+  end
+end
diff --git a/test/unit/danbooru_http_test.rb b/test/unit/danbooru_http_test.rb
index ae347bfd8..b6ef81e7c 100644
--- a/test/unit/danbooru_http_test.rb
+++ b/test/unit/danbooru_http_test.rb
@@ -51,6 +51,18 @@ class DanbooruHttpTest < ActiveSupport::TestCase
         assert_equal(true, response.parse[:slideshow].present?)
       end
 
+      should "track cookies between requests" do
+        http = Danbooru::Http.use(:session)
+
+        resp1 = http.get("https://httpbin.org/cookies/set/abc/1")
+        resp2 = http.get("https://httpbin.org/cookies/set/def/2")
+        resp3 = http.get("https://httpbin.org/cookies")
+        assert_equal({ abc: "1", def: "2" }, resp3.parse["cookies"].symbolize_keys)
+
+        resp4 = http.cookies(def: 3, ghi: 4).get("https://httpbin.org/cookies")
+        assert_equal({ abc: "1", def: "3", ghi: "4" }, resp4.parse["cookies"].symbolize_keys)
+      end
+
       should "cache requests" do
         response1 = Danbooru::Http.cache(1.minute).get("https://httpbin.org/uuid")
         assert_equal(200, response1.status)