Skip to content

Commit cdf1047

Browse files
wbclaude
authored and
wb
committed
feat(net): optimize disconnectRandom by tracking block receive time per peer
- Add blockRcvTime/blockRcvTimeCmp fields to PeerConnection to track when a peer last delivered a valid block - Set blockRcvTime in BlockMsgHandler after each block is received - Fix lastInteractiveTime update in InventoryMsgHandler: only update for block inventories above current head block num, preventing attackers from forging activity via stale block hashes - Add getRandomDisconnectionPeers() to ResilienceService: narrows the disconnect candidate pool to the oldest half by blockRcvTime, so peers that recently delivered blocks are protected from random eviction Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent f0a8f0f commit cdf1047

9 files changed

Lines changed: 313 additions & 8 deletions

File tree

‎framework/src/main/java/org/tron/core/net/messagehandler/BlockMsgHandler.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,7 @@ private void processBlock(PeerConnection peer, BlockCapsule block) throws P2pExc
150150

151151
try {
152152
tronNetDelegate.processBlock(block, false);
153+
peer.setBlockRcvTime(System.currentTimeMillis());
153154
witnessProductBlockService.validWitnessProductTwoBlock(block);
154155

155156
Item item = new Item(blockId, InventoryType.BLOCK);

‎framework/src/main/java/org/tron/core/net/messagehandler/InventoryMsgHandler.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import org.springframework.beans.factory.annotation.Autowired;
55
import org.springframework.stereotype.Component;
66
import org.tron.common.utils.Sha256Hash;
7+
import org.tron.core.capsule.BlockCapsule.BlockId;
78
import org.tron.core.config.args.Args;
89
import org.tron.core.net.TronNetDelegate;
910
import org.tron.core.net.message.TronMessage;
@@ -40,7 +41,10 @@ public void processMessage(PeerConnection peer, TronMessage msg) {
4041
peer.getAdvInvReceive().put(item, System.currentTimeMillis());
4142
advService.addInv(item);
4243
if (type.equals(InventoryType.BLOCK) && peer.getAdvInvSpread().getIfPresent(item) == null) {
43-
peer.setLastInteractiveTime(System.currentTimeMillis());
44+
long headNum = tronNetDelegate.getHeadBlockId().getNum();
45+
if (new BlockId(id).getNum() > headNum) {
46+
peer.setLastInteractiveTime(System.currentTimeMillis());
47+
}
4448
}
4549
}
4650
}

