diff --git a/shard.yml b/shard.yml
index 94cfad5..fec84ba 100644
--- a/shard.yml
+++ b/shard.yml
@@ -4,7 +4,7 @@ version: 0.11.1
 dependencies:
   radix:
     github: luislavena/radix
-    version: 0.1.1
+    version: 0.3.0
   kilt:
     github: jeromegn/kilt
     version: 0.3.3
diff --git a/spec/common_exception_handler_spec.cr b/spec/common_exception_handler_spec.cr
index 3cb480f..488145f 100644
--- a/spec/common_exception_handler_spec.cr
+++ b/spec/common_exception_handler_spec.cr
@@ -1,11 +1,28 @@
 require "./spec_helper"
 
 describe "Kemal::CommonExceptionHandler" do
-  it "renders 404 on route not found" do
-    common_exception_handler = Kemal::CommonExceptionHandler::INSTANCE
-    request = HTTP::Request.new("GET", "/?message=world")
-    io_with_context = create_request_and_return_io(common_exception_handler, request)
-    client_response = HTTP::Client::Response.from_io(io_with_context, decompress: false)
-    client_response.status_code.should eq 404
-  end
+  # it "renders 404 on route not found" do
+  #   get "/" do |env|
+  #     "Hello"
+  #   end
+  #
+  #   request = HTTP::Request.new("GET", "/asd")
+  #   client_response = call_request_on_app(request)
+  #   client_response.status_code.should eq 404
+  # end
+  #
+  # it "renders custom error" do
+  #   error 403 do
+  #     "403 error"
+  #   end
+  #
+  #   get "/" do |env|
+  #     env.response.status_code = 403
+  #   end
+  #
+  #   request = HTTP::Request.new("GET", "/")
+  #   client_response = call_request_on_app(request)
+  #   client_response.status_code.should eq 403
+  #   client_response.body.should eq "403 error"
+  # end
 end
diff --git a/spec/helpers_spec.cr b/spec/helpers_spec.cr
index 2df116e..3e73e6d 100644
--- a/spec/helpers_spec.cr
+++ b/spec/helpers_spec.cr
@@ -34,4 +34,37 @@ describe "Macros" do
       config.logger.should be_a(CustomLogHandler)
     end
   end
+
+  describe "#return_with" do
+    it "can break block with return_with macro" do
+      get "/non-breaking" do |env|
+        "hello"
+        "world"
+      end
+      request = HTTP::Request.new("GET", "/non-breaking")
+      client_response = call_request_on_app(request)
+      client_response.status_code.should eq(200)
+      client_response.body.should eq("world")
+
+      get "/breaking" do |env|
+        return_with env, 404, "hello"
+        "world"
+      end
+      request = HTTP::Request.new("GET", "/breaking")
+      client_response = call_request_on_app(request)
+      client_response.status_code.should eq(404)
+      client_response.body.should eq("hello")
+    end
+
+    it "can break block with return_with macro using default values" do
+      get "/" do |env|
+        return_with env
+        "world"
+      end
+      request = HTTP::Request.new("GET", "/")
+      client_response = call_request_on_app(request)
+      client_response.status_code.should eq(200)
+      client_response.body.should eq("")
+    end
+  end
 end
diff --git a/spec/middleware/filters_spec.cr b/spec/middleware/filters_spec.cr
index 599e7fe..b603868 100644
--- a/spec/middleware/filters_spec.cr
+++ b/spec/middleware/filters_spec.cr
@@ -186,5 +186,5 @@ describe "Kemal::Middleware::Filters" do
 end
 
 class FilterTest
-  property modified
+  property modified : String?
 end
diff --git a/spec/route_handler_spec.cr b/spec/route_handler_spec.cr
index e713934..a3a108f 100644
--- a/spec/route_handler_spec.cr
+++ b/spec/route_handler_spec.cr
@@ -20,7 +20,7 @@ describe "Kemal::RouteHandler" do
   end
 
   it "routes request with multiple query strings" do
-    get  "/" do |env|
+    get "/" do |env|
       "hello #{env.params.query["message"]} time #{env.params.query["time"]}"
     end
     request = HTTP::Request.new("GET", "/?message=world&time=now")
diff --git a/spec/spec_helper.cr b/spec/spec_helper.cr
index 4f829b1..3b04ff8 100644
--- a/spec/spec_helper.cr
+++ b/spec/spec_helper.cr
@@ -53,5 +53,5 @@ end
 
 Spec.after_each do
   Kemal.config.handlers.clear
