Merge pull request #271 from onflow/fxamacker/port-hash-collision-limit

Add MaxCollisionLimitPerDigest

errors.go

@@ -434,3 +434,33 @@ func (e UnreachableError) Error() string {
 func NewUnreachableError() *UnreachableError {
  return &UnreachableError{Stack: debug.Stack()}
 }
+
+// CollisionLimitError is a fatal error returned when a noncryptographic hash collision
+// would exceed collision limit (per digest per map) we enforce in the first level.
+type CollisionLimitError struct {
+ collisionLimitPerDigest uint32 // limit <= 255 is recommended, larger values are useful for tests
+}
+
+// NewCollisionLimitError constructs a CollisionLimitError
+func NewCollisionLimitError(collisionLimitPerDigest uint32) error {
+ return NewFatalError(&CollisionLimitError{collisionLimitPerDigest: collisionLimitPerDigest})
+}
+
+func (e *CollisionLimitError) Error() string {
+ return fmt.Sprintf("collision limit per digest %d already reached", e.collisionLimitPerDigest)
+}
+
+// MapElementCountError is a fatal error returned when element count is unexpected.
+// It is an implemtation error.
+type MapElementCountError struct {
+ msg string
+}
+
+// NewMapElementCountError constructs a MapElementCountError.
+func NewMapElementCountError(msg string) error {
+ return NewFatalError(&MapElementCountError{msg: msg})
+}
+
+func (e *MapElementCountError) Error() string {
+ return e.msg
+}

map.go

@@ -75,6 +75,13 @@ const (
  typicalRandomConstant = uint64(0x1BD11BDAA9FC1A22) // DO NOT MODIFY
 )
 
+// MaxCollisionLimitPerDigest is the noncryptographic hash collision limit
+// (per digest per map) we enforce in the first level. In the same map
+// for the same digest, having a non-intentional collision should be rare and
+// several collisions should be extremely rare. The default limit should
+// be high enough to ignore accidental collisions while mitigating attacks.
+var MaxCollisionLimitPerDigest = uint32(255)
+
 type MapKey Storable
 
 type MapValue Storable
@@ -123,6 +130,8 @@ type element interface {
 
  Size() uint32
 
+ Count(storage SlabStorage) (uint32, error)
+
  PopIterate(SlabStorage, MapPopIterationFunc) error
 }
 
@@ -635,6 +644,10 @@ func (e *singleElement) Size() uint32 {
  return e.size
 }
 
+func (e *singleElement) Count(_ SlabStorage) (uint32, error) {
+ return 1, nil
+}
+
 func (e *singleElement) PopIterate(_ SlabStorage, fn MapPopIterationFunc) error {
  fn(e.key, e.value)
  return nil
@@ -787,6 +800,10 @@ func (e *inlineCollisionGroup) Elements(_ SlabStorage) (elements, error) {
  return e.elements, nil
 }
 
+func (e *inlineCollisionGroup) Count(_ SlabStorage) (uint32, error) {
+ return e.elements.Count(), nil
+}
+
 func (e *inlineCollisionGroup) PopIterate(storage SlabStorage, fn MapPopIterationFunc) error {
  return e.elements.PopIterate(storage, fn)
 }
@@ -946,6 +963,14 @@ func (e *externalCollisionGroup) Elements(storage SlabStorage) (elements, error)
  return dataSlab.elements, nil
 }
 
+func (e *externalCollisionGroup) Count(storage SlabStorage) (uint32, error) {
+ elements, err := e.Elements(storage)
+ if err != nil {
+ return 0, err
+ }
+ return elements.Count(), nil
+}
+
 func (e *externalCollisionGroup) PopIterate(storage SlabStorage, fn MapPopIterationFunc) error {
  elements, err := e.Elements(storage)
  if err != nil {
@@ -1248,11 +1273,48 @@ func (e *hkeyElements) Set(storage SlabStorage, address Address, b DigesterBuild
  }
  }
 
- // Has matching hkey
+ // hkey digest has collision.
  if equalIndex != -1 {
-
+ // New element has the same digest as existing elem.
+ // elem is existing element before new element is inserted.
  elem := e.elems[equalIndex]
 
+ // Enforce MaxCollisionLimitPerDigest at the first level (noncryptographic hash).
+ if e.level == 0 {
+
+ // Before new element with colliding digest is inserted,
+ // existing elem is a single element or a collision group.
+ // elem.Count() returns 1 for single element,
+ // and returns > 1 for collision group.
+ elementCount, err := elem.Count(storage)
+ if err != nil {
+ return nil, err
+ }
+ if elementCount == 0 {
+ return nil, NewMapElementCountError("expect element count > 0, got element count == 0")
+ }
+
+ // collisionCount is elementCount-1 because:
+ // - if elem is single element, collision count is 0 (no collsion yet)
+ // - if elem is collision group, collision count is 1 less than number
+ // of elements in collision group.
+ collisionCount := elementCount - 1
+
+ // Check if existing collision count reached MaxCollisionLimitPerDigest
+ if collisionCount >= MaxCollisionLimitPerDigest {
+ // Enforce collision limit on inserts and ignore updates.
+ _, err = elem.Get(storage, digester, level, hkey, comparator, key)
+ if err != nil {
+ var knfe *KeyNotFoundError
+ if errors.As(err, &knfe) {
+ // Don't allow any more collisions for a digest that
+ // already reached MaxCollisionLimitPerDigest.
+ return nil, NewCollisionLimitError(MaxCollisionLimitPerDigest)
+ }
+ }
+ }
+ }
+
  oldElemSize := elem.Size()
 
  elem, existingValue, err := elem.Set(storage, address, b, digester, level, hkey, comparator, hip, key, value)

map_test.go

@@ -21,6 +21,7 @@ package atree
 import (
  "errors"
  "fmt"
+ "math"
  "math/rand"
  "reflect"
  "sort"
@@ -368,14 +369,20 @@ func TestMapSetAndGet(t *testing.T) {
 
  t.Run("unique keys with hash collision", func(t *testing.T) {
 
- SetThreshold(256)
- defer SetThreshold(1024)
-
  const (
  mapSize = 1024
  keyStringSize = 16
  )
 
+ SetThreshold(256)
+ defer SetThreshold(1024)
+
+ savedMaxCollisionLimitPerDigest := MaxCollisionLimitPerDigest
+ MaxCollisionLimitPerDigest = uint32(math.Ceil(float64(mapSize) / 10))
+ defer func() {
+ MaxCollisionLimitPerDigest = savedMaxCollisionLimitPerDigest
+ }()
+
  r := newRand(t)
 
  digesterBuilder := &mockDigesterBuilder{}
@@ -410,14 +417,20 @@ func TestMapSetAndGet(t *testing.T) {
  })
 
  t.Run("replicate keys with hash collision", func(t *testing.T) {
- SetThreshold(256)
- defer SetThreshold(1024)
-
  const (
  mapSize = 1024
  keyStringSize = 16
  )
 
+ SetThreshold(256)
+ defer SetThreshold(1024)
+
+ savedMaxCollisionLimitPerDigest := MaxCollisionLimitPerDigest
+ MaxCollisionLimitPerDigest = uint32(math.Ceil(float64(mapSize) / 10))
+ defer func() {
+ MaxCollisionLimitPerDigest = savedMaxCollisionLimitPerDigest
+ }()
+
  r := newRand(t)
 
  digesterBuilder := &mockDigesterBuilder{}
@@ -430,7 +443,7 @@ func TestMapSetAndGet(t *testing.T) {
  i++
 
  digests := []Digest{
- Digest(1 % 10),
+ Digest(i % 10),
  }
  digesterBuilder.On("Digest", k).Return(mockDigester{digests})
  }
@@ -3391,10 +3404,16 @@ func TestMapFromBatchData(t *testing.T) {
 
  t.Run("collision", func(t *testing.T) {
 
+ const mapSize = 1024
+
  SetThreshold(512)
  defer SetThreshold(1024)
 
- const mapSize = 1024
+ savedMaxCollisionLimitPerDigest := MaxCollisionLimitPerDigest
+ defer func() {
+ MaxCollisionLimitPerDigest = savedMaxCollisionLimitPerDigest
+ }()
+ MaxCollisionLimitPerDigest = mapSize / 2
 
  typeInfo := testTypeInfo{42}
 
@@ -3863,3 +3882,150 @@ func TestMapSlabDump(t *testing.T) {
  require.Equal(t, want, dumps)
  })
 }
+
+func TestMaxCollisionLimitPerDigest(t *testing.T) {
+ savedMaxCollisionLimitPerDigest := MaxCollisionLimitPerDigest
+ defer func() {
+ MaxCollisionLimitPerDigest = savedMaxCollisionLimitPerDigest
+ }()
+
+ t.Run("collision limit 0", func(t *testing.T) {
+ const mapSize = 1024
+
+ SetThreshold(256)
+ defer SetThreshold(1024)
+
+ // Set noncryptographic hash collision limit as 0,
+ // meaning no collision is allowed at first level.
+ MaxCollisionLimitPerDigest = uint32(0)
+
+ digesterBuilder := &mockDigesterBuilder{}
+ keyValues := make(map[Value]Value, mapSize)
+ for i := uint64(0); i < mapSize; i++ {
+ k := Uint64Value(i)
+ v := Uint64Value(i)
+ keyValues[k] = v
+
+ digests := []Digest{Digest(i)}
+ digesterBuilder.On("Digest", k).Return(mockDigester{digests})
+ }
+
+ typeInfo := testTypeInfo{42}
+ address := Address{1, 2, 3, 4, 5, 6, 7, 8}
+ storage := newTestPersistentStorage(t)
+
+ m, err := NewMap(storage, address, digesterBuilder, typeInfo)
+ require.NoError(t, err)
+
+ // Insert elements within collision limits
+ for k, v := range keyValues {
+ existingStorable, err := m.Set(compare, hashInputProvider, k, v)
+ require.NoError(t, err)
+ require.Nil(t, existingStorable)
+ }
+
+ verifyMap(t, storage, typeInfo, address, m, keyValues, nil, false)
+
+ // Insert elements exceeding collision limits
+ collisionKeyValues := make(map[Value]Value, mapSize)
+ for i := uint64(0); i < mapSize; i++ {
+ k := Uint64Value(mapSize + i)
+ v := Uint64Value(mapSize + i)
+ collisionKeyValues[k] = v
+
+ digests := []Digest{Digest(i)}
+ digesterBuilder.On("Digest", k).Return(mockDigester{digests})
+ }
+
+ for k, v := range collisionKeyValues {
+ existingStorable, err := m.Set(compare, hashInputProvider, k, v)
+ var collisionLimitError *CollisionLimitError
+ require.ErrorAs(t, err, &collisionLimitError)
+ require.Nil(t, existingStorable)
+ }
+
+ // Verify that no new elements exceeding collision limit inserted
+ verifyMap(t, storage, typeInfo, address, m, keyValues, nil, false)
+
+ // Update elements within collision limits
+ for k := range keyValues {
+ v := Uint64Value(0)
+ keyValues[k] = v
+ existingStorable, err := m.Set(compare, hashInputProvider, k, v)
+ require.NoError(t, err)
+ require.NotNil(t, existingStorable)
+ }
+
+ verifyMap(t, storage, typeInfo, address, m, keyValues, nil, false)
+ })
+
+ t.Run("collision limit > 0", func(t *testing.T) {
+ const mapSize = 1024
+
+ SetThreshold(256)
+ defer SetThreshold(1024)
+
+ // Set noncryptographic hash collision limit as 7,
+ // meaning at most 8 elements in collision group per digest at first level.
+ MaxCollisionLimitPerDigest = uint32(7)
+
+ digesterBuilder := &mockDigesterBuilder{}
+ keyValues := make(map[Value]Value, mapSize)
+ for i := uint64(0); i < mapSize; i++ {
+ k := Uint64Value(i)
+ v := Uint64Value(i)
+ keyValues[k] = v
+
+ digests := []Digest{Digest(i % 128)}
+ digesterBuilder.On("Digest", k).Return(mockDigester{digests})
+ }
+
+ typeInfo := testTypeInfo{42}
+ address := Address{1, 2, 3, 4, 5, 6, 7, 8}
+ storage := newTestPersistentStorage(t)
+
+ m, err := NewMap(storage, address, digesterBuilder, typeInfo)
+ require.NoError(t, err)
+
+ // Insert elements within collision limits
+ for k, v := range keyValues {
+ existingStorable, err := m.Set(compare, hashInputProvider, k, v)
+ require.NoError(t, err)
+ require.Nil(t, existingStorable)
+ }
+
+ verifyMap(t, storage, typeInfo, address, m, keyValues, nil, false)
+
+ // Insert elements exceeding collision limits
+ collisionKeyValues := make(map[Value]Value, mapSize)
+ for i := uint64(0); i < mapSize; i++ {
+ k := Uint64Value(mapSize + i)
+ v := Uint64Value(mapSize + i)
+ collisionKeyValues[k] = v
+
+ digests := []Digest{Digest(i % 128)}
+ digesterBuilder.On("Digest", k).Return(mockDigester{digests})
+ }
+
+ for k, v := range collisionKeyValues {
+ existingStorable, err := m.Set(compare, hashInputProvider, k, v)
+ var collisionLimitError *CollisionLimitError
+ require.ErrorAs(t, err, &collisionLimitError)
+ require.Nil(t, existingStorable)
+ }
+
+ // Verify that no new elements exceeding collision limit inserted
+ verifyMap(t, storage, typeInfo, address, m, keyValues, nil, false)
+
+ // Update elements within collision limits
+ for k := range keyValues {
+ v := Uint64Value(0)
+ keyValues[k] = v
+ existingStorable, err := m.Set(compare, hashInputProvider, k, v)
+ require.NoError(t, err)
+ require.NotNil(t, existingStorable)
+ }
+
+ verifyMap(t, storage, typeInfo, address, m, keyValues, nil, false)
+ })
+}