From ef61d09e59ef415f696b2779d60376ebf60623f8 Mon Sep 17 00:00:00 2001
From: acolytec3 <17355484+acolytec3@users.noreply.github.com>
Date: Thu, 9 Nov 2023 16:35:28 -0500
Subject: [PATCH] Update statemanager to use state.getBytecode

---
 .../src/networks/state/stateManager.ts              |  7 +++++--
 .../portalnetwork/src/networks/state/statedb.ts     |  8 ++++----
 .../test/networks/state/stateManager.spec.ts        | 13 ++++++++++++-
 3 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/packages/portalnetwork/src/networks/state/stateManager.ts b/packages/portalnetwork/src/networks/state/stateManager.ts
index 3ca74d0b1..6bbbdae29 100644
--- a/packages/portalnetwork/src/networks/state/stateManager.ts
+++ b/packages/portalnetwork/src/networks/state/stateManager.ts
@@ -2,6 +2,7 @@ import { EVMStateManagerInterface, Proof, StorageDump, StorageRange } from '@eth
 import { Address, Account, bytesToBigInt, bytesToHex } from '@ethereumjs/util'
 
 import { StateNetwork } from './state.js'
+import { toHexString } from '@chainsafe/ssz'
 
 export class UltralightStateManager implements EVMStateManagerInterface {
   originalStorageCache: {
@@ -34,7 +35,7 @@ export class UltralightStateManager implements EVMStateManagerInterface {
     return new UltralightStateManager(this.state)
   }
   getAccount(address: Address): Promise<Account | undefined> {
-    return this.state.stateDB.getAccount(address.toString(), this.stateRoot)
+    return this.state.getAccount(address.toString(), this.stateRoot)
   }
   putAccount = async (address: Address, account?: Account | undefined): Promise<void> => {
     return undefined
@@ -52,7 +53,9 @@ export class UltralightStateManager implements EVMStateManagerInterface {
     throw new Error('Method not implemented.')
   }
   getContractCode = async (address: Address): Promise<Uint8Array> => {
-    const code = await this.state.stateDB.getCode(address.toString(), this.stateRoot)
+    const account = await this.state.getAccount(address.toString(), this.stateRoot)
+    const code =
+      account && (await this.state.getBytecode(toHexString(account.codeHash), address.toString()))
     return code ?? new Uint8Array()
   }
   getContractStorage = async (address: Address, key: Uint8Array): Promise<Uint8Array> => {
diff --git a/packages/portalnetwork/src/networks/state/statedb.ts b/packages/portalnetwork/src/networks/state/statedb.ts
index f0b314118..6277d34dc 100644
--- a/packages/portalnetwork/src/networks/state/statedb.ts
+++ b/packages/portalnetwork/src/networks/state/statedb.ts
@@ -214,13 +214,13 @@ export class StateDB {
 
   /**
    * Get account data by address and state root
-   * @param address account address
-   * @param stateRoot state root
+   * @param account address as hex prefixed string
+   * @param stateRoot state root as hex prefixed string
    * @returns account data
    */
-  async getAccount(address: Address, stateRoot: StateRoot) {
+  async getAccount(address: string, stateRoot: StateRoot) {
     const trie = this.getAccountTrie(stateRoot)
-    const key = fromHexString(address.toString())
+    const key = fromHexString(address)
     const accountRLP = await trie.get(key)
     if (accountRLP === null) {
       return undefined
diff --git a/packages/portalnetwork/test/networks/state/stateManager.spec.ts b/packages/portalnetwork/test/networks/state/stateManager.spec.ts
index ad25c645d..80ab98e20 100644
--- a/packages/portalnetwork/test/networks/state/stateManager.spec.ts
+++ b/packages/portalnetwork/test/networks/state/stateManager.spec.ts
@@ -108,10 +108,14 @@ describe('UltralightStateManager', () => {
     const codehash = keccak256(byteCode)
     const account = Account.fromAccountData({ balance: 0n, nonce: 1n, codeHash: codehash })
 
+    const zero = Address.zero()
+    const zeroAccount = new Account()
     const trie = new Trie({ useKeyHashing: true })
     await trie.put(address.toBytes(), account.serialize())
+    await trie.put(zero.bytes, zeroAccount.serialize())
 
     const proof = await trie.createProof(address.toBytes())
+    const zeroProof = await trie.createProof(zero.bytes)
     const content = AccountTrieProofType.serialize({
       balance: account!.balance,
       nonce: account!.nonce,
@@ -119,8 +123,15 @@ describe('UltralightStateManager', () => {
       storageRoot: account!.storageRoot,
       witnesses: proof,
     })
+    const zeroContent = AccountTrieProofType.serialize({
+      balance: zeroAccount!.balance,
+      nonce: zeroAccount!.nonce,
+      codeHash: zeroAccount!.codeHash,
+      storageRoot: zeroAccount!.storageRoot,
+      witnesses: zeroProof,
+    })
     await network.stateDB.inputAccountTrieProof(address.toBytes(), trie.root(), content)
-
+    await network.stateDB.inputAccountTrieProof(zero.bytes, trie.root(), zeroContent)
     const byteCodeContent = ContractByteCodeType.serialize(byteCode)
     await network.stateDB.inputContractByteCode(address.toBytes(), codehash, byteCodeContent)