fluke/devops-service/sinatra/methods_with_headers.rb

require "sinatra/base"
require 'rack/accept_media_types'

module Sinatra

  class Base
    class << self

      # TODO: add protect! method
      def get_with_headers path, opt={}, &block
        headers = opt.delete(:headers) || []
        before path do
          check_headers *headers
          protect!
        end

        get path, opt, &block
      end

      def post_with_headers path, opt={}, &block
        headers = opt.delete(:headers) || []
        before path do
          check_headers *headers
          protect!
        end
        post_with_statistic path, opt, &block
      end

      def post_with_statistic path, opt={}, &block
        post path, opt, &block

        after path do
          statistic
        end
      end

      def put_with_headers path, opt={}, &block
        headers = opt.delete(:headers) || []
        before path do
          check_headers *headers
          protect!
        end

        put path, opt, &block

        after path do
          statistic
        end
      end

      def delete_with_headers path, opt={}, &block
        headers = opt.delete(:headers) || []
        before path do
          check_headers *headers
          protect!
        end

        delete path, opt, &block

        after path do
          statistic
        end
      end

      def multi_routes path, opts={}, hash={}
        headers = opts.delete(:headers) || []
        before path do
          if request.get?
            check_headers :accept
          else
            check_headers *headers
          end
          protect!
        end

        hash.each do |method, block|
          route(method, path, opts, &block)
        end

        after path do
          statistic
        end
      end
    end

    # Check request headers
    def check_headers *headers
      ha = (headers.empty? ? [:accept, :content_type] : headers)
      ha.each do |h|
        case h
        when :accept, "accept"
          accept_json
        when :content_type, "content_type"
          request_json
        end
      end
    end

    # Check Accept header
    #
    # Can client works with JSON?
    def accept_json
      types = request.accept_media_types
      unless types.include?('application/json') or types.include?("*/*")
        response.headers['Accept'] = 'application/json'
        halt_response("Accept header should contains 'application/json' type", 406)
      end
    end

    # Check Content-Type header
    def request_json
      halt_response("Content-Type should be 'application/json'", 415) if request.media_type.nil? or request.media_type != 'application/json'
    end

  end
end