From bd25be95f52f99e69e573e44d295751478bfed63 Mon Sep 17 00:00:00 2001
From: evazion <noizave@gmail.com>
Date: Sun, 21 Jun 2020 18:27:32 -0500
Subject: [PATCH] danbooru::http: factor out cache feature.

Fixes a bug with cookies stored by the `session` feature not being sent
with cached requests.
---
 app/logical/danbooru/http.rb       | 32 +++++++-----------------------
 app/logical/danbooru/http/cache.rb | 30 ++++++++++++++++++++++++++++
 test/unit/danbooru_http_test.rb    | 31 +++++++++++++++++++++++++----
 3 files changed, 64 insertions(+), 29 deletions(-)
 create mode 100644 app/logical/danbooru/http/cache.rb

diff --git a/app/logical/danbooru/http.rb b/app/logical/danbooru/http.rb
index 8e8655965..6bbc7db87 100644
--- a/app/logical/danbooru/http.rb
+++ b/app/logical/danbooru/http.rb
@@ -1,5 +1,6 @@
 require "danbooru/http/html_adapter"
 require "danbooru/http/xml_adapter"
+require "danbooru/http/cache"
 require "danbooru/http/redirector"
 require "danbooru/http/retriable"
 require "danbooru/http/session"
@@ -12,7 +13,7 @@ module Danbooru
     DEFAULT_TIMEOUT = 10
     MAX_REDIRECTS = 5
 
-    attr_accessor :cache, :max_size, :http
+    attr_accessor :max_size, :http
 
     class << self
       delegate :get, :head, :put, :post, :delete, :cache, :follow, :max_size, :timeout, :auth, :basic_auth, :headers, :cookies, :use, :public_only, :download_media, to: :new
@@ -49,10 +50,6 @@ module Danbooru
       request(:delete, url, **options)
     end
 
-    def cache(expiry)
-      dup.tap { |o| o.cache = expiry.to_i }
-    end
-
     def follow(*args)
       dup.tap { |o| o.http = o.http.follow(*args) }
     end
@@ -85,6 +82,10 @@ module Danbooru
       dup.tap { |o| o.http = o.http.use(*args) }
     end
 
+    def cache(expires_in)
+      use(cache: { expires_in: expires_in })
+    end
+
     # allow requests only to public IPs, not to local or private networks.
     def public_only
       dup.tap do |o|
@@ -128,11 +129,7 @@ module Danbooru
     protected
 
     def request(method, url, **options)
-      if @cache.present?
-        cached_request(method, url, **options)
-      else
-        raw_request(method, url, **options)
-      end
+      http.send(method, url, **options)
     rescue ValidatingSocket::ProhibitedIpError
       fake_response(597, "")
     rescue HTTP::Redirector::TooManyRedirectsError
@@ -141,21 +138,6 @@ module Danbooru
       fake_response(599, "")
     end
 
-    def cached_request(method, url, **options)
-      key = Cache.hash({ method: method, url: url, headers: http.default_options.headers.to_h, **options }.to_json)
-
-      cached_response = Cache.get(key, @cache) do
-        response = raw_request(method, url, **options)
-        { status: response.status, body: response.to_s, headers: response.headers.to_h, version: "1.1" }
-      end
-
-      ::HTTP::Response.new(**cached_response)
-    end
-
-    def raw_request(method, url, **options)
-      http.send(method, url, **options)
-    end
-
     def fake_response(status, body)
       ::HTTP::Response.new(status: status, version: "1.1", body: ::HTTP::Response::Body.new(body))
     end
diff --git a/app/logical/danbooru/http/cache.rb b/app/logical/danbooru/http/cache.rb
new file mode 100644
index 000000000..43932bd28
--- /dev/null
+++ b/app/logical/danbooru/http/cache.rb
@@ -0,0 +1,30 @@
+module Danbooru
+  class Http
+    class Cache < HTTP::Feature
+      HTTP::Options.register_feature :cache, self
+
+      attr_reader :expires_in
+
+      def initialize(expires_in:)
+        @expires_in = expires_in
+      end
+
+      def perform(request, &block)
+        ::Cache.get(cache_key(request), expires_in) do
+          response = yield request
+
+          # XXX hack to remove connection state from response body so we can serialize it for caching.
+          response.flush
+          response.body.instance_variable_set(:@connection, nil)
+          response.body.instance_variable_set(:@stream, nil)
+
+          response
+        end
+      end
+
+      def cache_key(request)
+        "http:" + ::Cache.hash({ method: request.verb, url: request.uri.to_s, headers: request.headers.sort }.to_json)
+      end
+    end
+  end
+end
diff --git a/test/unit/danbooru_http_test.rb b/test/unit/danbooru_http_test.rb
index b6ef81e7c..aa73976ea 100644
--- a/test/unit/danbooru_http_test.rb
+++ b/test/unit/danbooru_http_test.rb
@@ -62,14 +62,37 @@ class DanbooruHttpTest < ActiveSupport::TestCase
         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
+    end
 
-      should "cache requests" do
-        response1 = Danbooru::Http.cache(1.minute).get("https://httpbin.org/uuid")
+    context "cache feature" do
+      should "cache multiple requests to the same url" do
+        http = Danbooru::Http.cache(1.hour)
+
+        response1 = http.get("https://httpbin.org/uuid")
         assert_equal(200, response1.status)
 
-        response2 = Danbooru::Http.cache(1.minute).get("https://httpbin.org/uuid")
+        response2 = http.get("https://httpbin.org/uuid")
         assert_equal(200, response2.status)
-        assert_equal(response2.body, response1.body)
+        assert_equal(response2.to_s, response1.to_s)
+      end
+
+      should "cache cookies correctly" do
+        http = Danbooru::Http.cache(1.hour)
+
+        resp1 = http.get("https://httpbin.org/cookies")
+        resp2 = http.get("https://httpbin.org/cookies/set/abc/1")
+        resp3 = http.get("https://httpbin.org/cookies/set/def/2")
+        resp4 = http.get("https://httpbin.org/cookies")
+
+        assert_equal(200, resp1.status)
+        assert_equal(200, resp2.status)
+        assert_equal(200, resp3.status)
+        assert_equal(200, resp4.status)
+
+        assert_equal({}, resp1.parse["cookies"].symbolize_keys)
+        assert_equal({ abc: "1" }, resp2.parse["cookies"].symbolize_keys)
+        assert_equal({ abc: "1", def: "2" }, resp3.parse["cookies"].symbolize_keys)
+        assert_equal({ abc: "1", def: "2" }, resp4.parse["cookies"].symbolize_keys)
       end
     end