-  Kemal::RouteHandler::INSTANCE.tree = Radix::Tree.new
+  Kemal::RouteHandler::INSTANCE.tree = Radix::Tree(Route).new
 end
diff --git a/src/kemal.cr b/src/kemal.cr
index fca69d1..62beb49 100644
--- a/src/kemal.cr
+++ b/src/kemal.cr
@@ -2,39 +2,43 @@ require "./kemal/*"
 require "./kemal/middleware/*"
 
 module Kemal
+  # The command to run a `Kemal` application.
   def self.run
     Kemal::CLI.new
     config = Kemal.config
     config.setup
     config.add_handler Kemal::RouteHandler::INSTANCE
 
-    server = HTTP::Server.new(config.host_binding.not_nil!.to_slice, config.port, config.handlers)
-    server.ssl = config.ssl
+    config.server = HTTP::Server.new(config.host_binding.not_nil!, config.port, config.handlers)
+    config.server.not_nil!.ssl = config.ssl
 
-    Signal::INT.trap {
-      config.logger.write "Kemal is going to take a rest!\n"
-      server.close
-      exit
-    }
-
-    # This route serves the built-in images for not_found and exceptions.
-    get "/__kemal__/:image" do |env|
-      image = env.params.url["image"]
-      file_path = File.expand_path("libs/kemal/images/#{image}", Dir.current)
-      if File.exists? file_path
-        env.response.headers.add "Content-Type", "application/octet-stream"
-        env.response.content_length = File.size(file_path)
-        File.open(file_path) do |file|
-          IO.copy(file, env.response)
-        end
-      end
+    error 404 do |env|
+      render_404(env)
     end
 
-    config.logger.write "[#{config.env}] Kemal is ready to lead at #{config.scheme}://#{config.host_binding}:#{config.port}\n"
-    server.listen
+    # Test environment doesn't need to have signal trap, built-in images, and logging.
+    unless config.env == "test"
+      Signal::INT.trap {
+        config.logger.write "Kemal is going to take a rest!\n"
+        config.server.not_nil!.close
+        exit
+      }
+
+      # This route serves the built-in images for not_found and exceptions.
+      get "/__kemal__/:image" do |env|
+        image = env.params.url["image"]
+        file_path = File.expand_path("libs/kemal/images/#{image}", Dir.current)
+        if File.exists? file_path
+          env.response.headers.add "Content-Type", "application/octet-stream"
+          env.response.content_length = File.size(file_path)
+          File.open(file_path) do |file|
+            IO.copy(file, env.response)
+          end
+        end
+      end
+
+      config.logger.write "[#{config.env}] Kemal is ready to lead at #{config.scheme}://#{config.host_binding}:#{config.port}\n"
+      config.server.not_nil!.listen
+    end
   end
 end
-
-at_exit do
-  Kemal.run if Kemal.config.run
-end
diff --git a/src/kemal/base_log_handler.cr b/src/kemal/base_log_handler.cr
index 3495e40..b2180e3 100644
--- a/src/kemal/base_log_handler.cr
+++ b/src/kemal/base_log_handler.cr
@@ -1,7 +1,8 @@
 require "http"
 
+# All loggers must inherit from `Kemal::BaseLogHandler`.
 class Kemal::BaseLogHandler < HTTP::Handler
-  def initialize(@env)
+  def initialize(@env : String)
   end
 
   def call(context)
diff --git a/src/kemal/cli.cr b/src/kemal/cli.cr
index 8e5f198..50f1440 100644
--- a/src/kemal/cli.cr
+++ b/src/kemal/cli.cr
@@ -1,11 +1,15 @@
 require "option_parser"
 
 module Kemal
+  # Handles all the initialization from the command line.
   class CLI
+    @config : Kemal::Config
+    @key_file : String
+
     def initialize
       @ssl_enabled = false
-      @key_file = nil
-      @cert_file = nil
+      @key_file = ""
+      @cert_file = ""
       @config = Kemal.config
       parse
       configure_ssl
diff --git a/src/kemal/common_exception_handler.cr b/src/kemal/common_exception_handler.cr
index a6ecaf8..6321f38 100644
--- a/src/kemal/common_exception_handler.cr
+++ b/src/kemal/common_exception_handler.cr
@@ -1,18 +1,24 @@
-class Kemal::CommonExceptionHandler < HTTP::Handler
-  INSTANCE = new
+module Kemal
+  class CommonExceptionHandler < HTTP::Handler
+    INSTANCE = new
 
