Skip to content

Commit 69aa5b8

Browse files
ARR4NARR4N
committed
(WIP) feat: state-key transformatin w/ override
1 parent 8ae4631 commit 69aa5b8

File tree

6 files changed

+103
-8
lines changed

6 files changed

+103
-8
lines changed

‎core/state/statedb.go‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -341,18 +341,20 @@ func (s *StateDB) GetCodeHash(addr common.Address) common.Hash {
341341
}
342342

343343
// GetState retrieves a value from the given account's storage trie.
344-
func (s *StateDB) GetState(addr common.Address, hash common.Hash) common.Hash {
344+
func (s *StateDB) GetState(addr common.Address, hash common.Hash, opts ...stateconf.StateDBStateOption) common.Hash {
345345
stateObject := s.getStateObject(addr)
346346
if stateObject != nil {
347+
hash = transformStateKey(addr, hash, opts...)
347348
return stateObject.GetState(hash)
348349
}
349350
return common.Hash{}
350351
}
351352

352353
// GetCommittedState retrieves a value from the given account's committed storage trie.
353-
func (s *StateDB) GetCommittedState(addr common.Address, hash common.Hash) common.Hash {
354+
func (s *StateDB) GetCommittedState(addr common.Address, hash common.Hash, opts ...stateconf.StateDBStateOption) common.Hash {
354355
stateObject := s.getStateObject(addr)
355356
if stateObject != nil {
357+
hash = transformStateKey(addr, hash, opts...)
356358
return stateObject.GetCommittedState(hash)
357359
}
358360
return common.Hash{}
@@ -412,9 +414,10 @@ func (s *StateDB) SetCode(addr common.Address, code []byte) {
412414
}
413415
}
414416

415-
func (s *StateDB) SetState(addr common.Address, key, value common.Hash) {
417+
func (s *StateDB) SetState(addr common.Address, key, value common.Hash, opts ...stateconf.StateDBStateOption) {
416418
stateObject := s.getOrNewStateObject(addr)
417419
if stateObject != nil {
420+
key = transformStateKey(addr, key, opts...)
418421
stateObject.SetState(key, value)
419422
}
420423
}

‎core/state/statedb.libevm.go‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121

2222
"github.com/ava-labs/libevm/common"
2323
"github.com/ava-labs/libevm/core/state/snapshot"
24+
"github.com/ava-labs/libevm/libevm/register"
2425
"github.com/ava-labs/libevm/libevm/stateconf"
2526
)
2627

@@ -52,3 +53,25 @@ func clearTypedNilPointer(snaps SnapshotTree) SnapshotTree {
5253
}
5354
return snaps
5455
}
56+
57+
type StateDBHooks interface {
58+
TransformStateKey(common.Address, common.Hash) common.Hash
59+
}
60+
61+
func RegisterExtras(s StateDBHooks) {
62+
registeredExtras.MustRegister(s)
63+
}
64+
65+
func TestOnlyClearRegisteredExtras() {
66+
registeredExtras.TestOnlyClear()
67+
}
68+
69+
var registeredExtras register.AtMostOnce[StateDBHooks]
70+
71+
func transformStateKey(addr common.Address, key common.Hash, opts ...stateconf.StateDBStateOption) common.Hash {
72+
r := &registeredExtras
73+
if !r.Registered() || !stateconf.ShouldTransformStateKey(opts...) {
74+
return key
75+
}
76+
return r.Get().TransformStateKey(addr, key)
77+
}

‎core/state/statedb.libevm_test.go‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,3 +126,54 @@ func (r *triedbRecorder) Update(
126126
func (r *triedbRecorder) Reader(_ common.Hash) (database.Reader, error) {
127127
return r.Database.Reader(common.Hash{})
128128
}
129+
130+
type highByteFlipper struct{}
131+
132+
func flipHighByte(h common.Hash) common.Hash {
133+
h[0] = ^h[0]
134+
return h
135+
}
136+
137+
func (highByteFlipper) TransformStateKey(_ common.Address, key common.Hash) common.Hash {
138+
return flipHighByte(key)
139+
}
140+
141+
func TestTransformStateKey(t *testing.T) {
142+
cache := NewDatabase(rawdb.NewMemoryDatabase())
143+
sdb, err := New(types.EmptyRootHash, cache, nil)
144+
require.NoErrorf(t, err, "New()")
145+
146+
addr := common.Address{1}
147+
regularKey := common.Hash{0, 'k', 'e', 'y'}
148+
flippedKey := flipHighByte(regularKey)
149+
regularVal := common.Hash{'r', 'e', 'g', 'u', 'l', 'a', 'r'}
150+
flippedVal := common.Hash{'f', 'l', 'i', 'p', 'p', 'e', 'd'}
151+
152+
sdb.SetState(addr, regularKey, regularVal)
153+
sdb.SetState(addr, flippedKey, flippedVal)
154+
155+
assertEq := func(t *testing.T, key, want common.Hash, opts ...stateconf.StateDBStateOption) {
156+
t.Helper()
157+
assert.Equal(t, want, sdb.GetState(addr, key, opts...))
158+
}
159+
160+
assertEq(t, regularKey, regularVal)
161+
assertEq(t, flippedKey, flippedVal)
162+
163+
// Typically the hook would be registered before any state access or
164+
// setting, but doing it here aids testing by showing the before-and-after
165+
// effects.
166+
RegisterExtras(highByteFlipper{})
167+
t.Cleanup(TestOnlyClearRegisteredExtras)
168+
169+
noTransform := stateconf.SkipStateKeyTransformation()
170+
assertEq(t, regularKey, flippedVal)
171+
assertEq(t, regularKey, regularVal, noTransform)
172+
assertEq(t, flippedKey, regularVal)
173+
assertEq(t, flippedKey, flippedVal, noTransform)
174+
175+
updatedVal := common.Hash{'u', 'p', 'd', 'a', 't', 'e', 'd'}
176+
sdb.SetState(addr, regularKey, updatedVal)
177+
assertEq(t, regularKey, updatedVal)
178+
assertEq(t, flippedKey, updatedVal, noTransform)
179+
}

‎core/vm/interface.go‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121

2222
"github.com/ava-labs/libevm/common"
2323
"github.com/ava-labs/libevm/core/types"
24+
"github.com/ava-labs/libevm/libevm/stateconf"
2425
"github.com/ava-labs/libevm/params"
2526
"github.com/holiman/uint256"
2627
)
@@ -45,9 +46,9 @@ type StateDB interface {
4546
SubRefund(uint64)
4647
GetRefund() uint64
4748

48-
GetCommittedState(common.Address, common.Hash) common.Hash
49-
GetState(common.Address, common.Hash) common.Hash
50-
SetState(common.Address, common.Hash, common.Hash)
49+
GetCommittedState(common.Address, common.Hash, ...stateconf.StateDBStateOption) common.Hash
50+
GetState(common.Address, common.Hash, ...stateconf.StateDBStateOption) common.Hash
51+
SetState(common.Address, common.Hash, common.Hash, ...stateconf.StateDBStateOption)
5152

5253
GetTransientState(addr common.Address, key common.Hash) common.Hash
5354
SetTransientState(addr common.Address, key, value common.Hash)

‎libevm/libevm.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020
"github.com/holiman/uint256"
2121

2222
"github.com/ava-labs/libevm/common"
23+
"github.com/ava-labs/libevm/libevm/stateconf"
2324
)
2425

2526
// PrecompiledContract is an exact copy of vm.PrecompiledContract, mirrored here
@@ -43,8 +44,8 @@ type StateReader interface {
4344

4445
GetRefund() uint64
4546

46-
GetCommittedState(common.Address, common.Hash) common.Hash
47-
GetState(common.Address, common.Hash) common.Hash
47+
GetCommittedState(common.Address, common.Hash, ...stateconf.StateDBStateOption) common.Hash
48+
GetState(common.Address, common.Hash, ...stateconf.StateDBStateOption) common.Hash
4849

4950
GetTransientState(addr common.Address, key common.Hash) common.Hash
5051

‎libevm/stateconf/conf.go‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,3 +113,19 @@ func ExtractTrieDBUpdatePayload(opts ...TrieDBUpdateOption) (common.Hash, common
113113
}
114114
return *conf.parentBlockHash, *conf.currentBlockHash, true
115115
}
116+
117+
type StateDBStateOption = options.Option[stateDBStateConfig]
118+
119+
type stateDBStateConfig struct {
120+
skipKeyTransformation bool
121+
}
122+
123+
func SkipStateKeyTransformation() StateDBStateOption {
124+
return options.Func[stateDBStateConfig](func(c *stateDBStateConfig) {
125+
c.skipKeyTransformation = true
126+
})
127+
}
128+
129+
func ShouldTransformStateKey(opts ...StateDBStateOption) bool {
130+
return !options.As(opts...).skipKeyTransformation
131+
}

0 commit comments

Comments
 (0)