From 5d4057da0343740a7310476f543090c0da89c936 Mon Sep 17 00:00:00 2001
From: fatkodima <fatkodima123@gmail.com>
Date: Thu, 17 Oct 2019 14:30:45 +0300
Subject: [PATCH] Add ability to use multiple instances of middleware

---
 lib/rack/attack.rb                       | 11 ++++++--
 spec/acceptance/middleware_spec.rb       | 36 ++++++++++++++++++++++++
 spec/acceptance/rails_middleware_spec.rb | 20 -------------
 3 files changed, 45 insertions(+), 22 deletions(-)
 create mode 100644 spec/acceptance/middleware_spec.rb
 delete mode 100644 spec/acceptance/rails_middleware_spec.rb

diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb
index c9094b21..de273f4c 100644
--- a/lib/rack/attack.rb
+++ b/lib/rack/attack.rb
@@ -96,9 +96,16 @@ def reset!
 
     attr_reader :configuration
 
-    def initialize(app)
+    def initialize(app, &block)
       @app = app
-      @configuration = self.class.configuration
+      @configuration =
+        if block_given?
+          configuration = Configuration.new
+          configuration.instance_exec(&block)
+          configuration
+        else
+          self.class.configuration
+        end
     end
 
     def call(env)
diff --git a/spec/acceptance/middleware_spec.rb b/spec/acceptance/middleware_spec.rb
new file mode 100644
index 00000000..5036e519
--- /dev/null
+++ b/spec/acceptance/middleware_spec.rb
@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+require_relative "../spec_helper"
+
+describe "Middleware" do
+  it "is used in Rails by default" do
+    skip unless defined?(Rails)
+
+    app = Class.new(Rails::Application) do
+      config.eager_load = false
+      config.logger = Logger.new(nil) # avoid creating the log/ directory automatically
+      config.cache_store = :null_store # avoid creating tmp/ directory for cache
+    end
+
+    app.initialize!
+    assert app.middleware.include?(Rack::Attack)
+  end
+
+  def app
+    Rack::Builder.new do
+      use Rack::Attack do
+        blocklist_ip("1.2.3.4")
+      end
+
+      run lambda { |_env| [200, {}, ['Hello World']] }
+    end.to_app
+  end
+
+  it "can be configured via a block" do
+    get "/", {}, "REMOTE_ADDR" => "1.2.3.4"
+    assert_equal 403, last_response.status
+
+    get "/", {}, "REMOTE_ADDR" => "4.3.2.1"
+    assert_equal 200, last_response.status
+  end
+end
diff --git a/spec/acceptance/rails_middleware_spec.rb b/spec/acceptance/rails_middleware_spec.rb
deleted file mode 100644
index 31dc6209..00000000
--- a/spec/acceptance/rails_middleware_spec.rb
+++ /dev/null
@@ -1,20 +0,0 @@
-# frozen_string_literal: true
-
-require_relative "../spec_helper"
-
-if defined?(Rails::Application)
-  describe "Middleware for Rails" do
-    before do
-      @app = Class.new(Rails::Application) do
-        config.eager_load = false
-        config.logger = Logger.new(nil) # avoid creating the log/ directory automatically
-        config.cache_store = :null_store # avoid creating tmp/ directory for cache
-      end
-    end
-
-    it "is used by default" do
-      @app.initialize!
-      assert @app.middleware.include?(Rack::Attack)
-    end
-  end
-end