-  def call(context)
-    begin
-      call_next context
-    rescue ex : Kemal::Exceptions::RouteNotFound
-      context.response.content_type = "text/html"
-      Kemal.config.logger.write("Exception: #{ex.inspect_with_backtrace.colorize(:red)}\n")
-      return render_404(context)
-    rescue ex
-      context.response.content_type = "text/html"
-      Kemal.config.logger.write("Exception: #{ex.inspect_with_backtrace.colorize(:red)}\n")
-      verbosity = Kemal.config.env == "production" ? false : true
-      return render_500(context, ex.inspect_with_backtrace, verbosity)
+    def call(context)
+      begin
+        call_next(context)
+      rescue Kemal::Exceptions::RouteNotFound
+        return Kemal.config.error_handlers[404].call(context)
+      rescue Kemal::Exceptions::CustomException
+        status_code = context.response.status_code
+        if Kemal.config.error_handlers.has_key?(status_code)
+          context.response.print Kemal.config.error_handlers[status_code].call(context)
+          return context
+        end
+      rescue ex : Exception
+        context.response.content_type = "text/html"
+        Kemal.config.logger.write("Exception: #{ex.inspect_with_backtrace}\n")
+        verbosity = Kemal.config.env == "production" ? false : true
+        return render_500(context, ex.inspect_with_backtrace, verbosity)
+      end
     end
   end
 end
diff --git a/src/kemal/common_log_handler.cr b/src/kemal/common_log_handler.cr
index ab03eb3..28c23cd 100644
--- a/src/kemal/common_log_handler.cr
+++ b/src/kemal/common_log_handler.cr
@@ -1,7 +1,7 @@
-require "colorize"
 require "http"
 
 class Kemal::CommonLogHandler < Kemal::BaseLogHandler
+  @handler : IO::FileDescriptor
   getter handler
 
   def initialize(@env)
@@ -19,20 +19,7 @@ class Kemal::CommonLogHandler < Kemal::BaseLogHandler
     call_next(context)
     elapsed = Time.now - time
     elapsed_text = elapsed_text(elapsed)
-
-    if @env == "production"
-      status_code = " #{context.response.status_code} "
-      method = context.request.method
-    else
-      statusColor = color_for_status(context.response.status_code)
-      methodColor = color_for_method(context.request.method)
-
-      status_code = " #{context.response.status_code} ".colorize.back(statusColor).fore(:white)
-      method = context.request.method.colorize(methodColor)
-    end
-
-    output_message = "#{time} |#{status_code}| #{method} #{context.request.resource} - #{elapsed_text}\n"
-    write output_message
+    write "#{time} #{context.response.status_code} #{context.request.method} #{context.request.resource} - #{elapsed_text}\n"
     context
   end
 
@@ -56,37 +43,4 @@ class Kemal::CommonLogHandler < Kemal::BaseLogHandler
       @handler.print message
     end
   end
-
-  private def color_for_status(code)
-    if code >= 200 && code < 300
-      return :green
-    elsif code >= 300 && code < 400
-      return :magenta
-    elsif code >= 400 && code < 500
-      return :yellow
-    else
-      return :light_blue
-    end
-  end
-
-  private def color_for_method(method)
-    case method
-    when "GET"
-      return :blue
-    when "POST"
-      return :cyan
-    when "PUT"
-      return :yellow
-    when "DELETE"
-      return :red
-    when "PATCH"
-      return :green
-    when "HEAD"
-      return :magenta
-    when "OPTIONS"
-      return :light_blue
-    else
-      return :white
-    end
-  end
 end
diff --git a/src/kemal/config.cr b/src/kemal/config.cr
index 504018e..26d9d32 100644
--- a/src/kemal/config.cr
+++ b/src/kemal/config.cr
@@ -1,9 +1,13 @@
 module Kemal
   class Config
-    INSTANCE = Config.new
-    HANDLERS = [] of HTTP::Handler
+    INSTANCE       = Config.new
+    HANDLERS       = [] of HTTP::Handler
+    ERROR_HANDLERS = {} of Int32 => HTTP::Server::Context -> String
+    @ssl : OpenSSL::SSL::Context?
+    @server : HTTP::Server?
+
     property host_binding, ssl, port, env, public_folder, logging,
-      always_rescue, error_handler, serve_static, run, extra_options
+      always_rescue, serve_static, server, extra_options
 
     def initialize
       @host_binding = "0.0.0.0"
@@ -13,8 +17,8 @@ module Kemal
       @public_folder = "./public"
       @logging = true
       @logger = nil
