From d84d4342a437a3708bf29dffa85245904a65518e Mon Sep 17 00:00:00 2001
From: "Mr.Cyprinus" <38321534+Cyprinus12138@users.noreply.github.com>
Date: Sat, 6 Apr 2024 15:13:08 +0800
Subject: [PATCH] Fix #1 middleware eats req body issue  (#2)

* - Fix req body restoring
- Add unittests for calc the body size.

* - Update CHANGELOG.md
---
 CHANGELOG.md     |  5 ++++
 gintrace.go      |  3 ++-
 gintrace_test.go | 69 ++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 76 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b816d5c..e972f93 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,3 +16,8 @@ N/A
 - Implement basic http server metrics. 
 - Add example of initialization of `MeterProvider` in the example.
 - Add README.md.
+
+## [v1.0.1] - 2024-04-06
+
+### Fixed
+- Fix the bug when middleware calculating the size of the request body. (#1)
diff --git a/gintrace.go b/gintrace.go
index 56c75c7..e9300a6 100644
--- a/gintrace.go
+++ b/gintrace.go
@@ -6,6 +6,7 @@
 package otelgin // import "go.opentelemetry.io/contrib/instrumentation/github.com/gin-gonic/gin/otelgin"
 
 import (
+	"bytes"
 	"fmt"
 	otelmetric "go.opentelemetry.io/otel/metric"
 	"go.opentelemetry.io/otel/metric/noop"
@@ -192,7 +193,7 @@ func calcReqSize(c *gin.Context) int {
 	}
 
 	// Restore the request body for further processing
-	c.Request.Body = io.NopCloser(c.Request.Body)
+	c.Request.Body = io.NopCloser(bytes.NewReader(body))
 
 	// Calculate the size of headers
 	headerSize := 0
diff --git a/gintrace_test.go b/gintrace_test.go
index 8b02ac4..d6c1232 100644
--- a/gintrace_test.go
+++ b/gintrace_test.go
@@ -6,7 +6,9 @@
 package otelgin
 
 import (
+	"bytes"
 	"context"
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -95,3 +97,70 @@ func TestPropagationWithCustomPropagators(t *testing.T) {
 
 	router.ServeHTTP(w, r)
 }
+
+// TestCalcReqSize tests the calcReqSize function.
+func TestCalcReqSize(t *testing.T) {
+	// Create a sample request with a body and headers
+	body := []byte("sample body")
+	req, err := http.NewRequest("POST", "/test", bytes.NewReader(body))
+	if err != nil {
+		t.Fatalf("Failed to create request: %v", err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Authorization", "Bearer token")
+
+	// Create a Gin context with the request
+	w := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(w)
+	c.Request = req
+
+	// Call the function to calculate the request size
+	size := calcReqSize(c)
+
+	// Calculate the expected size (body + headers + extra bytes for header formatting)
+	expectedSize := len(body) + len("Content-Type") + len("application/json") + len("Authorization") + len("Bearer token") + 4 // 4 extra bytes for ": " and "\r\n"
+
+	// Check if the calculated size matches the expected size
+	if size != expectedSize {
+		t.Errorf("Expected request size %d, got %d", expectedSize, size)
+	}
+}
+
+// TestCalcReqSizeWithBodyRead tests the calcReqSize function and ensures the request body can still be read afterward.
+func TestCalcReqSizeWithBodyRead(t *testing.T) {
+	// Create a sample request with a body and headers
+	body := []byte("sample body")
+	req, err := http.NewRequest("POST", "/test", bytes.NewReader(body))
+	if err != nil {
+		t.Fatalf("Failed to create request: %v", err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Authorization", "Bearer token")
+
+	// Create a Gin context with the request
+	w := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(w)
+	c.Request = req
+
+	// Call the function to calculate the request size
+	size := calcReqSize(c)
+
+	// Calculate the expected size (body + headers + extra bytes for header formatting)
+	expectedSize := len(body) + len("Content-Type") + len("application/json") + len("Authorization") + len("Bearer token") + 4 // 4 extra bytes for ": " and "\r\n"
+
+	// Check if the calculated size matches the expected size
+	if size != expectedSize {
+		t.Errorf("Expected request size %d, got %d", expectedSize, size)
+	}
+
+	// Read the request body again
+	newBody, err := io.ReadAll(c.Request.Body)
+	if err != nil {
+		t.Fatalf("Failed to read request body: %v", err)
+	}
+
+	// Check if the body is unchanged
+	if !bytes.Equal(newBody, body) {
+		t.Errorf("Expected request body %q, got %q", body, newBody)
+	}
+}