Protect against accounting overflow in Arena allocation

This change protects against large allocations causing Arena's
internal size accounting to overflow and produce incorrect results.
This overflow could allow for invalid offsets being returned from
`Arena.Alloc`, leading to slice bounds and index out of range panics
when passed to `Arena.GetBytes` or `Arena.GetPointer`.

Specifically, if `Arena.n` overflowed in `Arena.Alloc`, which was
possible because it was a uint32 and we allow up to MaxUint32 size
allocations, then `Arena.Alloc` could bypass the length check against
`Arena.buf`. The "padded" size would then be subtracted from the offset,
allowing the offset to underflow back around to an offset that was out
of `Arena.buf`'s bounds.

I believe this is the underlying cause of the following two crashes:
- cockroachdb/cockroach#31624
- cockroachdb/cockroach#35557

The reason for this is subtle and has to do with failed allocations
slowly building up and eventually overflowing `Arena.n`. I'll explain
on the PR that vendors this fix.

We now protect against this overflow in two ways. We perform an initial
size check before performing the atomic addition in `Arena.Alloc` to
prevent failed allocations from "building up" to an overflow. We also
use 64-bit arithmetic so that it would take 2^32 concurrent allocations
with the maximum allocation size (math.MaxUint32) to overflow the
accounting. An alternative would be to use a CAS loop to strictly
serialize all additions to `Arena.n`, but that would limit concurrency.

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
 }
 

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())
+}