Merge pull request #1390 from Rosa/insert-middlewares-before-after

Allow inserting middleware at arbitrary points in the middleware stack

Author: dblock

CHANGELOG.md

@@ -5,6 +5,7 @@
 
 #### Features
 
+* [#1390](https://github.com/ruby-grape/grape/pull/1390): Allow inserting middleware at arbitrary points in the middleware stack - [@Rosa](https://github.com/Rosa).
 * [#1366](https://github.com/ruby-grape/grape/pull/1366): Store `message_key` on `Grape::Exceptions::Validation` - [@mkou](https://github.com/mkou).
 
 #### Fixes

README.md

@@ -2807,7 +2807,27 @@ Your middleware can overwrite application response as follows, except error case
 ```ruby
 class Overwriter < Grape::Middleware::Base
   def after
-    [200, { 'Content-Type' => 'text/plain' }, ['Overwrited.']]
+    [200, { 'Content-Type' => 'text/plain' }, ['Overwritten.']]
+  end
+end
+```
+
+You can add your custom middleware with `use`, that push the middleware onto the stack, and you can also control where the middleware is inserted using `insert`, `insert_before` and `insert_after`.
+
+```ruby
+class CustomOverwriter < Grape::Middleware::Base
+  def after
+    [200, { 'Content-Type' => 'text/plain' }, [@options[:message]]]
+  end
+end
+
+
+class API < Grape::API
+  use Overwriter
+  insert_before Overwriter, CustomOverwriter, message: 'Overwritten again.'
+  insert 0, CustomOverwriter, message: 'Overwrites all other middleware.'
+
+  get '/' do
   end
 end
 ```

lib/grape/dsl/middleware.rb

@@ -15,7 +15,21 @@ module ClassMethods
         # @param middleware_class [Class] The class of the middleware you'd like
         #   to inject.
         def use(middleware_class, *args, &block)
-          arr = [middleware_class, *args]
+          arr = [:use, middleware_class, *args]
+          arr << block if block_given?
+
+          namespace_stackable(:middleware, arr)
+        end
+
+        def insert_before(*args, &block)
+          arr = [:insert_before, *args]
+          arr << block if block_given?
+
+          namespace_stackable(:middleware, arr)
+        end
+
+        def insert_after(*args, &block)
+          arr = [:insert_after, *args]
           arr << block if block_given?
 
           namespace_stackable(:middleware, arr)

lib/grape/endpoint.rb

@@ -1,3 +1,5 @@
+require 'grape/middleware/stack'
+
 module Grape
   # An Endpoint is the proxy scope in which all routing
   # blocks are executed. In other words, any methods
@@ -243,45 +245,43 @@ def run
     end
 
     def build_stack
-      b = Rack::Builder.new
-
-      b.use Rack::Head
-      b.use Grape::Middleware::Error,
-            format: namespace_inheritable(:format),
-            content_types: namespace_stackable_with_hash(:content_types),
-            default_status: namespace_inheritable(:default_error_status),
-            rescue_all: namespace_inheritable(:rescue_all),
-            default_error_formatter: namespace_inheritable(:default_error_formatter),
-            error_formatters: namespace_stackable_with_hash(:error_formatters),
-            rescue_options: namespace_stackable_with_hash(:rescue_options) || {},
-            rescue_handlers: namespace_stackable_with_hash(:rescue_handlers) || {},
-            base_only_rescue_handlers: namespace_stackable_with_hash(:base_only_rescue_handlers) || {},
-            all_rescue_handler: namespace_inheritable(:all_rescue_handler)
-
-      (namespace_stackable(:middleware) || []).each do |m|
-        m = m.dup
-        block = m.pop if m.last.is_a?(Proc)
-        block ? b.use(*m, &block) : b.use(*m)
-      end
+      ms = Grape::Middleware::Stack.new
+
+      ms.use Rack::Head
+      ms.use Grape::Middleware::Error,
+             format: namespace_inheritable(:format),
+             content_types: namespace_stackable_with_hash(:content_types),
+             default_status: namespace_inheritable(:default_error_status),
+             rescue_all: namespace_inheritable(:rescue_all),
+             default_error_formatter: namespace_inheritable(:default_error_formatter),
+             error_formatters: namespace_stackable_with_hash(:error_formatters),
+             rescue_options: namespace_stackable_with_hash(:rescue_options) || {},
+             rescue_handlers: namespace_stackable_with_hash(:rescue_handlers) || {},
+             base_only_rescue_handlers: namespace_stackable_with_hash(:base_only_rescue_handlers) || {},
+             all_rescue_handler: namespace_inheritable(:all_rescue_handler)
+
+      ms.merge_with(namespace_stackable(:middleware) || [])
 
       if namespace_inheritable(:version)
-        b.use Grape::Middleware::Versioner.using(namespace_inheritable(:version_options)[:using]),
-              versions: namespace_inheritable(:version) ? namespace_inheritable(:version).flatten : nil,
-              version_options: namespace_inheritable(:version_options),
-              prefix: namespace_inheritable(:root_prefix)
-
+        ms.use Grape::Middleware::Versioner.using(namespace_inheritable(:version_options)[:using]),
+               versions: namespace_inheritable(:version) ? namespace_inheritable(:version).flatten : nil,
+               version_options: namespace_inheritable(:version_options),
+               prefix: namespace_inheritable(:root_prefix)
       end
 
-      b.use Grape::Middleware::Formatter,
-            format: namespace_inheritable(:format),
-            default_format: namespace_inheritable(:default_format) || :txt,
-            content_types: namespace_stackable_with_hash(:content_types),
-            formatters: namespace_stackable_with_hash(:formatters),
-            parsers: namespace_stackable_with_hash(:parsers)
+      ms.use Grape::Middleware::Formatter,
+             format: namespace_inheritable(:format),
+             default_format: namespace_inheritable(:default_format) || :txt,
+             content_types: namespace_stackable_with_hash(:content_types),
+             formatters: namespace_stackable_with_hash(:formatters),
+             parsers: namespace_stackable_with_hash(:parsers)
+
+      builder = Rack::Builder.new
+      ms.build(builder)
 
-      b.run ->(env) { env[Grape::Env::API_ENDPOINT].run }
+      builder.run ->(env) { env[Grape::Env::API_ENDPOINT].run }
 
-      b.to_app
+      builder.to_app
     end
 
     def build_helpers

lib/grape/middleware/stack.rb

@@ -0,0 +1,97 @@
+module Grape
+  module Middleware
+    # Class to handle the stack of middlewares based on ActionDispatch::MiddlewareStack
+    # It allows to insert and insert after
+    class Stack
+      class Middleware
+        attr_reader :args, :block, :klass
+
+        def initialize(klass, *args, &block)
+          @klass = klass
+          @args = args
+          @block = block
+        end
+
+        def name
+          klass.name
+        end
+
+        def ==(other)
+          case other
+          when Middleware
+            klass == other.klass
+          when Class
+            klass == other
+          end
+        end
+
+        def inspect
+          klass.to_s
+        end
+      end
+
+      include Enumerable
+
+      attr_accessor :middlewares
+
+      def initialize
+        @middlewares = []
+      end
+
+      def each
+        @middlewares.each { |x| yield x }
+      end
+
+      def size
+        middlewares.size
+      end
+
+      def last
+        middlewares.last
+      end
+
+      def [](i)
+        middlewares[i]
+      end
+
+      def insert(index, *args, &block)
+        index = assert_index(index, :before)
+        middleware = self.class::Middleware.new(*args, &block)
+        middlewares.insert(index, middleware)
+      end
+
+      alias insert_before insert
+
+      def insert_after(index, *args, &block)
+        index = assert_index(index, :after)
+        insert(index + 1, *args, &block)
+      end
+
+      def use(*args, &block)
+        middleware = self.class::Middleware.new(*args, &block)
+        middlewares.push(middleware)
+      end
+
+      def merge_with(other)
+        other.each do |operation, *args|
+          block = args.pop if args.last.is_a?(Proc)
+          block ? send(operation, *args, &block) : send(operation, *args)
+        end
+      end
+
+      def build(builder)
+        middlewares.each do |m|
+          m.block ? builder.use(m.klass, *m.args, &m.block) : builder.use(m.klass, *m.args)
+        end
+      end
+
+      protected
+
+      def assert_index(index, where)
+        i = index.is_a?(Integer) ? index : middlewares.index(index)
+        raise "No such middleware to insert #{where}: #{index.inspect}" unless i
+        i
+      end
+    end
+  end
+end

lib/grape/util/stackable_values.rb

@@ -3,16 +3,16 @@ module Util
     class StackableValues
       attr_accessor :inherited_values
       attr_reader :new_values
-      attr_reader :froozen_values
+      attr_reader :frozen_values
 
       def initialize(inherited_values = {})
         @inherited_values = inherited_values
         @new_values = {}
-        @froozen_values = {}
+        @frozen_values = {}
       end
 
       def [](name)
-        return @froozen_values[name] if @froozen_values.key? name
+        return @frozen_values[name] if @frozen_values.key? name
 
         value = []
         value.concat(@inherited_values[name]) if @inherited_values[name]
@@ -21,7 +21,7 @@ def [](name)
       end
 
       def []=(name, value)
-        raise if @froozen_values.key? name
+        raise if @frozen_values.key? name
         @new_values[name] ||= []
         @new_values[name].push value
       end
@@ -43,7 +43,7 @@ def to_hash
       end
 
       def freeze_value(key)
-        @froozen_values[key] = self[key].freeze
+        @frozen_values[key] = self[key].freeze
       end
 
       def initialize_copy(other)

spec/grape/api_spec.rb

@@ -1044,7 +1044,7 @@ def call(env)
     describe '.middleware' do
       it 'includes middleware arguments from settings' do
         subject.use ApiSpec::PhonyMiddleware, 'abc', 123
-        expect(subject.middleware).to eql [[ApiSpec::PhonyMiddleware, 'abc', 123]]
+        expect(subject.middleware).to eql [[:use, ApiSpec::PhonyMiddleware, 'abc', 123]]
       end
 
       it 'includes all middleware from stacked settings' do
@@ -1053,17 +1053,17 @@ def call(env)
         subject.use ApiSpec::PhonyMiddleware, 'foo'
 
         expect(subject.middleware).to eql [
-          [ApiSpec::PhonyMiddleware, 123],
-          [ApiSpec::PhonyMiddleware, 'abc'],
-          [ApiSpec::PhonyMiddleware, 'foo']
+          [:use, ApiSpec::PhonyMiddleware, 123],
+          [:use, ApiSpec::PhonyMiddleware, 'abc'],
+          [:use, ApiSpec::PhonyMiddleware, 'foo']
         ]
       end
     end
 
     describe '.use' do
       it 'adds middleware' do
         subject.use ApiSpec::PhonyMiddleware, 123
-        expect(subject.middleware).to eql [[ApiSpec::PhonyMiddleware, 123]]
+        expect(subject.middleware).to eql [[:use, ApiSpec::PhonyMiddleware, 123]]
       end
 
       it 'does not show up outside the namespace' do
@@ -1074,8 +1074,8 @@ def call(env)
           inner_middleware = middleware
         end
 
-        expect(subject.middleware).to eql [[ApiSpec::PhonyMiddleware, 123]]
-        expect(inner_middleware).to eql [[ApiSpec::PhonyMiddleware, 123], [ApiSpec::PhonyMiddleware, 'abc']]
+        expect(subject.middleware).to eql [[:use, ApiSpec::PhonyMiddleware, 123]]
+        expect(inner_middleware).to eql [[:use, ApiSpec::PhonyMiddleware, 123], [:use, ApiSpec::PhonyMiddleware, 'abc']]
       end
 
       it 'calls the middleware' do
@@ -1091,7 +1091,7 @@ def call(env)
       it 'adds a block if one is given' do
         block = -> {}
         subject.use ApiSpec::PhonyMiddleware, &block
-        expect(subject.middleware).to eql [[ApiSpec::PhonyMiddleware, block]]
+        expect(subject.middleware).to eql [[:use, ApiSpec::PhonyMiddleware, block]]
       end
 
       it 'uses a block if one is given' do
@@ -1132,7 +1132,50 @@ def before
         expect(last_response.body).to eq('Caught in the Net')
       end
     end
+
+    describe '.insert_before' do
+      it 'runs before a given middleware' do
+        m = Class.new(Grape::Middleware::Base) do
+          def call(env)
+            env['phony.args'] ||= []
+            env['phony.args'] << @options[:message]
+            @app.call(env)
+          end
+        end
+
+        subject.use ApiSpec::PhonyMiddleware, 'hello'
+        subject.insert_before ApiSpec::PhonyMiddleware, m, message: 'bye'
+        subject.get '/' do
+          env['phony.args'].join(' ')
+        end
+
+        get '/'
+        expect(last_response.body).to eql 'bye hello'
+      end
+    end
+
+    describe '.insert_after' do
+      it 'runs after a given middleware' do
+        m = Class.new(Grape::Middleware::Base) do
+          def call(env)
+            env['phony.args'] ||= []
+            env['phony.args'] << @options[:message]
+            @app.call(env)
+          end
+        end
+
+        subject.use ApiSpec::PhonyMiddleware, 'hello'
+        subject.insert_after ApiSpec::PhonyMiddleware, m, message: 'bye'
+        subject.get '/' do
+          env['phony.args'].join(' ')
+        end
+
+        get '/'
+        expect(last_response.body).to eql 'hello bye'
+      end
+    end
   end
+
   describe '.http_basic' do
     it 'protects any resources on the same scope' do
       subject.http_basic do |u, _p|