Add ability to add handlers for raised exceptions

spec/exception_handler_spec.cr

@@ -59,6 +59,99 @@ describe "Kemal::ExceptionHandler" do
     response.body.should eq "Something happened"
   end
 
+  it "renders custom error for a crystal exception" do
+    error RuntimeError do
+      "A RuntimeError has occured"
+    end
+
+    get "/" do
+      raise RuntimeError.new
+    end
+
+    request = HTTP::Request.new("GET", "/")
+    io = IO::Memory.new
+    response = HTTP::Server::Response.new(io)
+    context = HTTP::Server::Context.new(request, response)
+    Kemal::ExceptionHandler::INSTANCE.next = Kemal::RouteHandler::INSTANCE
+    Kemal::ExceptionHandler::INSTANCE.call(context)
+    response.close
+    io.rewind
+    response = HTTP::Client::Response.from_io(io, decompress: false)
+    response.status_code.should eq 500
+    response.headers["Content-Type"].should eq "text/html"
+    response.body.should eq "A RuntimeError has occured"
+  end
+
+  it "renders custom error for a custom exception" do
+    error CustomExceptionType do
+      "A custom exception of CustomExceptionType has occurred"
+    end
+
+    get "/" do
+      raise CustomExceptionType.new
+    end
+
+    request = HTTP::Request.new("GET", "/")
+    io = IO::Memory.new
+    response = HTTP::Server::Response.new(io)
+    context = HTTP::Server::Context.new(request, response)
+    Kemal::ExceptionHandler::INSTANCE.next = Kemal::RouteHandler::INSTANCE
+    Kemal::ExceptionHandler::INSTANCE.call(context)
+    response.close
+    io.rewind
+    response = HTTP::Client::Response.from_io(io, decompress: false)
+    response.status_code.should eq 500
+    response.headers["Content-Type"].should eq "text/html"
+    response.body.should eq "A custom exception of CustomExceptionType has occurred"
+  end
+
+  it "renders custom error for a custom exception with a specific HTTP status code" do
+    error CustomExceptionType do |env|
+      env.response.status_code = 503
+      "A custom exception of CustomExceptionType has occurred"
+    end
+
+    get "/" do
+      raise CustomExceptionType.new
+    end
+
+    request = HTTP::Request.new("GET", "/")
+    io = IO::Memory.new
+    response = HTTP::Server::Response.new(io)
+    context = HTTP::Server::Context.new(request, response)
+    Kemal::ExceptionHandler::INSTANCE.next = Kemal::RouteHandler::INSTANCE
+    Kemal::ExceptionHandler::INSTANCE.call(context)
+    response.close
+    io.rewind
+    response = HTTP::Client::Response.from_io(io, decompress: false)
+    response.status_code.should eq 503
+    response.headers["Content-Type"].should eq "text/html"
+    response.body.should eq "A custom exception of CustomExceptionType has occurred"
+  end
+
+  it "renders custom error for a child of a custom exception" do
+    error CustomExceptionType do |env, error|
+      "A custom exception of #{error.class} has occurred"
+    end
+
+    get "/" do
+      raise ChildCustomExceptionType.new
+    end
+
+    request = HTTP::Request.new("GET", "/")
+    io = IO::Memory.new
+    response = HTTP::Server::Response.new(io)
+    context = HTTP::Server::Context.new(request, response)
+    Kemal::ExceptionHandler::INSTANCE.next = Kemal::RouteHandler::INSTANCE
+    Kemal::ExceptionHandler::INSTANCE.call(context)
+    response.close
+    io.rewind
+    response = HTTP::Client::Response.from_io(io, decompress: false)
+    response.status_code.should eq 500
+    response.headers["Content-Type"].should eq "text/html"
+    response.body.should eq "A custom exception of ChildCustomExceptionType has occurred"
+  end
+
   it "overrides the content type for filters" do
     before_get do |env|
       env.response.content_type = "application/json"

spec/spec_helper.cr

@@ -26,6 +26,12 @@ class AnotherContextStorageType
   @name = "kemal-context"
 end
 
+class CustomExceptionType < Exception
+end
+
+class ChildCustomExceptionType < CustomExceptionType
+end
+
 add_context_storage_type(TestContextStorageType)
 add_context_storage_type(AnotherContextStorageType)
 

src/kemal/config.cr

@@ -8,11 +8,12 @@ module Kemal
   # Kemal.config
   # ```
   class Config
