From 81aafb09f33349a87e179237c8eabd597da2df63 Mon Sep 17 00:00:00 2001
From: Ho <noelwei@gmail.com>
Date: Tue, 10 Jan 2023 19:48:53 +0800
Subject: [PATCH] [FIX]: zkproof writer handle reading non-existed address
 (#197)

zkproof writer: fix issue in reading unexist address
---
 trie/zkproof/writer.go | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/trie/zkproof/writer.go b/trie/zkproof/writer.go
index 520bce1173..a49592dbf1 100644
--- a/trie/zkproof/writer.go
+++ b/trie/zkproof/writer.go
@@ -406,13 +406,12 @@ func (w *zktrieProofWriter) traceAccountUpdate(addr common.Address, updateAccDat
 			return nil, fmt.Errorf("update zktrie account state fail: %s", err)
 		}
 		w.tracingAccounts[addr] = accData
-	} else {
+	} else if accDataBefore != nil {
 		if err := w.tracingZktrie.TryDelete(addr.Bytes32()); err != nil {
 			return nil, fmt.Errorf("delete zktrie account state fail: %s", err)
 		}
 		w.tracingAccounts[addr] = nil
-
-	}
+	} // notice if both before/after is nil, we do not touch zktrie
 
 	proof = proofList{}
 	if err := w.tracingZktrie.Prove(s_key.Bytes(), 0, &proof); err != nil {
@@ -459,6 +458,7 @@ func (w *zktrieProofWriter) traceStorageUpdate(addr common.Address, key, value [
 	storeKey := zkt.NewByte32FromBytesPaddingZero(common.BytesToHash(key).Bytes())
 	storeValueBefore := trie.Get(storeKey[:])
 	storeValue := zkt.NewByte32FromBytes(value)
+	valZero := zkt.Byte32{}
 
 	if storeValueBefore != nil && !bytes.Equal(storeValueBefore[:], common.Hash{}.Bytes()) {
 		stateUpdate[0] = &StateStorage{
@@ -503,7 +503,11 @@ func (w *zktrieProofWriter) traceStorageUpdate(addr common.Address, key, value [
 	out, err := w.traceAccountUpdate(addr,
 		func(acc *types.StateAccount) *types.StateAccount {
 			if acc == nil {
-				panic("unexpected")
+				// in case we read an unexist account
+				if !bytes.Equal(valZero.Bytes(), value) {
+					panic(fmt.Errorf("write to an unexist account [%s] which is not allowed", addr))
+				}
+				return nil
 			}
 
 			//sanity check
@@ -533,6 +537,11 @@ func (w *zktrieProofWriter) traceStorageUpdate(addr common.Address, key, value [
 		} else {
 			out.StateKey = zkt.NewHashFromBigInt(h)[:]
 		}
+		stateUpdate[1] = &StateStorage{
+			Key:   storeKey.Bytes(),
+			Value: valZero.Bytes(),
+		}
+		stateUpdate[0] = stateUpdate[1]
 	}
 
 	out.StatePath = statePath