From 22b407e2e7838c08939315b67ec982519e8f4679 Mon Sep 17 00:00:00 2001
From: Hakan Kutluay <77051856+hakankutluay@users.noreply.github.com>
Date: Sun, 9 Apr 2023 17:05:51 +0300
Subject: [PATCH] :bug: [Bug-Fix] add original timeout middleware (#2367)

* add original timeout middleware

* fix linter issues

* deprecate original timeout middleware

* update timeout middleware documentation
---
 docs/api/middleware/timeout.md     | 19 ++++++++++--
 middleware/timeout/timeout.go      | 47 ++++++++++++++++++++++++++++--
 middleware/timeout/timeout_test.go | 12 ++++----
 3 files changed, 67 insertions(+), 11 deletions(-)

diff --git a/docs/api/middleware/timeout.md b/docs/api/middleware/timeout.md
index aa0f2eb3c4..7f8adf9336 100644
--- a/docs/api/middleware/timeout.md
+++ b/docs/api/middleware/timeout.md
@@ -3,16 +3,29 @@ id: timeout
 title: Timeout
 ---
 
-Timeout middleware for [Fiber](https://github.com/gofiber/fiber). As a `fiber.Handler` wrapper, it creates a context with `context.WithTimeout` and pass it in `UserContext`. 
+There exist two distinct implementations of timeout middleware [Fiber](https://github.com/gofiber/fiber).
 
+**New**
+
+ Wraps a `fiber.Handler` with a timeout. If the handler takes longer than the given duration to return, the timeout error is set and forwarded to the centralized [ErrorHandler](https://docs.gofiber.io/error-handling).
+
+	Note: This has been depreciated since it raises race conditions.
+
+**NewWithContext**
+
+ As a `fiber.Handler` wrapper, it creates a context with `context.WithTimeout` and pass it in `UserContext`. 
+ 
 If the context passed executions (eg. DB ops, Http calls) takes longer than the given duration to return, the timeout error is set and forwarded to the centralized `ErrorHandler`.
 
+
 It does not cancel long running executions. Underlying executions must handle timeout by using `context.Context` parameter.
 
 ## Signatures
 
 ```go
 func New(handler fiber.Handler, timeout time.Duration, timeoutErrors ...error) fiber.Handler
+
+func NewWithContext(handler fiber.Handler, timeout time.Duration, timeoutErrors ...error) fiber.Handler
 ```
 
 ## Examples
@@ -85,7 +98,7 @@ func main() {
 		return nil
 	}
 
-	app.Get("/foo/:sleepTime", timeout.New(h, 2*time.Second, ErrFooTimeOut))
+	app.Get("/foo/:sleepTime", timeout.NewWithContext(h, 2*time.Second, ErrFooTimeOut))
 	_ = app.Listen(":3000")
 }
 
@@ -124,7 +137,7 @@ func main() {
 		return nil
 	}
 
-	app.Get("/foo", timeout.New(handler, 10*time.Second))
+	app.Get("/foo", timeout.NewWithContext(handler, 10*time.Second))
 	app.Listen(":3000")
 }
 ```
diff --git a/middleware/timeout/timeout.go b/middleware/timeout/timeout.go
index 9f1fd21997..1485704dff 100644
--- a/middleware/timeout/timeout.go
+++ b/middleware/timeout/timeout.go
@@ -3,13 +3,56 @@ package timeout
 import (
 	"context"
 	"errors"
+	"log"
+	"sync"
 	"time"
 
 	"github.com/gofiber/fiber/v2"
 )
 
-// New implementation of timeout middleware. Set custom errors(context.DeadlineExceeded vs) for get fiber.ErrRequestTimeout response.
-func New(h fiber.Handler, t time.Duration, tErrs ...error) fiber.Handler {
+var once sync.Once
+
+// New wraps a handler and aborts the process of the handler if the timeout is reached.
+//
+// Deprecated: This implementation contains data race issues. Use NewWithContext instead.
+// Find documentation and sample usage on https://docs.gofiber.io/api/middleware/timeout
+func New(handler fiber.Handler, timeout time.Duration) fiber.Handler {
+	once.Do(func() {
+		log.Printf("[Warning] timeout contains data race issues, not ready for production!")
+	})
+
+	if timeout <= 0 {
+		return handler
+	}
+
+	// logic is from fasthttp.TimeoutWithCodeHandler https://github.com/valyala/fasthttp/blob/master/server.go#L418
+	return func(ctx *fiber.Ctx) error {
+		ch := make(chan struct{}, 1)
+
+		go func() {
+			defer func() {
+				if err := recover(); err != nil {
+					log.Printf("[Warning] recover error %v", err)
+				}
+			}()
+			if err := handler(ctx); err != nil {
+				log.Printf("[Warning] handler error %v", err)
+			}
+			ch <- struct{}{}
+		}()
+
+		select {
+		case <-ch:
+		case <-time.After(timeout):
+			return fiber.ErrRequestTimeout
+		}
+
+		return nil
+	}
+}
+
+// NewWithContext implementation of timeout middleware. Set custom errors(context.DeadlineExceeded vs) for get fiber.ErrRequestTimeout response.
+func NewWithContext(h fiber.Handler, t time.Duration, tErrs ...error) fiber.Handler {
 	return func(ctx *fiber.Ctx) error {
 		timeoutContext, cancel := context.WithTimeout(ctx.UserContext(), t)
 		defer cancel()
diff --git a/middleware/timeout/timeout_test.go b/middleware/timeout/timeout_test.go
index 8498cb816d..413ac0da2d 100644
--- a/middleware/timeout/timeout_test.go
+++ b/middleware/timeout/timeout_test.go
@@ -12,12 +12,12 @@ import (
 	"github.com/gofiber/fiber/v2/utils"
 )
 
-// go test -run Test_Timeout
-func Test_Timeout(t *testing.T) {
+// go test -run Test_WithContextTimeout
+func Test_WithContextTimeout(t *testing.T) {
 	t.Parallel()
 	// fiber instance
 	app := fiber.New()
-	h := New(func(c *fiber.Ctx) error {
+	h := NewWithContext(func(c *fiber.Ctx) error {
 		sleepTime, err := time.ParseDuration(c.Params("sleepTime") + "ms")
 		utils.AssertEqual(t, nil, err)
 		if err := sleepWithContext(c.UserContext(), sleepTime, context.DeadlineExceeded); err != nil {
@@ -44,12 +44,12 @@ func Test_Timeout(t *testing.T) {
 
 var ErrFooTimeOut = errors.New("foo context canceled")
 
-// go test -run Test_TimeoutWithCustomError
-func Test_TimeoutWithCustomError(t *testing.T) {
+// go test -run Test_WithContextTimeoutWithCustomError
+func Test_WithContextTimeoutWithCustomError(t *testing.T) {
 	t.Parallel()
 	// fiber instance
 	app := fiber.New()
-	h := New(func(c *fiber.Ctx) error {
+	h := NewWithContext(func(c *fiber.Ctx) error {
 		sleepTime, err := time.ParseDuration(c.Params("sleepTime") + "ms")
 		utils.AssertEqual(t, nil, err)
 		if err := sleepWithContext(c.UserContext(), sleepTime, ErrFooTimeOut); err != nil {