‎framework/src/main/java/org/tron/core/net/peer/PeerConnection.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,14 @@ public class PeerConnection {
8888
@Setter
8989
private volatile long lastInteractiveTime;
9090

91+
@Setter
92+
@Getter
93+
private volatile long blockRcvTime;
94+
95+
@Setter
96+
@Getter
97+
private volatile long blockRcvTimeCmp;
98+
9199
@Getter
92100
@Setter
93101
private volatile TronState tronState = TronState.INIT;

‎framework/src/main/java/org/tron/core/net/service/effective/ResilienceService.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ public class ResilienceService {
4444

4545
@Autowired
4646
private ChainBaseManager chainBaseManager;
47-
47+
4848
public void init() {
4949
if (Args.getInstance().isOpenFullTcpDisconnect) {
5050
executor.scheduleWithFixedDelay(() -> {
@@ -86,6 +86,7 @@ private void disconnectRandom() {
8686
.collect(Collectors.toList());
8787

8888
if (peers.size() >= minBroadcastPeerSize) {
89+
peers = getRandomDisconnectionPeers(peers);
8990
long now = System.currentTimeMillis();
9091
Map<Object, Integer> weights = new HashMap<>();
9192
peers.forEach(peer -> {
@@ -121,6 +122,12 @@ private void disconnectRandom() {
121122
}
122123

123124

125+
private List<PeerConnection> getRandomDisconnectionPeers(List<PeerConnection> peers) {
126+
peers.forEach(p -> p.setBlockRcvTimeCmp(p.getBlockRcvTime()));
127+
peers.sort(Comparator.comparingLong(PeerConnection::getBlockRcvTimeCmp));
128+
return peers.subList(0, peers.size() / 2);
129+
}
130+
124131
private void disconnectLan() {
125132
if (!isLanNode()) {
126133
return;

‎framework/src/main/java/org/tron/core/net/service/sync/SyncService.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -305,6 +305,7 @@ private void processSyncBlock(BlockCapsule block, PeerConnection peerConnection)
305305
try {
306306
tronNetDelegate.validSignature(block);
307307
tronNetDelegate.processBlock(block, true);
308+
peerConnection.setBlockRcvTime(System.currentTimeMillis());
308309
pbftDataSyncHandler.processPBFTCommitData(block);
309310
} catch (P2pException p2pException) {
310311
logger.error("Process sync block {} failed, type: {}",

‎framework/src/test/java/org/tron/core/net/messagehandler/BlockMsgHandlerTest.java‎

Lines changed: 72 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -132,12 +132,12 @@ public void testProcessMessage() {
132132
}
133133

134134
@Test
135-
public void testProcessBlock() {
135+
public void testProcessBlock() throws Exception {
136136
TronNetDelegate tronNetDelegate = Mockito.mock(TronNetDelegate.class);
137-
137+
Field field = handler.getClass().getDeclaredField("tronNetDelegate");
138+
field.setAccessible(true);
139+
Object origin = field.get(handler);
138140
try {
139-
Field field = handler.getClass().getDeclaredField("tronNetDelegate");
140-
field.setAccessible(true);
141141
field.set(handler, tronNetDelegate);
142142

143143
BlockCapsule blockCapsule0 = new BlockCapsule(1,
@@ -164,8 +164,74 @@ public void testProcessBlock() {
164164
.getDeclaredMethod("processBlock", PeerConnection.class, BlockCapsule.class);
165165
method.setAccessible(true);
166166
method.invoke(handler, peer, blockCapsule0);
167-
} catch (Exception e) {
168-
Assert.fail();
167+
} finally {
168+
field.set(handler, origin);
169+
}
170+
}
171+
172+
@Test
173+
public void testBlockRcvTimeSetAfterProcessBlockSuccess() throws Exception {
174+
TronNetDelegate tronNetDelegate = Mockito.mock(TronNetDelegate.class);
175+
Field field = handler.getClass().getDeclaredField("tronNetDelegate");
176+
field.setAccessible(true);
177+
Object origin = field.get(handler);
178+
try {
179+
field.set(handler, tronNetDelegate);
180+
181+
BlockCapsule blockCapsule = new BlockCapsule(1,
182+
Sha256Hash.wrap(ByteString.copyFrom(ByteArray.fromHexString(
183+
"9938a342238077182498b464ac0292229938a342238077182498b464ac029222"))),
184+
1234, ByteString.copyFrom("1234567".getBytes()));
185+
186+
Mockito.doReturn(true).when(tronNetDelegate).validBlock(any(BlockCapsule.class));
187+
Mockito.doReturn(true).when(tronNetDelegate).containBlock(any(BlockId.class));
188+
Mockito.doReturn(blockCapsule.getBlockId()).when(tronNetDelegate).getHeadBlockId();
189+
Mockito.doNothing().when(tronNetDelegate).processBlock(any(BlockCapsule.class), anyBoolean());
190+
Mockito.doReturn(new ArrayList<PeerConnection>()).when(tronNetDelegate).getActivePeer();
191+
192+
peer.setBlockRcvTime(0L);
193+
Method method = handler.getClass()
194+
.getDeclaredMethod("processBlock", PeerConnection.class, BlockCapsule.class);
195+
method.setAccessible(true);
196+
197+
long before = System.currentTimeMillis();
198+
method.invoke(handler, peer, blockCapsule);
199+
long after = System.currentTimeMillis();
200+
201+
Assert.assertTrue("blockRcvTime should be set after successful processBlock",
202+
peer.getBlockRcvTime() >= before && peer.getBlockRcvTime() <= after);
203+
} finally {
204+
field.set(handler, origin);
205+
}
206+
}
207+
208+
@Test
209+
public void testBlockRcvTimeNotSetWhenValidationFails() throws Exception {
210+
TronNetDelegate tronNetDelegate = Mockito.mock(TronNetDelegate.class);
211+
Field field = handler.getClass().getDeclaredField("tronNetDelegate");
212+
field.setAccessible(true);
213+
Object origin = field.get(handler);
214+
try {
215+
field.set(handler, tronNetDelegate);
216+
217+
BlockCapsule blockCapsule = new BlockCapsule(1,
218+
Sha256Hash.wrap(ByteString.copyFrom(ByteArray.fromHexString(
219+
"9938a342238077182498b464ac0292229938a342238077182498b464ac029222"))),
220+
1234, ByteString.copyFrom("1234567".getBytes()));
221+
222+
// validBlock returns false → processBlock short-circuits before setBlockRcvTime
223+
Mockito.doReturn(false).when(tronNetDelegate).validBlock(any(BlockCapsule.class));
224+
225+
peer.setBlockRcvTime(0L);
226+
Method method = handler.getClass()
227+
.getDeclaredMethod("processBlock", PeerConnection.class, BlockCapsule.class);
228+
method.setAccessible(true);
229+
method.invoke(handler, peer, blockCapsule);
230+
231+
Assert.assertEquals("blockRcvTime must stay 0 when block fails validation",
232+
0L, peer.getBlockRcvTime());
233+
} finally {
234+
field.set(handler, origin);
169235
}
170236
}
171237
}

‎framework/src/test/java/org/tron/core/net/messagehandler/InventoryMsgHandlerTest.java‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,24 @@
11
package org.tron.core.net.messagehandler;
22

3+
import static org.mockito.ArgumentMatchers.any;
34
import static org.mockito.Mockito.mock;
45

56
import java.lang.reflect.Field;
67
import java.net.InetAddress;
78
import java.net.InetSocketAddress;
89
import java.util.ArrayList;
10+
import java.util.Arrays;
11+
import org.junit.Assert;
912
import org.junit.Test;
1013
import org.mockito.Mockito;
1114
import org.tron.common.TestConstants;
15+
import org.tron.common.utils.Sha256Hash;
16+
import org.tron.core.capsule.BlockCapsule.BlockId;
1217
import org.tron.core.config.args.Args;
1318
import org.tron.core.net.TronNetDelegate;
1419
import org.tron.core.net.message.adv.InventoryMessage;
1520
import org.tron.core.net.peer.PeerConnection;
21+
import org.tron.core.net.service.adv.AdvService;
1622
import org.tron.p2p.connection.Channel;
1723
import org.tron.protos.Protocol.Inventory.InventoryType;
1824

@@ -51,6 +57,85 @@ public void testProcessMessage() throws Exception {
5157
handler.processMessage(peer, msg);
5258
}
5359

60+
@Test
61+
public void testLastInteractiveTimeNotUpdatedForSolidifiedBlock() throws Exception {
62+
InventoryMsgHandler handler = new InventoryMsgHandler();
63+
Args.setParam(new String[]{}, TestConstants.TEST_CONF);
64+
65+
TronNetDelegate tronNetDelegate = mock(TronNetDelegate.class);
66+
AdvService advService = mock(AdvService.class);
67+
Mockito.when(advService.addInv(any())).thenReturn(true);
68+
// block num 100 is at head boundary — should NOT update
69+
Mockito.when(tronNetDelegate.getHeadBlockId())
70+
.thenReturn(new BlockId(Sha256Hash.ZERO_HASH, 100L));
71+
72+
Field delegateField = handler.getClass().getDeclaredField("tronNetDelegate");
73+
delegateField.setAccessible(true);
74+
delegateField.set(handler, tronNetDelegate);
75+
Field advField = handler.getClass().getDeclaredField("advService");
76+
advField.setAccessible(true);
77+
advField.set(handler, advService);
78+
79+
PeerConnection peer = new PeerConnection();
80+
peer.setChannel(getChannel("1.0.0.4", 1001));
81+
peer.setNeedSyncFromPeer(false);
82+
peer.setNeedSyncFromUs(false);
83+
peer.setLastInteractiveTime(0L);
84+
85+
// Block hash encodes num=100 (at solidified boundary — should NOT update)
86+
Sha256Hash blockHash = new BlockId(Sha256Hash.ZERO_HASH, 100L);
87+
InventoryMessage msg = new InventoryMessage(Arrays.asList(blockHash), InventoryType.BLOCK);
88+
handler.processMessage(peer, msg);
89+
90+
Assert.assertEquals("lastInteractiveTime should NOT be updated for solidified block",
91+
0L, peer.getLastInteractiveTime());
92+
}
93+
94+
@Test
95+
public void testLastInteractiveTimeUpdatedForBothPeersWithSameAboveSolidifiedBlock()
96+
throws Exception {
97+
InventoryMsgHandler handler = new InventoryMsgHandler();
98+
Args.setParam(new String[]{}, TestConstants.TEST_CONF);
99+
100+
TronNetDelegate tronNetDelegate = mock(TronNetDelegate.class);
101+
AdvService advService = mock(AdvService.class);
102+
// First call returns true (peer1), second call returns false (peer2 — already in cache)
103+
Mockito.when(advService.addInv(any())).thenReturn(true).thenReturn(false);
104+
Mockito.when(tronNetDelegate.getHeadBlockId())
105+
.thenReturn(new BlockId(Sha256Hash.ZERO_HASH, 99L));
106+
107+
Field delegateField = handler.getClass().getDeclaredField("tronNetDelegate");
108+
delegateField.setAccessible(true);
109+
delegateField.set(handler, tronNetDelegate);
110+
Field advField = handler.getClass().getDeclaredField("advService");
111+
advField.setAccessible(true);
112+
advField.set(handler, advService);
113+
114+
PeerConnection peer1 = new PeerConnection();
115+
peer1.setChannel(getChannel("1.0.0.5", 1002));
116+
peer1.setNeedSyncFromPeer(false);
117+
peer1.setNeedSyncFromUs(false);
118+
peer1.setLastInteractiveTime(0L);
119+
120+
PeerConnection peer2 = new PeerConnection();
121+
peer2.setChannel(getChannel("1.0.0.6", 1003));
122+
peer2.setNeedSyncFromPeer(false);
123+
peer2.setNeedSyncFromUs(false);
124+
peer2.setLastInteractiveTime(0L);
125+
126+
// block num 100 > solidified 99 — both peers should update
127+
Sha256Hash blockHash = new BlockId(Sha256Hash.ZERO_HASH, 100L);
128+
InventoryMessage msg = new InventoryMessage(Arrays.asList(blockHash), InventoryType.BLOCK);
129+
130+
handler.processMessage(peer1, msg);
131+
handler.processMessage(peer2, msg);
132+
133+
Assert.assertTrue("peer1 lastInteractiveTime should be updated",
134+
peer1.getLastInteractiveTime() > 0L);
135+
Assert.assertTrue("peer2 lastInteractiveTime should be updated even when addInv returns false",
136+
peer2.getLastInteractiveTime() > 0L);
137+
}
138+
54139
private Channel getChannel(String host, int port) throws Exception {
55140
Channel channel = new Channel();
56141
InetSocketAddress inetSocketAddress = new InetSocketAddress(host, port);

‎framework/src/test/java/org/tron/core/net/services/ResilienceServiceTest.java‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import java.io.IOException;
88
import java.net.InetSocketAddress;
99
import java.util.HashSet;
10+
import java.util.List;
1011
import java.util.Set;
1112
import javax.annotation.Resource;
1213
import org.junit.After;
@@ -97,6 +98,57 @@ public void testDisconnectRandom() {
9798
Assert.assertEquals(maxConnection - 1, PeerManager.getPeers().size());
9899
}
99100

101+
@Test
102+
public void testDisconnectRandomPreservesRecentBlockRcvTimePeer() {
103+
int maxConnection = 30;
104+
Assert.assertEquals(0, PeerManager.getPeers().size());
105+
106+
ApplicationContext ctx = (ApplicationContext) ReflectUtils.getFieldObject(p2pEventHandler,
107+
"ctx");
108+
109+
// Create maxConnection + 1 peers (triggers disconnectRandom)
110+
for (int i = 0; i < maxConnection + 1; i++) {
111+
InetSocketAddress inetSocketAddress = new InetSocketAddress("202.0.0." + i, 10001);
112+
Channel c1 = spy(Channel.class);
113+
ReflectUtils.setFieldValue(c1, "inetSocketAddress", inetSocketAddress);
114+
ReflectUtils.setFieldValue(c1, "inetAddress", inetSocketAddress.getAddress());
115+
ReflectUtils.setFieldValue(c1, "ctx", spy(ChannelHandlerContext.class));
116+
Mockito.doNothing().when(c1).send((byte[]) any());
117+
PeerManager.add(ctx, c1);
118+
}
119+
120+
// Set first minBroadcastPeerSize peers as broadcast-state
121+
List<PeerConnection> peers = PeerManager.getPeers();
122+
for (PeerConnection peer : peers.subList(0, ResilienceService.minBroadcastPeerSize)) {
123+
peer.setNeedSyncFromPeer(false);
124+
peer.setNeedSyncFromUs(false);
125+
peer.setLastInteractiveTime(System.currentTimeMillis() - 1000);
126+
}
127+
for (PeerConnection peer : peers.subList(ResilienceService.minBroadcastPeerSize,
128+
maxConnection + 1)) {
129+
peer.setNeedSyncFromPeer(false);
130+
peer.setNeedSyncFromUs(true);
131+
}
132+
133+
// Give the LAST broadcast peer a very recent blockRcvTime — it must NOT be disconnected
134+
PeerConnection bestPeer = peers.stream()
135+
.filter(p -> !p.isNeedSyncFromUs() && !p.isNeedSyncFromPeer())
136+
.reduce((a, b) -> b) // last broadcast peer
137+
.orElseThrow(() -> new AssertionError("no broadcast peer"));
138+
bestPeer.setBlockRcvTime(System.currentTimeMillis());
139+
140+
InetSocketAddress bestPeerAddress = bestPeer.getChannel().getInetSocketAddress();
141+
142+
// With minBroadcastPeerSize=3 broadcast peers, getRandomDisconnectionPeers returns
143+
// the 1 peer with oldest blockRcvTime (0). bestPeer has most recent time → exempt.
144+
ReflectUtils.invokeMethod(service, "disconnectRandom");
145+
146+
boolean bestPeerStillConnected = PeerManager.getPeers().stream()
147+
.anyMatch(p -> p.getChannel().getInetSocketAddress().equals(bestPeerAddress));
148+
Assert.assertTrue("Peer with most recent blockRcvTime should not be disconnected",
149+
bestPeerStillConnected);
150+
}
151+
100152
@Test
101153
public void testDisconnectLan() {
102154
int minConnection = 8;

0 commit comments

Comments
 (0)