diff --git a/spec/lucky/maximum_request_size_handler_spec.cr b/spec/lucky/maximum_request_size_handler_spec.cr new file mode 100644 index 000000000..17f81f8d4 --- /dev/null +++ b/spec/lucky/maximum_request_size_handler_spec.cr @@ -0,0 +1,54 @@ +require "../spec_helper" +require "http/server" + +include ContextHelper + +describe Lucky::MaximumRequestSizeHandler do + context "when the handler is disabled" do + it "simply serves the request" do + context = build_small_request_context("/path") + Lucky::MaximumRequestSizeHandler.temp_config(enabled: false) do + run_request_size_handler(context) + end + context.response.status.should eq(HTTP::Status::OK) + end + end + + context "when the handler is enabled" do + it "with a small request, serve the request" do + context = build_small_request_context("/path") + Lucky::MaximumRequestSizeHandler.temp_config(enabled: true) do + run_request_size_handler(context) + end + context.response.status.should eq(HTTP::Status::OK) + end + + it "with a large request, deny the request" do + context = build_large_request_context("/path") + Lucky::MaximumRequestSizeHandler.temp_config( + enabled: true, + max_size: 10, + ) do + run_request_size_handler(context) + end + context.response.status.should eq(HTTP::Status::PAYLOAD_TOO_LARGE) + end + end +end + +private def run_request_size_handler(context) + handler = Lucky::MaximumRequestSizeHandler.new + handler.next = ->(_ctx : HTTP::Server::Context) {} + handler.call(context) +end + +private def build_small_request_context(path : String) : HTTP::Server::Context + build_context(path: path) +end + +private def build_large_request_context(path : String) : HTTP::Server::Context + build_context(path: path).tap do |context| + context.request.headers["Content-Length"] = "1000000" + context.request.body = "a" * 1000000 + end +end diff --git a/src/lucky/maximum_request_size_handler.cr b/src/lucky/maximum_request_size_handler.cr new file mode 100644 index 000000000..6a14a17f1 --- /dev/null +++ b/src/lucky/maximum_request_size_handler.cr @@ -0,0 +1,52 @@ +# Allows a maximum request size to be set for incoming requests. +# +# Configure the max_size to the maximum size in bytes that you +# want to allow. +# +# ``` +# Lucky::MaximumRequestSizeHandler.configure do |settings| +# settings.enabled = true +# settings.max_size = 1_048_576 # 1MB +# end +# ``` + +class Lucky::MaximumRequestSizeHandler + include HTTP::Handler + + Habitat.create do + setting enabled : Bool = false + setting max_size : Int64 = 1_048_576_i64 # 1MB + end + + def call(context) + return call_next(context) unless settings.enabled + + body_size = 0 + body = IO::Memory.new + + begin + buffer = Bytes.new(8192) # 8KB buffer + while read_bytes = context.request.body.try(&.read(buffer)) + body_size += read_bytes + body.write(buffer[0, read_bytes]) + + if body_size > settings.max_size + context.response.status = HTTP::Status::PAYLOAD_TOO_LARGE + context.response.print("Request entity too large") + return context + end + + break if read_bytes < buffer.size # End of body + end + rescue IO::Error + context.response.status = HTTP::Status::BAD_REQUEST + context.response.print("Error reading request body") + return context + end + + # Reset the request body for downstream handlers + context.request.body = IO::Memory.new(body.to_s) + + call_next(context) + end +end