Adding a MaximumRequestSizeHandler

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

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