From 1dad6be36fbc79b2c05b3b57455a1f86b9aa615d Mon Sep 17 00:00:00 2001 From: wb Date: Fri, 30 Oct 2026 15:54:54 +0800 Subject: [PATCH] fix(db): validate block witness signatures against committed state Recheck received block witness authorization under the Manager lock after pending state is reset, so concurrent permission changes and pending transactions cannot authorize block admission. Keep witness validation before advertisement broadcast and defer sync validation to Manager. Update the shutdown target only after witness authorization succeeds. --- .../main/java/org/tron/core/db/Manager.java | 21 +- .../org/tron/core/net/TronNetDelegate.java | 6 +- .../core/db/ManagerWitnessSignatureTest.java | 401 ++++++++++++++++++ .../tron/core/net/TronNetDelegateTest.java | 16 +- .../messagehandler/BlockMsgHandlerTest.java | 52 +++ 5 files changed, 481 insertions(+), 15 deletions(-) create mode 100644 framework/src/test/java/org/tron/core/db/ManagerWitnessSignatureTest.java diff --git a/framework/src/main/java/org/tron/core/db/Manager.java b/framework/src/main/java/org/tron/core/db/Manager.java index 9e5d2db7547..25097f2e5ac 100644 --- a/framework/src/main/java/org/tron/core/db/Manager.java +++ b/framework/src/main/java/org/tron/core/db/Manager.java @@ -1228,7 +1228,8 @@ public List getVerifyTxs(BlockCapsule block) { /** * Save a block. Untrusted callers must first use network prevalidation to check its timestamp, - * transaction structure, Merkle root and block signature. + * transaction structure and Merkle root. Check witness signatures before broadcast as well. + * Received block witness authorization is checked against committed state before admission. */ public void pushBlock(final BlockCapsule block) throws ValidateSignatureException, ContractValidateException, ContractExeException, @@ -1261,18 +1262,24 @@ public void pushBlock(final BlockCapsule block) block.getNum(), rePushTransactions.size(), pendingTransactions.size(), block.getTransactions().size(), txs.size()); - if (CommonParameter.getInstance().getShutdownBlockTime() != null - && CommonParameter.getInstance().getShutdownBlockTime() - .isSatisfiedBy(new Date(block.getTimeStamp()))) { - latestSolidityNumShutDown = block.getNum(); - } - try (PendingManager pm = new PendingManager(this)) { if (!block.generatedByMyself) { + // Pending changes have been reset; use committed witness permissions for admission. + // switchFork also verifies admitted fork blocks against their own parent state. + if (!block.validateSignature(getDynamicPropertiesStore(), getAccountStore())) { + throw new ValidateSignatureException( + "block " + block.getNum() + " signature invalid"); + } consensus.receiveBlock(block); } + if (CommonParameter.getInstance().getShutdownBlockTime() != null + && CommonParameter.getInstance().getShutdownBlockTime() + .isSatisfiedBy(new Date(block.getTimeStamp()))) { + latestSolidityNumShutDown = block.getNum(); + } + if (block.getTransactions().stream() .filter(tran -> isShieldedTransaction(tran.getInstance())) .count() > SHIELDED_TRANS_IN_BLOCK_COUNTS) { diff --git a/framework/src/main/java/org/tron/core/net/TronNetDelegate.java b/framework/src/main/java/org/tron/core/net/TronNetDelegate.java index 2948862b181..6b5daf75dd0 100644 --- a/framework/src/main/java/org/tron/core/net/TronNetDelegate.java +++ b/framework/src/main/java/org/tron/core/net/TronNetDelegate.java @@ -353,6 +353,10 @@ public void validSignature(BlockCapsule block) throws P2pException { } } + /** + * Validate block structure and Merkle root. Advertisement also checks witness authorization; + * sync defers it to Manager, which checks committed state and revalidates forks after rewind. + */ public void validBlock(BlockCapsule block) throws P2pException { for (Transaction transaction : block.getInstance().getTransactionsList()) { int contractCount = transaction.getRawData().getContractCount(); @@ -373,7 +377,6 @@ public void validBlock(BlockCapsule block) throws P2pException { } catch (BadBlockException e) { throw new P2pException(TypeEnum.BLOCK_MERKLE_INVALID, e.getMessage()); } - validSignature(block); } public boolean validAdvBlock(PeerConnection peer, BlockCapsule block) throws P2pException { @@ -391,6 +394,7 @@ public boolean validAdvBlock(PeerConnection peer, BlockCapsule block) throws P2p } } validBlock(block); + validSignature(block); return activeWitness; } diff --git a/framework/src/test/java/org/tron/core/db/ManagerWitnessSignatureTest.java b/framework/src/test/java/org/tron/core/db/ManagerWitnessSignatureTest.java new file mode 100644 index 00000000000..cb6ff5aa2ba --- /dev/null +++ b/framework/src/test/java/org/tron/core/db/ManagerWitnessSignatureTest.java @@ -0,0 +1,401 @@ +package org.tron.core.db; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.protobuf.ByteString; +import java.lang.reflect.Method; +import java.math.BigInteger; +import java.net.InetSocketAddress; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.Date; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Test; +import org.tron.common.BaseMethodTest; +import org.tron.common.cron.CronExpression; +import org.tron.common.crypto.ECKey; +import org.tron.common.parameter.CommonParameter; +import org.tron.common.utils.ReflectUtils; +import org.tron.common.utils.Sha256Hash; +import org.tron.core.capsule.AccountCapsule; +import org.tron.core.capsule.BlockCapsule; +import org.tron.core.capsule.TransactionCapsule; +import org.tron.core.capsule.WitnessCapsule; +import org.tron.core.consensus.ConsensusService; +import org.tron.core.exception.UnLinkedBlockException; +import org.tron.core.exception.ValidateSignatureException; +import org.tron.core.net.TronNetDelegate; +import org.tron.core.net.message.adv.BlockMessage; +import org.tron.core.net.messagehandler.BlockMsgHandler; +import org.tron.core.net.messagehandler.PbftDataSyncHandler; +import org.tron.core.net.peer.PeerConnection; +import org.tron.core.net.service.adv.AdvService; +import org.tron.core.net.service.sync.SyncService; +import org.tron.core.services.WitnessProductBlockService; +import org.tron.core.store.AccountStore; +import org.tron.core.store.DynamicPropertiesStore; +import org.tron.protos.Protocol.Account; +import org.tron.protos.Protocol.ReasonCode; +import org.tron.protos.Protocol.Transaction.Contract.ContractType; +import org.tron.protos.contract.AccountContract.AccountPermissionUpdateContract; + +public class ManagerWitnessSignatureTest extends BaseMethodTest { + + private final ECKey owner = ECKey.fromPrivate(BigInteger.valueOf(101)); + private final ECKey signer = ECKey.fromPrivate(BigInteger.valueOf(102)); + private final ECKey replacement = ECKey.fromPrivate(BigInteger.valueOf(103)); + private ByteString witness; + + @Override + protected void afterInit() { + context.getBean(ConsensusService.class).start(); + witness = ByteString.copyFrom(owner.getAddress()); + chainBaseManager.getDynamicPropertiesStore().saveAllowMultiSign(1); + chainBaseManager.getAccountStore().put(witness.toByteArray(), new AccountCapsule( + Account.newBuilder().setAddress(witness).setIsWitness(true).setBalance(1_000_000_000_000L) + .setOwnerPermission(AccountCapsule.createDefaultOwnerPermission(witness)) + .addActivePermission(AccountCapsule.createDefaultActivePermission(witness, + chainBaseManager.getDynamicPropertiesStore())) + .setWitnessPermission(AccountCapsule.createDefaultWitnessPermission( + ByteString.copyFrom(signer.getAddress()))).build())); + chainBaseManager.getWitnessStore().put(witness.toByteArray(), new WitnessCapsule(witness)); + chainBaseManager.getWitnessScheduleStore() + .saveActiveWitnesses(Collections.singletonList(witness)); + } + + @Test + public void acceptsCommittedWitnessPermission() throws Exception { + BlockCapsule block = nextBlock(signer); + context.getBean(TronNetDelegate.class).validSignature(block); + + dbManager.pushBlock(block); + + assertEquals(block.getBlockId(), chainBaseManager.getHeadBlockId()); + assertEquals(block.getBlockId(), chainBaseManager.getKhaosDb().getHeadId()); + assertTrue(chainBaseManager.getBlockStore().has(block.getBlockId().getBytes())); + } + + @Test + public void rejectsWitnessRevokedWhilePushWaitsForManagerLock() throws Exception { + BlockCapsule block = spy(nextBlock(signer)); + AtomicInteger validations = new AtomicInteger(); + doAnswer(invocation -> { + if (validations.incrementAndGet() > 1) { + assertTrue(Thread.holdsLock(dbManager)); + } + return invocation.callRealMethod(); + }).when(block).validateSignature(any(DynamicPropertiesStore.class), any(AccountStore.class)); + CountDownLatch ingressValidated = new CountDownLatch(1); + ExecutorService executor = Executors.newSingleThreadExecutor(); + Future pushed; + try { + synchronized (dbManager) { + pushed = executor.submit(() -> { + context.getBean(TronNetDelegate.class).validSignature(block); + ingressValidated.countDown(); + return assertThrows(ValidateSignatureException.class, () -> dbManager.pushBlock(block)); + }); + assertTrue(ingressValidated.await(10, TimeUnit.SECONDS)); + setWitnessPermission(replacement); + } + pushed.get(10, TimeUnit.SECONDS); + assertRejected(block); + } finally { + executor.shutdownNow(); + assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } + } + + @Test + public void rejectsWitnessAuthorizedOnlyByPendingSession() throws Exception { + setWitnessPermission(replacement); + dbManager.getSession().setValue(dbManager.getRevokingStore().buildSession()); + setWitnessPermission(signer); + BlockCapsule block = nextBlock(signer); + context.getBean(TronNetDelegate.class).validSignature(block); + + assertThrows(ValidateSignatureException.class, () -> dbManager.pushBlock(block)); + + assertRejected(block); + assertFalse(dbManager.getSession().valid()); + assertEquals(ByteString.copyFrom(replacement.getAddress()), + ByteString.copyFrom(chainBaseManager.getAccountStore().get(witness.toByteArray()) + .getWitnessPermissionAddress())); + } + + @Test + public void acceptsCommittedPermissionAfterPendingRevocationReset() throws Exception { + BlockCapsule block = nextBlock(signer); + context.getBean(TronNetDelegate.class).validSignature(block); + dbManager.getSession().setValue(dbManager.getRevokingStore().buildSession()); + setWitnessPermission(replacement); + assertFalse(block.validateSignature(dbManager.getDynamicPropertiesStore(), + dbManager.getAccountStore())); + + dbManager.pushBlock(block); + + assertEquals(block.getBlockId(), chainBaseManager.getHeadBlockId()); + assertEquals(ByteString.copyFrom(signer.getAddress()), + ByteString.copyFrom(chainBaseManager.getAccountStore().get(witness.toByteArray()) + .getWitnessPermissionAddress())); + } + + @Test + public void rejectsWellFormedUnauthorizedWitness() throws Exception { + BlockCapsule block = nextBlock(replacement); + assertFalse(block.validateSignature(dbManager.getDynamicPropertiesStore(), + dbManager.getAccountStore())); + + assertThrows(ValidateSignatureException.class, () -> dbManager.pushBlock(block)); + + assertRejected(block); + } + + @Test + public void rejectsMalformedWitnessSignature() throws Exception { + chainBaseManager.getDynamicPropertiesStore().saveAllowStrictEcdsaValidation(1); + BlockCapsule signed = nextBlock(signer); + BlockCapsule block = new BlockCapsule(signed.getInstance().toBuilder() + .setBlockHeader(signed.getInstance().getBlockHeader().toBuilder() + .setWitnessSignature(ByteString.copyFrom(new byte[64]))).build()); + + assertThrows(ValidateSignatureException.class, () -> dbManager.pushBlock(block)); + + assertRejected(block); + } + + @Test + public void preservesTrustedGeneratedBlockBypassAndDuplicateHandling() throws Exception { + BlockCapsule block = nextBlock(replacement); + block.generatedByMyself = true; + setMatchingShutdownTime(block); + dbManager.pushBlock(block); + assertEquals(block.getNum(), dbManager.getLatestSolidityNumShutDown()); + setWitnessPermission(owner); + + dbManager.pushBlock(new BlockCapsule(block.getData())); + + assertEquals(block.getBlockId(), chainBaseManager.getHeadBlockId()); + assertEquals(block.getBlockId(), chainBaseManager.getKhaosDb().getHeadId()); + } + + @Test + public void rejectsForkWitnessNotAuthorizedAtCurrentHead() throws Exception { + BlockCapsule parent = nextBlock(signer); + dbManager.pushBlock(parent); + BlockCapsule oldHead = blockRevokingWitness(parent); + BlockCapsule fork = signedBlock(parent.getNum() + 1, parent.getBlockId(), + parent.getTimeStamp() + 3001, signer); + + assertThrows(ValidateSignatureException.class, () -> dbManager.pushBlock(fork)); + + assertEquals(oldHead.getBlockId(), chainBaseManager.getHeadBlockId()); + assertEquals(oldHead.getBlockId(), chainBaseManager.getKhaosDb().getHeadId()); + assertFalse(chainBaseManager.getKhaosDb().containBlockInMiniStore(fork.getBlockId())); + assertFalse(chainBaseManager.getBlockStore().has(fork.getBlockId().getBytes())); + } + + @Test + public void revalidatesAdmittedForkAgainstRewoundParentPermission() throws Exception { + BlockCapsule parent = nextBlock(signer); + dbManager.pushBlock(parent); + BlockCapsule oldHead = blockRevokingWitness(parent); + BlockCapsule fork = signedBlock(parent.getNum() + 1, parent.getBlockId(), + parent.getTimeStamp() + 3001, replacement); + dbManager.pushBlock(fork); + assertTrue(chainBaseManager.getKhaosDb().containBlockInMiniStore(fork.getBlockId())); + BlockCapsule forkHead = signedBlock(fork.getNum() + 1, fork.getBlockId(), + parent.getTimeStamp() + 6000, replacement); + + synchronized (dbManager) { + assertThrows(ValidateSignatureException.class, () -> dbManager.pushBlock(forkHead)); + + assertEquals(oldHead.getBlockId(), chainBaseManager.getHeadBlockId()); + assertEquals(oldHead.getBlockId(), chainBaseManager.getKhaosDb().getHeadId()); + assertFalse(chainBaseManager.getKhaosDb().containBlockInMiniStore(fork.getBlockId())); + assertFalse(chainBaseManager.getKhaosDb().containBlockInMiniStore(forkHead.getBlockId())); + assertEquals(ByteString.copyFrom(replacement.getAddress()), + ByteString.copyFrom(chainBaseManager.getAccountStore().get(witness.toByteArray()) + .getWitnessPermissionAddress())); + } + } + + @Test + public void syncValidatesWitnessOnceUnderManagerLock() throws Exception { + List validationLocks = new ArrayList<>(); + BlockCapsule block = recordValidationLocks(nextBlock(signer), validationLocks); + + processSyncBlock(block); + + assertEquals(Collections.singletonList(true), validationLocks); + assertEquals(block.getBlockId(), chainBaseManager.getHeadBlockId()); + } + + @Test + public void syncRejectsUnauthorizedWitnessBeforeChainOrKhaosAdmission() throws Exception { + List validationLocks = new ArrayList<>(); + BlockCapsule block = recordValidationLocks(nextBlock(replacement), validationLocks); + + PeerConnection peer = processSyncBlock(block); + + assertEquals(Collections.singletonList(true), validationLocks); + verify(peer).disconnect(ReasonCode.BAD_BLOCK); + assertRejected(block); + } + + @Test + public void invalidSyncWitnessCannotChangeMatchingShutdownTarget() throws Exception { + BlockCapsule block = nextBlock(replacement); + setMatchingShutdownTime(block); + long shutdownTarget = dbManager.getLatestSolidityNumShutDown(); + + PeerConnection peer = processSyncBlock(block); + + assertEquals(shutdownTarget, dbManager.getLatestSolidityNumShutDown()); + verify(peer).disconnect(ReasonCode.BAD_BLOCK); + assertRejected(block); + } + + @Test + public void validSyncWitnessUpdatesMatchingShutdownTarget() throws Exception { + BlockCapsule block = nextBlock(signer); + setMatchingShutdownTime(block); + + processSyncBlock(block); + + assertEquals(block.getNum(), dbManager.getLatestSolidityNumShutDown()); + assertEquals(block.getBlockId(), chainBaseManager.getHeadBlockId()); + } + + @Test + public void authorizedBlockRetainsShutdownSchedulingBeforeLaterAdmissionFailure() + throws Exception { + BlockCapsule block = signedBlock(1, Sha256Hash.of(new byte[]{1}), 1533529947000L, signer); + setMatchingShutdownTime(block); + + assertThrows(UnLinkedBlockException.class, () -> dbManager.pushBlock(block)); + + assertEquals(block.getNum(), dbManager.getLatestSolidityNumShutDown()); + assertRejected(block); + } + + @Test + public void advertisedBlockValidatesBeforeBroadcastAndAgainUnderManagerLock() throws Exception { + List validationLocks = new ArrayList<>(); + BlockCapsule block = recordValidationLocks(nextBlock(signer), validationLocks); + AdvService adv = mock(AdvService.class); + doAnswer(invocation -> { + assertEquals(Collections.singletonList(false), validationLocks); + assertEquals(0, chainBaseManager.getHeadBlockNum()); + return null; + }).when(adv).broadcast(any(BlockMessage.class)); + BlockMsgHandler handler = new BlockMsgHandler(); + ReflectUtils.setFieldValue(handler, "tronNetDelegate", context.getBean(TronNetDelegate.class)); + ReflectUtils.setFieldValue(handler, "advService", adv); + ReflectUtils.setFieldValue(handler, "witnessProductBlockService", + mock(WitnessProductBlockService.class)); + Method process = BlockMsgHandler.class.getDeclaredMethod("processBlock", + PeerConnection.class, BlockCapsule.class); + process.setAccessible(true); + + process.invoke(handler, new PeerConnection(), block); + + assertEquals(Arrays.asList(false, true), validationLocks); + verify(adv).broadcast(any(BlockMessage.class)); + assertEquals(block.getBlockId(), chainBaseManager.getHeadBlockId()); + } + + private void setMatchingShutdownTime(BlockCapsule block) throws Exception { + CronExpression schedule = new CronExpression("* * * * * ?"); + assertTrue(schedule.isSatisfiedBy(new Date(block.getTimeStamp()))); + CommonParameter.getInstance().setShutdownBlockTime(schedule); + } + + private PeerConnection processSyncBlock(BlockCapsule block) throws Exception { + SyncService service = context.getBean(SyncService.class); + ReflectUtils.setFieldValue(service, "pbftDataSyncHandler", mock(PbftDataSyncHandler.class)); + PeerConnection peer = mock(PeerConnection.class); + when(peer.getInetSocketAddress()).thenReturn(new InetSocketAddress("127.0.0.2", 10001)); + Method process = SyncService.class.getDeclaredMethod("processSyncBlock", + BlockCapsule.class, PeerConnection.class); + process.setAccessible(true); + process.invoke(service, block, peer); + return peer; + } + + private BlockCapsule recordValidationLocks(BlockCapsule original, List locks) + throws Exception { + BlockCapsule block = spy(original); + doAnswer(invocation -> { + locks.add(Thread.holdsLock(dbManager)); + return invocation.callRealMethod(); + }).when(block).validateSignature(any(DynamicPropertiesStore.class), any(AccountStore.class)); + return block; + } + + private BlockCapsule blockRevokingWitness(BlockCapsule parent) throws Exception { + AccountPermissionUpdateContract contract = AccountPermissionUpdateContract.newBuilder() + .setOwnerAddress(witness) + .setOwner(AccountCapsule.createDefaultOwnerPermission(witness)) + .setWitness(AccountCapsule.createDefaultWitnessPermission( + ByteString.copyFrom(replacement.getAddress()))) + .addActives(AccountCapsule.createDefaultActivePermission(witness, + chainBaseManager.getDynamicPropertiesStore())).build(); + TransactionCapsule update = new TransactionCapsule(contract, + ContractType.AccountPermissionUpdateContract); + update.setReference(parent.getNum(), parent.getBlockId().getBytes()); + update.setExpiration(parent.getTimeStamp() + 1_000_000L); + update.sign(owner.getPrivKeyBytes()); + BlockCapsule oldHead = nextBlock(signer); + oldHead.addTransaction(update); + oldHead.setMerkleRoot(); + oldHead.sign(signer.getPrivKeyBytes()); + dbManager.pushBlock(oldHead); + return oldHead; + } + + private void setWitnessPermission(ECKey key) { + Account account = chainBaseManager.getAccountStore().get(witness.toByteArray()).getInstance(); + chainBaseManager.getAccountStore().put(witness.toByteArray(), new AccountCapsule(account + .toBuilder().setWitnessPermission(AccountCapsule.createDefaultWitnessPermission( + ByteString.copyFrom(key.getAddress()))).build())); + } + + private BlockCapsule nextBlock(ECKey key) { + long time = chainBaseManager.getHeadBlockNum() == 0 ? 1533529947000L + : chainBaseManager.getHeadBlockTimeStamp() + 3000; + return signedBlock(chainBaseManager.getHeadBlockNum() + 1, + chainBaseManager.getHeadBlockId(), time, key); + } + + private BlockCapsule signedBlock(long number, Sha256Hash parent, long time, ECKey key) { + BlockCapsule block = new BlockCapsule(number, parent, time, witness); + block.setMerkleRoot(); + block.sign(key.getPrivKeyBytes()); + return block; + } + + private void assertRejected(BlockCapsule block) { + assertEquals(0, chainBaseManager.getHeadBlockNum()); + assertEquals(chainBaseManager.getGenesisBlockId(), chainBaseManager.getHeadBlockId()); + assertEquals(chainBaseManager.getGenesisBlockId(), chainBaseManager.getKhaosDb().getHeadId()); + assertFalse(chainBaseManager.getKhaosDb().containBlockInMiniStore(block.getBlockId())); + assertFalse(chainBaseManager.getBlockStore().has(block.getBlockId().getBytes())); + } +} diff --git a/framework/src/test/java/org/tron/core/net/TronNetDelegateTest.java b/framework/src/test/java/org/tron/core/net/TronNetDelegateTest.java index 38c0b38f04b..45794f2465e 100644 --- a/framework/src/test/java/org/tron/core/net/TronNetDelegateTest.java +++ b/framework/src/test/java/org/tron/core/net/TronNetDelegateTest.java @@ -220,6 +220,7 @@ public void testValidBlockDoesNotCheckWitnessActivity() throws Exception { Mockito.doNothing().when(delegate).validSignature(block); delegate.validBlock(block); Mockito.verify(schedule, Mockito.never()).getActiveWitnesses(); + Mockito.verify(delegate, Mockito.never()).validSignature(block); } @Test @@ -320,7 +321,7 @@ public void testValidBlockAccepts80ByteBoundary() throws Exception { BlockCapsule block = blockWithTransaction(sizedTransaction(size)); signatureDelegate().validBlock(block); Mockito.verify(block).validateMerkleRoot(); - Mockito.verify(block).validateSignature(Mockito.any(), Mockito.any()); + Mockito.verify(block, Mockito.never()).validateSignature(Mockito.any(), Mockito.any()); } } @@ -411,7 +412,8 @@ public void testValidBlockAcceptsEmptyAndSingleContractTransactions() throws Exc delegate.validBlock(block); - Mockito.verify(block).validateSignature(Mockito.any(), Mockito.any()); + Mockito.verify(block).validateMerkleRoot(); + Mockito.verify(block, Mockito.never()).validateSignature(Mockito.any(), Mockito.any()); } } @@ -440,11 +442,11 @@ public void testCommonAndAdvertisedValidationDoNotRepeatIngressClockCheck() thro Assert.assertTrue(delegate.validAdvBlock(new PeerConnection(), block)); Mockito.verify(block, Mockito.times(2)).validateMerkleRoot(); - Mockito.verify(block, Mockito.times(2)).validateSignature(Mockito.any(), Mockito.any()); + Mockito.verify(block).validateSignature(Mockito.any(), Mockito.any()); } @Test - public void testValidBlockRejectsFalseSignature() throws Exception { + public void testValidSignatureRejectsFalseSignature() throws Exception { TronNetDelegate delegate = signatureDelegate(); BlockCapsule block = Mockito.spy(new BlockCapsule(1, Sha256Hash.ZERO_HASH, 0, ByteString.copyFromUtf8("witness"))); @@ -452,14 +454,14 @@ public void testValidBlockRejectsFalseSignature() throws Exception { Mockito.doReturn(false).when(block).validateSignature(Mockito.any(), Mockito.any()); P2pException error = Assert.assertThrows(P2pException.class, - () -> delegate.validBlock(block)); + () -> delegate.validSignature(block)); Assert.assertEquals(TypeEnum.BLOCK_SIGN_INVALID, error.getType()); Assert.assertEquals("valid signature failed.", error.getMessage()); } @Test - public void testValidBlockWrapsSignatureFailureWithCause() throws Exception { + public void testValidSignatureWrapsFailureWithCause() throws Exception { for (Exception failure : new Exception[]{new ValidateSignatureException("invalid signature"), new IllegalStateException("signature validation unavailable")}) { TronNetDelegate delegate = signatureDelegate(); @@ -469,7 +471,7 @@ public void testValidBlockWrapsSignatureFailureWithCause() throws Exception { Mockito.doThrow(failure).when(block).validateSignature(Mockito.any(), Mockito.any()); P2pException error = Assert.assertThrows(P2pException.class, - () -> delegate.validBlock(block)); + () -> delegate.validSignature(block)); Assert.assertEquals(TypeEnum.BLOCK_SIGN_INVALID, error.getType()); Assert.assertSame(failure, error.getCause()); diff --git a/framework/src/test/java/org/tron/core/net/messagehandler/BlockMsgHandlerTest.java b/framework/src/test/java/org/tron/core/net/messagehandler/BlockMsgHandlerTest.java index 78521764e9a..3a31ed0174b 100644 --- a/framework/src/test/java/org/tron/core/net/messagehandler/BlockMsgHandlerTest.java +++ b/framework/src/test/java/org/tron/core/net/messagehandler/BlockMsgHandlerTest.java @@ -237,6 +237,58 @@ public void testInvalidWitnessIngress() throws Exception { } } + @Test + public void testUnauthorizedWitnessRejectedBeforeBroadcast() throws Exception { + for (boolean fastForward : new boolean[]{false, true}) { + BlockCapsule block = signedWitnessBlock(); + block.sign(ECKey.fromPrivate(BigInteger.valueOf(11)).getPrivKeyBytes()); + TronNetDelegate delegate = witnessDelegate(false, block.getWitnessAddress()); + AdvService adv = Mockito.mock(AdvService.class); + RelayService relay = Mockito.mock(RelayService.class); + BlockMsgHandler ingress = witnessHandler(delegate, adv, Mockito.mock(SyncService.class)); + ReflectUtils.setFieldValue(ingress, "fastForward", fastForward); + ReflectUtils.setFieldValue(ingress, "relayService", relay); + BlockMessage message = new BlockMessage(block); + PeerConnection sender = witnessPeer(); + sender.getAdvInvRequest().put(new Item(message.getBlockId(), InventoryType.BLOCK), 0L); + + P2pException error = Assert.assertThrows(P2pException.class, + () -> ingress.processMessage(sender, message)); + + Assert.assertEquals(P2pException.TypeEnum.BLOCK_SIGN_INVALID, error.getType()); + Mockito.verify(adv, Mockito.never()).broadcast(any()); + Mockito.verify(relay, Mockito.never()).broadcast(any()); + Mockito.verify(delegate, Mockito.never()).processBlock(any(BlockCapsule.class), anyBoolean()); + } + } + + @Test + public void testWitnessValidatedBeforeBroadcastAndProcessing() throws Exception { + for (boolean fastForward : new boolean[]{false, true}) { + BlockCapsule block = signedWitnessBlock(); + TronNetDelegate delegate = witnessDelegate(false, block.getWitnessAddress()); + AdvService adv = Mockito.mock(AdvService.class); + RelayService relay = Mockito.mock(RelayService.class); + BlockMsgHandler ingress = witnessHandler(delegate, adv, Mockito.mock(SyncService.class)); + ReflectUtils.setFieldValue(ingress, "fastForward", fastForward); + ReflectUtils.setFieldValue(ingress, "relayService", relay); + BlockMessage message = new BlockMessage(block); + PeerConnection sender = witnessPeer(); + sender.getAdvInvRequest().put(new Item(message.getBlockId(), InventoryType.BLOCK), 0L); + + ingress.processMessage(sender, message); + + org.mockito.InOrder order = Mockito.inOrder(delegate, adv, relay); + order.verify(delegate).validSignature(any(BlockCapsule.class)); + if (fastForward) { + order.verify(relay).broadcast(any(BlockMessage.class)); + } else { + order.verify(adv).broadcast(any(BlockMessage.class)); + } + order.verify(delegate).processBlock(any(BlockCapsule.class), Mockito.eq(false)); + } + } + @Test public void testFutureBlockRejectedBeforeAdvertisedOrSyncProcessing() throws Exception { for (boolean syncResponse : new boolean[]{false, true}) {