From 6809685620d5c0352206cf1b59a46aae32757cb8 Mon Sep 17 00:00:00 2001
From: Russ Smith <russ@bashme.org>
Date: Sun, 6 Oct 2024 09:23:19 -0700
Subject: [PATCH] Adding a MaximumRequestSizeHandler

* Configurable to be on/off. Off by default.
* Configurable to set the maximum request size. Default is 1MB.

Ref #1143
---
 .../maximum_request_size_handler_spec.cr      | 51 +++++++++++++++++++
 spec/spec_helper.cr                           |  4 ++
 src/lucky/maximum_request_size_handler.cr     | 40 +++++++++++++++
 3 files changed, 95 insertions(+)
 create mode 100644 spec/lucky/maximum_request_size_handler_spec.cr
 create mode 100644 src/lucky/maximum_request_size_handler.cr

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..a551cd210
--- /dev/null
+++ b/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
diff --git a/spec/spec_helper.cr b/spec/spec_helper.cr
index df67856c0..1ac943359 100644
--- a/spec/spec_helper.cr
+++ b/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!
diff --git a/src/lucky/maximum_request_size_handler.cr b/src/lucky/maximum_request_size_handler.cr
new file mode 100644
index 000000000..7c0aaff3f
--- /dev/null
+++ b/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