feat: Enhance purechecks for CryptoGetAccountBalanceHandler (#12839)

Signed-off-by: Thomas Moran <152873392+thomas-swirlds-labs@users.noreply.github.com>
hashgraph · May 8, 2024 · 0970b99 · 0970b99
1 parent c5ab020
commit 0970b99
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 6 deletions.
diff --git a/.../java/com/hedera/node/app/service/token/impl/handlers/CryptoGetAccountBalanceHandler.java b/.../java/com/hedera/node/app/service/token/impl/handlers/CryptoGetAccountBalanceHandler.java
@@ -21,10 +21,14 @@
 import static com.hedera.hapi.node.base.ResponseCodeEnum.INVALID_ACCOUNT_ID;
 import static com.hedera.hapi.node.base.ResponseCodeEnum.INVALID_CONTRACT_ID;
 import static com.hedera.hapi.node.base.ResponseCodeEnum.OK;
+import static com.hedera.node.app.spi.validation.Validations.mustExist;
+import static com.hedera.node.app.spi.validation.Validations.validateAccountID;
 import static com.hedera.node.app.spi.workflows.PreCheckException.validateFalsePreCheck;
+import static com.hedera.node.app.spi.workflows.PreCheckException.validateTruePreCheck;
 import static java.util.Objects.requireNonNull;
 
 import com.hedera.hapi.node.base.AccountID;
+import com.hedera.hapi.node.base.ContractID;
 import com.hedera.hapi.node.base.HederaFunctionality;
 import com.hedera.hapi.node.base.QueryHeader;
 import com.hedera.hapi.node.base.ResponseHeader;
@@ -83,18 +87,40 @@ public void validate(@NonNull final QueryContext context) throws PreCheckExcepti
         final var accountStore = context.createStore(ReadableAccountStore.class);
         final CryptoGetAccountBalanceQuery op = query.cryptogetAccountBalanceOrThrow();
         if (op.hasAccountID()) {
-            final var account = accountStore.getAliasedAccountById(requireNonNull(op.accountID()));
-            validateFalsePreCheck(account == null, INVALID_ACCOUNT_ID);
-            validateFalsePreCheck(account.deleted(), ACCOUNT_DELETED);
+            validateAccountId(op, accountStore);
         } else if (op.hasContractID()) {
-            final var contract = accountStore.getContractById(requireNonNull(op.contractID()));
-            validateFalsePreCheck(contract == null || !contract.smartContract(), INVALID_CONTRACT_ID);
-            validateFalsePreCheck(contract.deleted(), CONTRACT_DELETED);
+            validateContractId(op, accountStore);
         } else {
             throw new PreCheckException(INVALID_ACCOUNT_ID);
         }
     }
 
+    private void validateContractId(CryptoGetAccountBalanceQuery op, ReadableAccountStore accountStore)
+            throws PreCheckException {
+        mustExist(op.contractID(), INVALID_CONTRACT_ID);
+        final ContractID contractId = (ContractID) op.balanceSource().value();
+        validateTruePreCheck(contractId.shardNum() == 0, INVALID_CONTRACT_ID);
+        validateTruePreCheck(contractId.realmNum() == 0, INVALID_CONTRACT_ID);
+        validateTruePreCheck(
+                (contractId.hasContractNum() && contractId.contractNumOrThrow() >= 0) || contractId.hasEvmAddress(),
+                INVALID_CONTRACT_ID);
+        final var contract = accountStore.getContractById(requireNonNull(op.contractID()));
+        validateFalsePreCheck(contract == null, INVALID_CONTRACT_ID);
+        validateTruePreCheck(contract.smartContract(), INVALID_CONTRACT_ID);
+        validateFalsePreCheck(contract.deleted(), CONTRACT_DELETED);
+    }
+
+    private void validateAccountId(CryptoGetAccountBalanceQuery op, ReadableAccountStore accountStore)
+            throws PreCheckException {
+        AccountID accountId = (AccountID) op.balanceSource().value();
+        validateTruePreCheck(accountId.shardNum() == 0, INVALID_ACCOUNT_ID);
+        validateTruePreCheck(accountId.realmNum() == 0, INVALID_ACCOUNT_ID);
+        validateAccountID(accountId, INVALID_ACCOUNT_ID);
+        final var account = accountStore.getAliasedAccountById(requireNonNull(op.accountID()));
+        validateFalsePreCheck(account == null, INVALID_ACCOUNT_ID);
+        validateFalsePreCheck(account.deleted(), ACCOUNT_DELETED);
+    }
+
     @Override
     public Response findResponse(@NonNull final QueryContext context, @NonNull final ResponseHeader header) {
         requireNonNull(context);

diff --git a/.../hedera/node/app/service/token/impl/test/handlers/CryptoGetAccountBalanceHandlerTest.java b/.../hedera/node/app/service/token/impl/test/handlers/CryptoGetAccountBalanceHandlerTest.java
@@ -177,6 +177,25 @@ void validatesQueryIfInvalidAccount() throws Throwable {
                 .has(responseCode(ResponseCodeEnum.INVALID_ACCOUNT_ID));
     }
 
+    @Test
+    @DisplayName("Account Id with valid header is needed during validate")
+    void validatesQueryIfInvalidAccountHeader() throws Throwable {
+        final var state =
+                MapReadableKVState.<AccountID, Account>builder(ACCOUNTS).build();
+        given(readableStates.<AccountID, Account>get(ACCOUNTS)).willReturn(state);
+        final var store = new ReadableAccountStoreImpl(readableStates);
+        final AccountID invalidRealmAccountId =
+                AccountID.newBuilder().accountNum(5).realmNum(-1L).build();
+
+        final var query = createGetAccountBalanceQueryWithInvalidHeader(invalidRealmAccountId.accountNumOrThrow());
+        when(context.query()).thenReturn(query);
+        when(context.createStore(ReadableAccountStore.class)).thenReturn(store);
+
+        assertThatThrownBy(() -> subject.validate(context))
+                .isInstanceOf(PreCheckException.class)
+                .has(responseCode(ResponseCodeEnum.INVALID_ACCOUNT_ID));
+    }
+
     @Test
     @DisplayName("Contract Id is needed during validate")
     void validatesQueryIfInvalidContract() throws Throwable {
@@ -466,6 +485,15 @@ private Query createGetAccountBalanceQuery(final long accountId) {
         return Query.newBuilder().cryptogetAccountBalance(data).build();
     }
 
+    private Query createGetAccountBalanceQueryWithInvalidHeader(final long accountId) {
+        final var data = CryptoGetAccountBalanceQuery.newBuilder()
+                .accountID(AccountID.newBuilder().accountNum(accountId).build())
+                .header((QueryHeader) null)
+                .build();
+
+        return Query.newBuilder().cryptogetAccountBalance(data).build();
+    }
+
     private Query createGetAccountBalanceQueryWithContract(final long contractId) {
         final var data = CryptoGetAccountBalanceQuery.newBuilder()
                 .contractID(ContractID.newBuilder().contractNum(contractId).build())