lib/grape/middleware/base.rb in grape-0.14.0 vs lib/grape/middleware/base.rb in grape-0.15.0

@@ -1,11 +1,15 @@
+require 'grape/dsl/headers'
+
 module Grape
   module Middleware
     class Base
       attr_reader :app, :env, :options
       TEXT_HTML = 'text/html'.freeze
 
+      include Grape::DSL::Headers
+
       # @param [Rack Application] app The standard argument for a Rack middleware.
       # @param [Hash] options A hash of options, simply stored for use by subclasses.
       def initialize(app, options = {})
         @app = app
         @options = default_options.merge(options)
@@ -20,12 +24,24 @@
       end
 
       def call!(env)
         @env = env
         before
-        @app_response = @app.call(@env)
-        after || @app_response
+        begin
+          @app_response = @app.call(@env)
+        ensure
+          begin
+            after_response = after
+          rescue StandardError => e
+            warn "caught error of type #{e.class} in after callback inside #{self.class.name} : #{e.message}"
+            raise e
+          end
+        end
+
+        response = after_response || @app_response
+        merge_headers response
+        response
       end
 
       # @abstract
       # Called before the application is called in the middleware lifecycle.
       def before
@@ -58,9 +74,19 @@
         types_without_params = {}
         content_types.each_pair do |k, v|
           types_without_params[v.split(';').first] = k
         end
         types_without_params
+      end
+
+      private
+
+      def merge_headers(response)
+        return unless headers.is_a?(Hash)
+        case response
+        when Rack::Response then response.headers.merge!(headers)
+        when Array          then response[1].merge!(headers)
+        end
       end
     end
   end
 end