test(middleware/csrf): fix Benchmark_Middleware_CSRF_*

gofiber · Mar 25, 2024 · 54eae2a · 54eae2a
1 parent 1607d87
commit 54eae2a
Showing 1 changed file with 19 additions and 6 deletions.
diff --git a/middleware/csrf/csrf_test.go b/middleware/csrf/csrf_test.go
@@ -992,7 +992,10 @@ func Benchmark_Middleware_CSRF_Check(b *testing.B) {
 		return c.SendStatus(fiber.StatusTeapot)
 	})
 
-	fctx := &fasthttp.RequestCtx{}
+	app.Post("/", func(c *fiber.Ctx) error {
+		return c.SendStatus(fiber.StatusTeapot)
+	})
+
 	h := app.Handler()
 	ctx := &fasthttp.RequestCtx{}
 
@@ -1002,17 +1005,27 @@ func Benchmark_Middleware_CSRF_Check(b *testing.B) {
 	token := string(ctx.Response.Header.Peek(fiber.HeaderSetCookie))
 	token = strings.Split(strings.Split(token, ";")[0], "=")[1]
 
+	// Test Correct Referer POST
+	ctx.Request.Reset()
+	ctx.Response.Reset()
 	ctx.Request.Header.SetMethod(fiber.MethodPost)
+	ctx.Request.Header.Set(fiber.HeaderXForwardedProto, "https")
+	ctx.Request.URI().SetScheme("https")
+	ctx.Request.URI().SetHost("example.com")
+	ctx.Request.Header.SetProtocol("https")
+	ctx.Request.Header.SetHost("example.com")
+	ctx.Request.Header.Set(fiber.HeaderReferer, "https://example.com")
 	ctx.Request.Header.Set(HeaderName, token)
+	ctx.Request.Header.SetCookie(ConfigDefault.CookieName, token)
 
 	b.ReportAllocs()
 	b.ResetTimer()
 
 	for n := 0; n < b.N; n++ {
-		h(fctx)
+		h(ctx)
 	}
 
-	utils.AssertEqual(b, fiber.StatusTeapot, fctx.Response.Header.StatusCode())
+	utils.AssertEqual(b, fiber.StatusTeapot, ctx.Response.Header.StatusCode())
 }
 
 // go test -v -run=^$ -bench=Benchmark_Middleware_CSRF_GenerateToken -benchmem -count=4
@@ -1024,7 +1037,6 @@ func Benchmark_Middleware_CSRF_GenerateToken(b *testing.B) {
 		return c.SendStatus(fiber.StatusTeapot)
 	})
 
-	fctx := &fasthttp.RequestCtx{}
 	h := app.Handler()
 	ctx := &fasthttp.RequestCtx{}
 
@@ -1034,8 +1046,9 @@ func Benchmark_Middleware_CSRF_GenerateToken(b *testing.B) {
 	b.ResetTimer()
 
 	for n := 0; n < b.N; n++ {
-		h(fctx)
+		h(ctx)
 	}
 
-	utils.AssertEqual(b, fiber.StatusTeapot, fctx.Response.Header.StatusCode())
+	// Ensure the GET request returns a 418 status code
+	utils.AssertEqual(b, fiber.StatusTeapot, ctx.Response.Header.StatusCode())
 }