Skip to content

Commit

Permalink
feat: add max slice alloc size config (#376)
Browse files Browse the repository at this point in the history
  • Loading branch information
brianshih1 authored Apr 20, 2024
1 parent 0b21284 commit b43fe48
Show file tree
Hide file tree
Showing 3 changed files with 38 additions and 7 deletions.
10 changes: 4 additions & 6 deletions codec_array.go
Original file line number Diff line number Diff line change
Expand Up @@ -39,10 +39,6 @@ type arrayDecoder struct {
decoder ValDecoder
}

// Max allocation size for an array due to the limit in number of bits in a heap address:
// https://github.com/golang/go/blob/7f76c00fc5678fa782708ba8fece63750cb89d03/src/runtime/malloc.go#L183
var maxAllocSize = uint64(1 << 48)

func (d *arrayDecoder) Decode(ptr unsafe.Pointer, r *Reader) {
var size int
sliceType := d.typ
Expand All @@ -55,10 +51,12 @@ func (d *arrayDecoder) Decode(ptr unsafe.Pointer, r *Reader) {

start := size
size += int(l)
if size > int(maxAllocSize) {
r.ReportError("decode array", "size exceeded max allocation size")

if size > r.cfg.getMaxSliceAllocSize() {
r.ReportError("decode array", "size is greater than `Config.MaxSliceAllocSize`")
return
}

sliceType.UnsafeGrow(ptr, size)

for i := start; i < size; i++ {
Expand Down
20 changes: 19 additions & 1 deletion config.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,12 @@ import (
"github.com/modern-go/reflect2"
)

const defaultMaxByteSliceSize = 1_048_576 // 1 MiB
const (
defaultMaxByteSliceSize = 1_048_576 // 1 MiB
// Max allocation size for an array due to the limit in number of bits in a heap address:
// https://github.com/golang/go/blob/7f76c00fc5678fa782708ba8fece63750cb89d03/src/runtime/malloc.go#L183
maxAllocSize = int(1 << 48)
)

// DefaultConfig is the default API.
var DefaultConfig = Config{}.Freeze()
Expand Down Expand Up @@ -49,6 +54,11 @@ type Config struct {
// MaxByteSliceSize is the maximum size of `bytes` or `string` types the Reader will create, defaulting to 1MiB.
// If this size is exceeded, the Reader returns an error. This can be disabled by setting a negative number.
MaxByteSliceSize int

// MaxSliceAllocSize is the maximum size that the decoder will allocate, set to the max heap
// allocation size by default.
// If this size is exceeded, the decoder returns an error.
MaxSliceAllocSize int
}

// Freeze makes the configuration immutable.
Expand Down Expand Up @@ -266,3 +276,11 @@ func (c *frozenConfig) getMaxByteSliceSize() int {
}
return size
}

func (c *frozenConfig) getMaxSliceAllocSize() int {
size := c.config.MaxSliceAllocSize
if size > maxAllocSize || size <= 0 {
return maxAllocSize
}
return size
}
15 changes: 15 additions & 0 deletions decoder_array_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -91,3 +91,18 @@ func TestDecoder_ArrayMaxAllocationError(t *testing.T) {

assert.Error(t, err)
}

func TestDecoder_ArrayExceedMaxSliceAllocationConfig(t *testing.T) {
avro.DefaultConfig = avro.Config{MaxSliceAllocSize: 5}.Freeze()

// 10 (long) gets encoded to 0x14
data := []byte{0x14}
schema := `{"type":"array", "items": { "type": "boolean" }}`
dec, err := avro.NewDecoder(schema, bytes.NewReader(data))
require.NoError(t, err)

var got []bool
err = dec.Decode(&got)

assert.Error(t, err)
}

0 comments on commit b43fe48

Please sign in to comment.