Add ability to customize response in case of error by passing callable object

README.md

@@ -44,6 +44,8 @@ $ gem install rack-jwt
 
 * `exclude` : optional : Array : An Array of path strings representing paths that should not be checked for the presence of a valid JWT token. Excludes sub-paths as of specified paths as well (e.g. `%w(/docs)` excludes `/docs/some/thing.html` also). Each path should start with a `/`. If a path matches the current request path this entire middleware is skipped and no authentication or verification of tokens takes place.
 
+* `on_error` : optional : Callable : An object which responds to `call` method with single `error` parameter. `error` parameter is one of `Rack::JWT::Auth::ERRORS_TO_RESCUE`. `on_error` callable object will be called if one of `Rack::JWT::Auth::ERRORS_TO_RESCUE` raised. For default handler check `Rack::JWT::Auth#default_on_error`.
+
 ## Example Server-Side Config
 
 Where `my_args` is a `Hash` containing valid keys. See `spec/example_spec.rb`

lib/rack/jwt/auth.rb

@@ -38,6 +38,25 @@ class Auth
         )$
       }x
 
+      JWT_DECODE_ERRORS = [
+        ::JWT::DecodeError,
+        ::JWT::VerificationError,
+        ::JWT::ExpiredSignature,
+        ::JWT::IncorrectAlgorithm,
+        ::JWT::ImmatureSignature,
+        ::JWT::InvalidIssuerError,
+        ::JWT::InvalidIatError,
+        ::JWT::InvalidAudError,
+        ::JWT::InvalidSubError,
+        ::JWT::InvalidJtiError,
+        ::JWT::InvalidPayload,
+      ].freeze
+
+      MissingAuthHeader = Class.new(StandardError)
+      InvalidAuthHeaderFormat = Class.new(StandardError)
+
+      ERRORS_TO_RESCUE = (JWT_DECODE_ERRORS + [MissingAuthHeader, InvalidAuthHeaderFormat]).freeze
+
       # Initialization should fail fast with an ArgumentError
       # if any args are invalid.
       def initialize(app, opts = {})
@@ -47,7 +66,9 @@ def initialize(app, opts = {})
         @options = opts.fetch(:options, {})
         @exclude = opts.fetch(:exclude, [])
 
-        @secret  = @secret.strip if @secret.is_a?(String)
+        @on_error = opts.fetch(:on_error, method(:default_on_error))
+
+        @secret = @secret.strip if @secret.is_a?(String)
         @options[:algorithm] = DEFAULT_ALGORITHM if @options[:algorithm].nil?
 
         check_secret_type!
@@ -57,15 +78,12 @@ def initialize(app, opts = {})
         check_options_type!
         check_valid_algorithm!
         check_exclude_type!
+        check_on_error_callable!
       end
 
       def call(env)
         if path_matches_excluded_path?(env)
           @app.call(env)
-        elsif missing_auth_header?(env)
-          return_error('Missing Authorization header')
-        elsif invalid_auth_header?(env)
-          return_error('Invalid Authorization header format')
         else
           verify_token(env)
         end
@@ -74,36 +92,19 @@ def call(env)
       private
 
       def verify_token(env)
+        raise MissingAuthHeader if missing_auth_header?(env)
+        raise InvalidAuthHeaderFormat if invalid_auth_header?(env)
+
         # extract the token from the Authorization: Bearer header
         # with a regex capture group.
         token = BEARER_TOKEN_REGEX.match(env['HTTP_AUTHORIZATION'])[1]
 
-        begin
-          decoded_token = Token.decode(token, @secret, @verify, @options)
-          env['jwt.payload'] = decoded_token.first
-          env['jwt.header'] = decoded_token.last
-          @app.call(env)
-        rescue ::JWT::VerificationError
-          return_error('Invalid JWT token : Signature Verification Error')
-        rescue ::JWT::ExpiredSignature
-          return_error('Invalid JWT token : Expired Signature (exp)')
-        rescue ::JWT::IncorrectAlgorithm
-          return_error('Invalid JWT token : Incorrect Key Algorithm')
-        rescue ::JWT::ImmatureSignature
-          return_error('Invalid JWT token : Immature Signature (nbf)')
-        rescue ::JWT::InvalidIssuerError
-          return_error('Invalid JWT token : Invalid Issuer (iss)')
-        rescue ::JWT::InvalidIatError
-          return_error('Invalid JWT token : Invalid Issued At (iat)')
-        rescue ::JWT::InvalidAudError
-          return_error('Invalid JWT token : Invalid Audience (aud)')
-        rescue ::JWT::InvalidSubError
-          return_error('Invalid JWT token : Invalid Subject (sub)')
-        rescue ::JWT::InvalidJtiError
-          return_error('Invalid JWT token : Invalid JWT ID (jti)')
-        rescue ::JWT::DecodeError
-          return_error('Invalid JWT token : Decode Error')
-        end
+        decoded_token = Token.decode(token, @secret, @verify, @options)
+        env['jwt.payload'] = decoded_token.first
+        env['jwt.header'] = decoded_token.last
+        @app.call(env)
+      rescue *ERRORS_TO_RESCUE => e
+        @on_error.call(e)
       end
 
       def check_secret_type!
