From 58677d5c8644eae4474ba5d5d0f283cbf0a9ae21 Mon Sep 17 00:00:00 2001
From: ACHMAD IRIANTO EKA PUTRA <i@ryanbekhen.dev>
Date: Mon, 23 Dec 2024 14:20:41 +0700
Subject: [PATCH] feat: Add Drop method to DefaultCtx for silent connection
 termination (#3257)

* Add Drop method to DefaultCtx and remove redundant checks

Introduced a Drop method in DefaultCtx for closing connections, enabling easier resource management. Removed unnecessary nil-checks for headers in manager_msgp to simplify code logic. Added a unit test to ensure the new Drop method behaves as expected.

* Add `Drop` method to Fiber context API documentation

The `Drop` method allows silently terminating client connections without sending HTTP headers or a response body. This is useful for scenarios like mitigating DDoS attacks or blocking unauthorized access to sensitive endpoints. Example usage and function signature are included in the updated documentation.

* Remove extraneous blank line in documentation.

Eliminated an unnecessary blank line in the API context documentation for improved readability and formatting consistency. No functional changes were made to the content.

* Update API documentation example to return "Hello World!"

Revised the example code in the API documentation to return a generic "Hello World!" string instead of a dynamic response. This improves consistency and simplifies the example for easier understanding.

* Refactor Drop method and extend test coverage.

Simplified the Drop method by inlining the connection close call. Added new test cases to ensure proper handling of no-response scenarios and improved overall test coverage.

* fix golangci-lint issue

* Add test for Ctx.Drop with middleware interaction

This test ensures the correct behavior of the Ctx.Drop method when used with middleware, including response handling and error scenarios. It verifies that the middleware and handler properly handle the Drop call and its resulting effects.

* Add Drop method to DefaultCtx for closing connections

The Drop method allows closing connections without sending a response, improving control over connection handling. Also updated a test assertion to use StatusOK for improved readability and consistency.

* Refine Drop method comments to clarify error handling.

Explain the rationale for not wrapping errors in the Drop method. Emphasize that the returned error is solely for logging and not for further propagation or processing.

* Update Drop method documentation for clarity

Clarified the `Drop` method's behavior, specifying that it closes the connection without sending headers or a body. Added examples of use cases, such as DDoS mitigation and blocking sensitive endpoints.

* Refactor response header setting in middleware.

Replaced the direct header setting with the `Set` method for consistency and improved clarity. Removed a test case checking for a panic on closed response body as it is no longer applicable.
---
 ctx.go                           |  8 +++++
 ctx_interface_gen.go             |  1 +
 ctx_test.go                      | 53 ++++++++++++++++++++++++++++++++
 docs/api/ctx.md                  | 21 +++++++++++++
 middleware/cache/manager_msgp.go |  6 ----
 5 files changed, 83 insertions(+), 6 deletions(-)

diff --git a/ctx.go b/ctx.go
index 5dde2f89..789a5315 100644
--- a/ctx.go
+++ b/ctx.go
@@ -1978,3 +1978,11 @@ func (c *DefaultCtx) setMatched(matched bool) {
 func (c *DefaultCtx) setRoute(route *Route) {
 	c.route = route
 }
+
+// Drop closes the underlying connection without sending any response headers or body.
+// This can be useful for silently terminating client connections, such as in DDoS mitigation
+// or when blocking access to sensitive endpoints.
+func (c *DefaultCtx) Drop() error {
+	//nolint:wrapcheck // error wrapping is avoided to keep the operation lightweight and focused on connection closure.
+	return c.RequestCtx().Conn().Close()
+}
diff --git a/ctx_interface_gen.go b/ctx_interface_gen.go
index 859563f6..d7f8bbc6 100644
--- a/ctx_interface_gen.go
+++ b/ctx_interface_gen.go
@@ -350,4 +350,5 @@ type Ctx interface {
 	setIndexRoute(route int)
 	setMatched(matched bool)
 	setRoute(route *Route)
+	Drop() error
 }
diff --git a/ctx_test.go b/ctx_test.go
index 72c415a9..d025c244 100644
--- a/ctx_test.go
+++ b/ctx_test.go
@@ -5847,6 +5847,59 @@ func Test_GenericParseTypeBoolean(t *testing.T) {
 	}
 }
 
+// go test -run Test_Ctx_Drop -v
+func Test_Ctx_Drop(t *testing.T) {
+	t.Parallel()
+
+	app := New()
+
+	// Handler that calls Drop
+	app.Get("/block-me", func(c Ctx) error {
+		return c.Drop()
+	})
+
+	// Additional handler that just calls return
+	app.Get("/no-response", func(_ Ctx) error {
+		return nil
+	})
+
+	// Test the Drop method
+	resp, err := app.Test(httptest.NewRequest(MethodGet, "/block-me", nil))
+	require.Error(t, err)
+	require.Nil(t, resp)
+
+	// Test the no-response handler
+	resp, err = app.Test(httptest.NewRequest(MethodGet, "/no-response", nil))
+	require.NoError(t, err)
+	require.NotNil(t, resp)
+	require.Equal(t, StatusOK, resp.StatusCode)
+	require.Equal(t, "0", resp.Header.Get("Content-Length"))
+}
+
+// go test -run Test_Ctx_DropWithMiddleware -v
+func Test_Ctx_DropWithMiddleware(t *testing.T) {
+	t.Parallel()
+
+	app := New()
+
+	// Middleware that calls Drop
+	app.Use(func(c Ctx) error {
+		err := c.Next()
+		c.Set("X-Test", "test")
+		return err
+	})
+
+	// Handler that calls Drop
+	app.Get("/block-me", func(c Ctx) error {
+		return c.Drop()
+	})
+
+	// Test the Drop method
+	resp, err := app.Test(httptest.NewRequest(MethodGet, "/block-me", nil))
+	require.Error(t, err)
+	require.Nil(t, resp)
+}
+
 // go test -run Test_GenericParseTypeString
 func Test_GenericParseTypeString(t *testing.T) {
 	t.Parallel()
diff --git a/docs/api/ctx.md b/docs/api/ctx.md
index 6f4b1ae4..85532a54 100644
--- a/docs/api/ctx.md
+++ b/docs/api/ctx.md
@@ -463,6 +463,27 @@ app.Get("/", func(c fiber.Ctx) error {
 })
 ```
 
+## Drop
+
+Terminates the client connection silently without sending any HTTP headers or response body.
+
+This can be used for scenarios where you want to block certain requests without notifying the client, such as mitigating
+DDoS attacks or protecting sensitive endpoints from unauthorized access.
+
+```go title="Signature"
+func (c fiber.Ctx) Drop() error
+```
+
+```go title="Example"
+app.Get("/", func(c fiber.Ctx) error {
+  if c.IP() == "192.168.1.1" {
+    return c.Drop()
+  }
+
+  return c.SendString("Hello World!")
+})
+```
+
 ## Format
 
 Performs content-negotiation on the [Accept](https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Accept) HTTP header. It uses [Accepts](ctx.md#accepts) to select a proper format from the supplied offers. A default handler can be provided by setting the `MediaType` to `"default"`. If no offers match and no default is provided, a 406 (Not Acceptable) response is sent. The Content-Type is automatically set when a handler is selected.
diff --git a/middleware/cache/manager_msgp.go b/middleware/cache/manager_msgp.go
index 492e9a88..bf5d6152 100644
--- a/middleware/cache/manager_msgp.go
+++ b/middleware/cache/manager_msgp.go
@@ -52,9 +52,6 @@ func (z *item) DecodeMsg(dc *msgp.Reader) (err error) {
 					err = msgp.WrapError(err, "headers", za0001)
 					return
 				}
-				if za0002 == nil {
-					za0002 = make([]byte, 0)
-				}
 				z.headers[za0001] = za0002
 			}
 		case "body":
@@ -270,9 +267,6 @@ func (z *item) UnmarshalMsg(bts []byte) (o []byte, err error) {
 					err = msgp.WrapError(err, "headers", za0001)
 					return
 				}
-				if za0002 == nil {
-					za0002 = make([]byte, 0)
-				}
 				z.headers[za0001] = za0002
 			}
 		case "body":