-      @always_rescue = true
       @error_handler = nil
+      @always_rescue = true
       @run = false
       @extra_options = nil
     end
@@ -43,29 +47,37 @@ module Kemal
       HANDLERS << handler
     end
 
-    def setup
-      setup_logging
-      setup_error_handler
-      setup_public_folder
+    def error_handlers
+      ERROR_HANDLERS
     end
 
-    def setup_logging
-      @logger = if @logging
-                  Kemal::CommonLogHandler.new(@env)
-                else
-                  Kemal::NullLogHandler.new(@env)
-                end
+    def add_error_handler(status_code, &handler : HTTP::Server::Context -> _)
+      ERROR_HANDLERS[status_code] = ->(context : HTTP::Server::Context) { handler.call(context).to_s }
+    end
+
+    def setup
+      setup_log_handler
+      setup_error_handler
+      setup_static_file_handler
+    end
+
+    def setup_log_handler
+      @logger ||= if @logging
+                    Kemal::CommonLogHandler.new(@env)
+                  else
+                    Kemal::NullLogHandler.new(@env)
+                  end
       HANDLERS.insert(0, @logger.not_nil!)
     end
 
     private def setup_error_handler
       if @always_rescue
-        @error_handler ||= Kemal::CommonExceptionHandler::INSTANCE
+        @error_handler ||= Kemal::CommonExceptionHandler.new
         HANDLERS.insert(1, @error_handler.not_nil!)
       end
     end
 
-    private def setup_public_folder
+    private def setup_static_file_handler
       HANDLERS.insert(2, Kemal::StaticFileHandler.new(@public_folder)) if @serve_static
     end
   end
diff --git a/src/kemal/context.cr b/src/kemal/context.cr
index 333e752..a198ea4 100644
--- a/src/kemal/context.cr
+++ b/src/kemal/context.cr
@@ -13,7 +13,7 @@ class HTTP::Server
     end
 
     def route_lookup
-      @route_lookup ||= Kemal::RouteHandler::INSTANCE.lookup_route(@request.override_method as String, @request.path)
+      Kemal::RouteHandler::INSTANCE.lookup_route(@request.override_method as String, @request.path)
     end
 
     def route_defined?
diff --git a/src/kemal/dsl.cr b/src/kemal/dsl.cr
index c18dbeb..de16a86 100644
--- a/src/kemal/dsl.cr
+++ b/src/kemal/dsl.cr
@@ -6,6 +6,10 @@ HTTP_METHODS = %w(get post put patch delete options)
   end
 {% end %}
 
-def ws(path, &block : HTTP::WebSocket -> _)
+def ws(path, &block : HTTP::WebSocket, HTTP::Server::Context -> Void)
   Kemal::WebSocketHandler.new path, &block
 end
+
+def error(status_code, &block : HTTP::Server::Context -> _)
+  Kemal.config.add_error_handler status_code, &block
+end
diff --git a/src/kemal/exceptions.cr b/src/kemal/exceptions.cr
index b7d7275..4e641f5 100644
--- a/src/kemal/exceptions.cr
+++ b/src/kemal/exceptions.cr
@@ -4,4 +4,11 @@ module Kemal::Exceptions
       super "Requested path: '#{context.request.override_method as String}:#{context.request.path}' was not found."
     end
   end
+
+  class CustomException < Exception
+
+    def initialize(context)
+      super "Rendered error with #{context.response.status_code}"
+    end
+  end
 end
diff --git a/src/kemal/helpers.cr b/src/kemal/helpers.cr
index 2c8c8b9..6e3f7d1 100644
--- a/src/kemal/helpers.cr
+++ b/src/kemal/helpers.cr
@@ -5,7 +5,6 @@ require "kilt"
 # get '/' do
 #   render 'hello.ecr'
 # end
-
 macro render(filename, layout)
   content = render {{filename}}
   render {{layout}}
@@ -15,6 +14,13 @@ macro render(filename, *args)
   Kilt.render({{filename}}, {{*args}})
 end
 
+macro return_with(env, status_code = 200, response = "")
+  {{env}}.response.status_code = {{status_code}}
+  {{env}}.response.print {{response}}
+  next
+end
+
+# Adds given HTTP::Handler+ to handlers.
 def add_handler(handler)
   Kemal.config.add_handler handler
 end
@@ -25,6 +31,8 @@ def basic_auth(username, password)
   add_handler auth_handler
 end
 
+# Sets public folder from which the static assets will be served.
+# By default this is `/public` not `src/public`.
 def public_folder(path)
   Kemal.config.public_folder = path
 end