@@ -166,6 +167,12 @@ def check_exclude_type!
         end
       end
 
+      def check_on_error_callable!
+        unless @on_error.respond_to?(:call)
+          raise ArgumentError, 'on_error argument must respond to call'
+        end
+      end
+
       def path_matches_excluded_path?(env)
         @exclude.any? { |ex| env['PATH_INFO'].start_with?(ex) }
       end
@@ -182,7 +189,23 @@ def missing_auth_header?(env)
         env['HTTP_AUTHORIZATION'].nil? || env['HTTP_AUTHORIZATION'].strip.empty?
       end
 
-      def return_error(message)
+      def default_on_error(error)
+        error_message = {
+          ::JWT::DecodeError => 'Invalid JWT token : Decode Error',
+          ::JWT::VerificationError => 'Invalid JWT token : Signature Verification Error',
+          ::JWT::ExpiredSignature => 'Invalid JWT token : Expired Signature (exp)',
+          ::JWT::IncorrectAlgorithm => 'Invalid JWT token : Incorrect Key Algorithm',
+          ::JWT::ImmatureSignature => 'Invalid JWT token : Immature Signature (nbf)',
+          ::JWT::InvalidIssuerError => 'Invalid JWT token : Invalid Issuer (iss)',
+          ::JWT::InvalidIatError => 'Invalid JWT token : Invalid Issued At (iat)',
+          ::JWT::InvalidAudError => 'Invalid JWT token : Invalid Audience (aud)',
+          ::JWT::InvalidSubError => 'Invalid JWT token : Invalid Subject (sub)',
+          ::JWT::InvalidJtiError => 'Invalid JWT token : Invalid JWT ID (jti)',
+          ::JWT::InvalidPayload => 'Invalid JWT token : Invalid Payload',
+          MissingAuthHeader => 'Missing Authorization header',
+          InvalidAuthHeaderFormat => 'Invalid Authorization header format'
+        }
+        message = error_message.fetch(error.class, 'Default')
         body    = { error: message }.to_json
         headers = { 'Content-Type' => 'application/json', 'Content-Length' => body.bytesize.to_s }
 

spec/auth_spec.rb

@@ -145,5 +145,16 @@
         end
       end
     end
+
+    describe 'on_error' do
+      let(:app) { Rack::JWT::Auth.new(inner_app, secret: secret) }
+
+      describe 'when non-callable type provided' do
+        it 'raises an exception' do
+          args = { secret: secret, on_error: [] }
+          expect { Rack::JWT::Auth.new(inner_app, args) }.to raise_error(ArgumentError)
+        end
+      end
+    end
   end
 end

spec/on_error_spec.rb

