diff --git a/arena.go b/arena.go index 1565fb1..201f075 100644 --- a/arena.go +++ b/arena.go @@ -19,13 +19,14 @@ package arenaskl import ( "errors" + "math" "sync/atomic" "unsafe" ) // Arena should be lock-free. type Arena struct { - n uint32 + n uint64 buf []byte } @@ -55,24 +56,37 @@ func NewArena(size uint32) *Arena { } func (a *Arena) Size() uint32 { - return atomic.LoadUint32(&a.n) + s := atomic.LoadUint64(&a.n) + if s > math.MaxUint32 { + // Saturate at MaxUint32. + return math.MaxUint32 + } + return uint32(s) } func (a *Arena) Reset() { - atomic.StoreUint32(&a.n, 1) + atomic.StoreUint64(&a.n, 1) } func (a *Arena) Alloc(size uint32, align Align) (uint32, error) { + // Verify that the arena isn't already full. + origSize := atomic.LoadUint64(&a.n) + if int(origSize) > len(a.buf) { + return 0, ErrArenaFull + } + // Pad the allocation with enough bytes to ensure the requested alignment. padded := uint32(size) + uint32(align) - newSize := atomic.AddUint32(&a.n, padded) + // Use 64-bit arithmetic to protect against overflow. + newSize := atomic.AddUint64(&a.n, uint64(padded)) if int(newSize) > len(a.buf) { + // Doubles as a check against newSize > math.MaxUint32. return 0, ErrArenaFull } // Return the aligned offset. - offset := (newSize - padded + uint32(align)) & ^uint32(align) + offset := (uint32(newSize) - padded + uint32(align)) & ^uint32(align) return offset, nil } diff --git a/arena_test.go b/arena_test.go new file mode 100644 index 0000000..2881cd6 --- /dev/null +++ b/arena_test.go @@ -0,0 +1,48 @@ +/* + * Copyright 2017 Dgraph Labs, Inc. and Contributors + * Modifications copyright (C) 2017 Andy Kimball and Contributors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package arenaskl + +import ( + "math" + "testing" + + "github.com/stretchr/testify/require" +) + +// TestArenaSizeOverflow tests that large allocations do not cause Arena's +// internal size accounting to overflow and produce incorrect results. +func TestArenaSizeOverflow(t *testing.T) { + a := NewArena(math.MaxUint32) + + // Allocating under the limit throws no error. + offset, err := a.Alloc(math.MaxUint16, Align1) + require.Nil(t, err) + require.Equal(t, uint32(1), offset) + require.Equal(t, uint32(math.MaxUint16)+1, a.Size()) + + // Allocating over the limit could cause an accounting + // overflow if 32-bit arithmetic was used. It shouldn't. + _, err = a.Alloc(math.MaxUint32, Align1) + require.Equal(t, ErrArenaFull, err) + require.Equal(t, uint32(math.MaxUint32), a.Size()) + + // Continuing to allocate continues to throw an error. + _, err = a.Alloc(math.MaxUint16, Align1) + require.Equal(t, ErrArenaFull, err) + require.Equal(t, uint32(math.MaxUint32), a.Size()) +}