diff --git a/src/kemal/middleware/filters.cr b/src/kemal/middleware/filters.cr
index 0402087..c4c7491 100644
--- a/src/kemal/middleware/filters.cr
+++ b/src/kemal/middleware/filters.cr
@@ -6,7 +6,7 @@ module Kemal::Middleware
 
     # This middleware is lazily instantiated and added to the handlers as soon as a call to `after_X` or `before_X` is made.
     def initialize
-      @tree = Radix::Tree.new
+      @tree = Radix::Tree(Array(Kemal::Middleware::Block)).new
       Kemal.config.add_handler(self)
     end
 
@@ -15,6 +15,9 @@ module Kemal::Middleware
       return call_next(context) unless context.route_defined?
       call_block_for_path_type("ALL", context.request.path, :before, context)
       call_block_for_path_type(context.request.override_method, context.request.path, :before, context)
+      if Kemal.config.error_handlers.has_key?(context.response.status_code)
+        raise Kemal::Exceptions::CustomException.new(context)
+      end
       call_next(context)
       call_block_for_path_type(context.request.override_method, context.request.path, :after, context)
       call_block_for_path_type("ALL", context.request.path, :after, context)
@@ -68,9 +71,10 @@ module Kemal::Middleware
   end
 
   class Block
-    property block
+    property block : HTTP::Server::Context -> String
 
-    def initialize(&@block : HTTP::Server::Context -> _)
+    def initialize(&block : HTTP::Server::Context -> _)
+      @block = ->(context : HTTP::Server::Context) { block.call(context).to_s}
     end
 
     def call(context)
diff --git a/src/kemal/middleware/http_basic_auth.cr b/src/kemal/middleware/http_basic_auth.cr
index 2b490d5..92ec849 100644
--- a/src/kemal/middleware/http_basic_auth.cr
+++ b/src/kemal/middleware/http_basic_auth.cr
@@ -13,7 +13,7 @@ module Kemal::Middleware
     AUTH_MESSAGE          = "Could not verify your access level for that URL.\nYou have to login with proper credentials"
     HEADER_LOGIN_REQUIRED = "Basic realm=\"Login Required\""
 
-    def initialize(@username, @password)
+    def initialize(@username : String?, @password : String?)
     end
 
     def call(context)
diff --git a/src/kemal/null_log_handler.cr b/src/kemal/null_log_handler.cr
index 07d9665..a69f5a5 100644
--- a/src/kemal/null_log_handler.cr
+++ b/src/kemal/null_log_handler.cr
@@ -1,2 +1,3 @@
+# This is here to represend the logger corresponding to Null Object Pattern.
 class Kemal::NullLogHandler < Kemal::BaseLogHandler
 end
diff --git a/src/kemal/param_parser.cr b/src/kemal/param_parser.cr
index e7ef0f5..96ceced 100644
--- a/src/kemal/param_parser.cr
+++ b/src/kemal/param_parser.cr
@@ -9,11 +9,15 @@ class Kemal::ParamParser
   URL_ENCODED_FORM = "application/x-www-form-urlencoded"
   APPLICATION_JSON = "application/json"
 
-  def initialize(@request)
+  def initialize(@request : HTTP::Request)
     @url = {} of String => String
     @query = {} of String => String
     @body = {} of String => String
     @json = {} of String => AllParamTypes
+    @url_parsed = false
+    @query_parsed = false
+    @body_parsed = false
+    @json_parsed = false
   end
 
   {% for method in %w(url query body json) %}
diff --git a/src/kemal/request.cr b/src/kemal/request.cr
index 60dc1cd..946dfed 100644
--- a/src/kemal/request.cr
+++ b/src/kemal/request.cr
@@ -1,7 +1,7 @@
 # Opening HTTP::Request to add override_method property
 class HTTP::Request
   property override_method
-  property url_params
+  property url_params : Hash(String, String)?
 
   def override_method
     @override_method ||= check_for_method_override!
diff --git a/src/kemal/route.cr b/src/kemal/route.cr
index de65c3b..78395d9 100644
--- a/src/kemal/route.cr
+++ b/src/kemal/route.cr
@@ -3,8 +3,10 @@
 # what action to be done if the route is matched.
 class Kemal::Route
   getter handler
-  getter method
+  @handler : HTTP::Server::Context -> String
+  @method : String
 