@@ -0,0 +1,160 @@
+require 'spec_helper'
+
+describe Rack::JWT::Auth do
+  let(:issuer)  { Rack::JWT::Token }
+  let(:secret)  { 'secret' } # use 'secret to match hardcoded 'secret' @ http://jwt.io'
+  let(:verify)  { true }
+  let(:payload) { { foo: 'bar' } }
+
+  let(:inner_app) do
+    ->(env) { [200, env, [payload.to_json]] }
+  end
+
+  let(:app) { Rack::JWT::Auth.new(inner_app, secret: secret) }
+
+  describe 'handles the exception' do
+    before(:each) do
+      header 'Authorization', "Bearer #{issuer.encode(payload, secret, 'HS256')}"
+      get('/')
+    end
+
+    let(:on_error) do
+      lambda do |error|
+        message =
+          if ::Rack::JWT::Auth::JWT_DECODE_ERRORS.include?(error.class)
+            'Invalid JWT token'
+          elsif error.is_a?(::Rack::JWT::Auth::MissingAuthHeader)
+            'Missing Authorization header'
+          elsif error.is_a?(::Rack::JWT::Auth::InvalidAuthHeaderFormat)
+            'Invalid Authorization header format'
+          end
+        body    = { error: message }.to_json
+        headers = { 'Content-Type' => 'application/json', 'Content-Length' => body.bytesize.to_s }
+
+        [401, headers, [body]]
+      end
+    end
+
+    let(:app) { Rack::JWT::Auth.new(inner_app, secret: 'secret', on_error: on_error) }
+
+    describe '::JWT::VerificationError' do
+      let(:inner_app) { ->(_env) { raise ::JWT::VerificationError } }
+
+      it 'returns a 401 and the correct error msg' do
+        expect(last_response.status).to eq 401
+        body = JSON.parse(last_response.body, symbolize_names: true)
+        expect(body).to eq(error: 'Invalid JWT token')
+      end
+    end
+
+    describe '::JWT::ExpiredSignature' do
+      let(:inner_app) { ->(_env) { raise ::JWT::ExpiredSignature } }
+
+      it 'returns a 401 and the correct error msg' do
+        expect(last_response.status).to eq 401
+        body = JSON.parse(last_response.body, symbolize_names: true)
+        expect(body).to eq(error: 'Invalid JWT token')
+      end
+    end
+
+    describe '::JWT::IncorrectAlgorithm' do
+      let(:inner_app) { ->(_env) { raise ::JWT::IncorrectAlgorithm } }
+
+      it 'returns a 401 and the correct error msg' do
+        expect(last_response.status).to eq 401
+        body = JSON.parse(last_response.body, symbolize_names: true)
+        expect(body).to eq(error: 'Invalid JWT token')
+      end
+    end
+
+    describe '::JWT::ImmatureSignature' do
+      let(:inner_app) { ->(_env) { raise ::JWT::ImmatureSignature } }
+
+      it 'returns a 401 and the correct error msg' do
+        expect(last_response.status).to eq 401
+        body = JSON.parse(last_response.body, symbolize_names: true)
+        expect(body).to eq(error: 'Invalid JWT token')
+      end
+    end
+
+    describe '::JWT::InvalidIssuerError' do
+      let(:inner_app) { ->(_env) { raise ::JWT::InvalidIssuerError } }
+
+      it 'returns a 401 and the correct error msg' do
+        expect(last_response.status).to eq 401
+        body = JSON.parse(last_response.body, symbolize_names: true)
+        expect(body).to eq(error: 'Invalid JWT token')
+      end
+    end
+
+    describe '::JWT::InvalidIatError' do
+      let(:inner_app) { ->(_env) { raise ::JWT::InvalidIatError } }
+
+      it 'returns a 401 and the correct error msg' do
+        expect(last_response.status).to eq 401
+        body = JSON.parse(last_response.body, symbolize_names: true)
+        expect(body).to eq(error: 'Invalid JWT token')
+      end
+    end
+
+    describe '::JWT::InvalidAudError' do
+      let(:inner_app) { ->(_env) { raise ::JWT::InvalidAudError } }
+
+      it 'returns a 401 and the correct error msg' do
+        expect(last_response.status).to eq 401
+        body = JSON.parse(last_response.body, symbolize_names: true)
+        expect(body).to eq(error: 'Invalid JWT token')
+      end
+    end
+
+    describe '::JWT::InvalidSubError' do
+      let(:inner_app) { ->(_env) { raise ::JWT::InvalidSubError } }
+
+      it 'returns a 401 and the correct error msg' do
+        expect(last_response.status).to eq 401
+        body = JSON.parse(last_response.body, symbolize_names: true)
+        expect(body).to eq(error: 'Invalid JWT token')
+      end
+    end
+
+    describe '::JWT::InvalidJtiError' do
+      let(:inner_app) { ->(_env) { raise ::JWT::InvalidJtiError } }
+
+      it 'returns a 401 and the correct error msg' do
+        expect(last_response.status).to eq 401
+        body = JSON.parse(last_response.body, symbolize_names: true)
+        expect(body).to eq(error: 'Invalid JWT token')
+      end
+    end
+
+    describe '::JWT::DecodeError' do
+      let(:inner_app) { ->(_env) { raise ::JWT::DecodeError } }
+
+      it 'returns a 401 and the correct error msg' do
+        expect(last_response.status).to eq 401
+        body = JSON.parse(last_response.body, symbolize_names: true)
+        expect(body).to eq(error: 'Invalid JWT token')
+      end
+    end
+
+    describe '::Rack::JWT::Auth::MissingAuthHeader' do
+      let(:inner_app) { ->(_env) { raise ::Rack::JWT::Auth::MissingAuthHeader } }
+
+      it 'returns a 401 and the correct error msg' do
+        expect(last_response.status).to eq 401
+        body = JSON.parse(last_response.body, symbolize_names: true)
+        expect(body).to eq(error: 'Missing Authorization header')
+      end
+    end
+
+    describe '::Rack::JWT::Auth::InvalidAuthHeaderFormat' do
+      let(:inner_app) { ->(_env) { raise ::Rack::JWT::Auth::InvalidAuthHeaderFormat } }
+
+      it 'returns a 401 and the correct error msg' do
+        expect(last_response.status).to eq 401
+        body = JSON.parse(last_response.body, symbolize_names: true)
+        expect(body).to eq(error: 'Invalid Authorization header format')
+      end
+    end
+  end
+end