From 6ca70c34fdfa46c52f7db99649d2e5f9fbbf051a Mon Sep 17 00:00:00 2001
From: "Daniel G. Taylor" <danielgtaylor@gmail.com>
Date: Thu, 15 Aug 2024 10:38:05 -0700
Subject: [PATCH] fix: add test & only defer when using RawBody

---
 huma.go      | 16 ++++++++++++++--
 huma_test.go | 24 ++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/huma.go b/huma.go
index 38ea470a..d0566e6b 100644
--- a/huma.go
+++ b/huma.go
@@ -1292,8 +1292,20 @@ func Register[I, O any](api API, op Operation, handler func(context.Context, *I)
 						}
 					}
 
-					defer bufPool.Put(buf)
-					defer buf.Reset()
+					if rawBodyIndex != -1 {
+						// If the raw body is used, then we must wait until *AFTER* the
+						// handler has run to return the body byte buffer to the pool, as
+						// the handler can read and modify this buffer. The safest way is
+						// to just wait until the end of this handler via defer.
+						defer bufPool.Put(buf)
+						defer buf.Reset()
+					} else {
+						// No raw body, and the body has already been unmarshalled above, so
+						// we can return the buffer to the pool now as we don't need the
+						// bytes any more.
+						bufPool.Put(buf)
+						buf.Reset()
+					}
 				}
 			}
 		}
diff --git a/huma_test.go b/huma_test.go
index 9e1f5704..d4e5ade4 100644
--- a/huma_test.go
+++ b/huma_test.go
@@ -2160,6 +2160,30 @@ func TestSchemaWithExample(t *testing.T) {
 	assert.Equal(t, 1, example)
 }
 
+func TestBodyRace(t *testing.T) {
+	// Run with the following:
+	// go test -run=TestBodyRace -race -parallel=100
+	_, api := humatest.New(t, huma.DefaultConfig("Test API", "1.0.0"))
+	huma.Post(api, "/ping", func(ctx context.Context, input *struct {
+		Body struct {
+			Value string `json:"value"`
+		}
+		RawBody []byte
+	}) (*struct{}, error) {
+		// Access/modify the raw input to detect races.
+		input.RawBody[1] = 'a'
+		return nil, nil
+	})
+
+	for i := 0; i < 100; i++ {
+		t.Run(fmt.Sprintf("test-%d", i), func(tt *testing.T) {
+			tt.Parallel()
+			resp := api.Post("/ping", map[string]any{"value": "hello"})
+			assert.Equal(tt, 204, resp.Result().StatusCode)
+		})
+	}
+}
+
 // func BenchmarkSecondDecode(b *testing.B) {
 // 	//nolint: musttag
 // 	type MediumSized struct {