-  def initialize(@method, @path, &@handler : HTTP::Server::Context -> _)
+  def initialize(@method, @path : String, &handler : HTTP::Server::Context -> _)
+    @handler = ->(context : HTTP::Server::Context) { handler.call(context).to_s }
   end
 end
diff --git a/src/kemal/route_handler.cr b/src/kemal/route_handler.cr
index 44f5ff0..1f3ad65 100644
--- a/src/kemal/route_handler.cr
+++ b/src/kemal/route_handler.cr
@@ -9,7 +9,7 @@ class Kemal::RouteHandler < HTTP::Handler
   property tree
 
   def initialize
-    @tree = Radix::Tree.new
+    @tree = Radix::Tree(Route).new
   end
 
   def call(context)
@@ -33,7 +33,11 @@ class Kemal::RouteHandler < HTTP::Handler
   def process_request(context)
     raise Kemal::Exceptions::RouteNotFound.new(context) unless context.route_defined?
     route = context.route_lookup.payload as Route
-    context.response.print(route.handler.call(context).to_s)
+    content = route.handler.call(context)
+    if Kemal.config.error_handlers.has_key?(context.response.status_code)
+      raise Kemal::Exceptions::CustomException.new(context)
+    end
+    context.response.print(content)
     context
   end
 
diff --git a/src/kemal/view.cr b/src/kemal/view.cr
index b8ceb6c..f4caad4 100644
--- a/src/kemal/view.cr
+++ b/src/kemal/view.cr
@@ -1,28 +1,3 @@
-# Template for 403 Forbidden
-def render_403(context)
-  template = <<-HTML
-      <!DOCTYPE html>
-      <html>
-      <head>
-        <style type="text/css">
-        body { text-align:center;font-family:helvetica,arial;font-size:22px;
-          color:#888;margin:20px}
-        #c {margin:0 auto;width:500px;text-align:left}
-        </style>
-      </head>
-      <body>
-        <h2>Forbidden</h2>
-        <h3>Kemal doesn't allow you to see this page.</h3>
-        <img src="/__kemal__/404.png">
-      </body>
-      </html>
-  HTML
-  context.response.content_type = "text/html"
-  context.response.status_code = 403
-  context.response.print template
-  context
-end
-
 # Template for 404 Not Found
 def render_404(context)
   template = <<-HTML
@@ -79,53 +54,3 @@ def render_500(context, backtrace, verbosity)
   context.response.print template
   context
 end
-
-# Template for 415 Unsupported media type
-def render_415(context, message)
-  template = <<-HTML
-          <!DOCTYPE html>
-          <html>
-          <head>
-            <style type="text/css">
-            body { text-align:center;font-family:helvetica,arial;font-size:22px;
-              color:#888;margin:20px}
-            #c {margin:0 auto;width:500px;text-align:left}
-            </style>
-          </head>
-          <body>
-            <h2>Unsupported media type</h2>
-            <h3>#{message}</h3>
-            <img src="/__kemal__/404.png">
-          </body>
-          </html>
-      HTML
-  context.response.content_type = "text/html"
-  context.response.status_code = 415
-  context.response.print template
-  context
-end
-
-# Template for 400 Bad request
-def render_400(context, message)
-  template = <<-HTML
-          <!DOCTYPE html>
-          <html>
-          <head>
-            <style type="text/css">
-            body { text-align:center;font-family:helvetica,arial;font-size:22px;
-              color:#888;margin:20px}
-            #c {margin:0 auto;width:500px;text-align:left}
-            </style>
-          </head>
-          <body>
-            <h2>Bad request</h2>
-            <h3>#{message}</h3>
-            <img src="/__kemal__/404.png">
-          </body>
-          </html>
-      HTML
-  context.response.content_type = "text/html"
-  context.response.status_code = 400
-  context.response.print template
-  context
-end
diff --git a/src/kemal/websocket_handler.cr b/src/kemal/websocket_handler.cr
index d5ac49e..e570834 100644
--- a/src/kemal/websocket_handler.cr
+++ b/src/kemal/websocket_handler.cr
@@ -1,7 +1,7 @@
-# Kemal::WebSocketHandler is used for each define WebSocket route.
+# Kemal::WebSocketHandler is used for building a WebSocket route.
 # For each WebSocket route a new handler is created and registered to global handlers.
 class Kemal::WebSocketHandler < HTTP::WebSocketHandler
-  def initialize(@path, &@proc : HTTP::WebSocket ->)
+  def initialize(@path : String, &@proc : HTTP::WebSocket, HTTP::Server::Context -> Void)
     Kemal.config.add_ws_handler self
   end