Merge pull request #542 from nunoo/nunoo/fix

fix: race by deferring the return of buf to sync.Pool when using RawBody

huma.go

@@ -1292,8 +1292,20 @@ func Register[I, O any](api API, op Operation, handler func(context.Context, *I)
 						}
 					}
 
-					buf.Reset()
-					bufPool.Put(buf)
+					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()
+					}
 				}
 			}
 		}

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 {