-    INSTANCE        = Config.new
-    HANDLERS        = [] of HTTP::Handler
-    CUSTOM_HANDLERS = [] of Tuple(Nil | Int32, HTTP::Handler)
-    FILTER_HANDLERS = [] of HTTP::Handler
-    ERROR_HANDLERS  = {} of Int32 => HTTP::Server::Context, Exception -> String
+    INSTANCE           = Config.new
+    HANDLERS           = [] of HTTP::Handler
+    CUSTOM_HANDLERS    = [] of Tuple(Nil | Int32, HTTP::Handler)
+    FILTER_HANDLERS    = [] of HTTP::Handler
+    ERROR_HANDLERS     = {} of Int32 => HTTP::Server::Context, Exception -> String
+    EXCEPTION_HANDLERS = {} of Exception.class => HTTP::Server::Context, Exception -> String
 
     {% if flag?(:without_openssl) %}
       @ssl : Bool?
@@ -88,14 +89,26 @@ module Kemal
       FILTER_HANDLERS << handler
     end
 
+    # Returns the defined error handlers for HTTP status codes
     def error_handlers
       ERROR_HANDLERS
     end
 
+    # Adds an error handler for the given HTTP status code
     def add_error_handler(status_code : Int32, &handler : HTTP::Server::Context, Exception -> _)
       ERROR_HANDLERS[status_code] = ->(context : HTTP::Server::Context, error : Exception) { handler.call(context, error).to_s }
     end
 
+    # Returns tje defined error handlers for exceptions
+    def exception_handlers
+      EXCEPTION_HANDLERS
+    end
+
+    # Adds an error handler for the given exception
+    def add_exception_handler(exception : Exception.class, &handler : HTTP::Server::Context, Exception -> _)
+      EXCEPTION_HANDLERS[exception] = ->(context : HTTP::Server::Context, error : Exception) { handler.call(context, error).to_s }
+    end
+
     def extra_options(&@extra_options : OptionParser ->)
     end
 

src/kemal/dsl.cr

@@ -21,10 +21,16 @@ def ws(path : String, &block : HTTP::WebSocket, HTTP::Server::Context -> Void)
   Kemal::WebSocketHandler::INSTANCE.add_route path, &block
 end
 
+# Defines an error handler to be called when route returns the given HTTP status code
 def error(status_code : Int32, &block : HTTP::Server::Context, Exception -> _)
   Kemal.config.add_error_handler status_code, &block
 end
 
+# Defines an error handler to be called when the given exception is raised
+def error(exception : Exception.class, &block : HTTP::Server::Context, Exception -> _)
+  Kemal.config.add_exception_handler exception, &block
+end
+
 # All the helper methods available are:
 #  - before_all, before_get, before_post, before_put, before_patch, before_delete, before_options
 #  - after_all, after_get, after_post, after_put, after_patch, after_delete, after_options

src/kemal/exception_handler.cr

@@ -11,12 +11,41 @@ module Kemal
     rescue ex : Kemal::Exceptions::CustomException
       call_exception_with_status_code(context, ex, context.response.status_code)
     rescue ex : Exception
+      # Matches an error handler for the given exception
+      #
+      # Matches based on order of declaration rather than inheritance relationship
+      # for child exceptions
+      Kemal.config.exception_handlers.each do |expected_exception, handler|
+        if ex.class <= expected_exception
+          return call_exception_with_exception(context, ex, handler, 500)
+        end
+      end
+
       log("Exception: #{ex.inspect_with_backtrace}")
+      # Else use generic 500 handler if defined
       return call_exception_with_status_code(context, ex, 500) if Kemal.config.error_handlers.has_key?(500)
       verbosity = Kemal.config.env == "production" ? false : true
       render_500(context, ex, verbosity)
     end
 
+    # Calls the given error handler with the current exception
+    #
+    # The logic for validating that the current exception should be handled
+    # by the given error handler should be done by the caller of this method.
+    private def call_exception_with_exception(
+      context : HTTP::Server::Context,
+      exception : Exception,
+      handler : Proc(HTTP::Server::Context, Exception, String),
+      status_code : Int32 = 500,
+    )
+      return if context.response.closed?
+
+      context.response.content_type = "text/html" unless context.response.headers.has_key?("Content-Type")
+      context.response.status_code = status_code
+      context.response.print handler.call(context, exception)
+      context
+    end
+
     private def call_exception_with_status_code(context : HTTP::Server::Context, exception : Exception, status_code : Int32)
       return if context.response.closed?
       if !Kemal.config.error_handlers.empty? && Kemal.config.error_handlers.has_key?(status_code)