Adding a MaximumRequestSizeHandler

spec/lucky/maximum_request_size_handler_spec.cr

@@ -0,0 +1,51 @@
+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) 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

spec/spec_helper.cr

@@ -48,4 +48,8 @@ Lucky::ForceSSLHandler.configure do |settings|
   settings.enabled = true
 end
 
+Lucky::MaximumRequestSizeHandler.configure do |settings|
+  settings.enabled = false
+end
+
 Habitat.raise_if_missing_settings!

src/lucky/maximum_request_size_handler.cr

@@ -0,0 +1,40 @@
+class Lucky::MaximumRequestSizeHandler
+  include HTTP::Handler
+
+  Habitat.create do
+    setting enabled : Bool = false
+    setting max_size : Int32 = 1_048_576 # 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