Skip to content

Commit

Permalink
fix: add test & only defer when using RawBody
Browse files Browse the repository at this point in the history
  • Loading branch information
danielgtaylor committed Aug 15, 2024
1 parent 0534ce8 commit 6ca70c3
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 2 deletions.
16 changes: 14 additions & 2 deletions huma.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}
}
}
}
Expand Down
24 changes: 24 additions & 0 deletions huma_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down

0 comments on commit 6ca70c3

Please sign in to comment.