Skip to content

Commit 0c741e3

Browse files
halibobo1205halibobo1205
authored
feat(net): normalize inbound messages (#6797)
* feat(net): normalize inbound messages * perf(net): skip wire-byte rewrite when sanitize is a no-op
1 parent af88269 commit 0c741e3

7 files changed

Lines changed: 244 additions & 1 deletion

File tree

‎chainbase/src/main/java/org/tron/core/capsule/BlockCapsule.java‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import com.google.protobuf.ByteString;
2222
import com.google.protobuf.CodedInputStream;
2323
import com.google.protobuf.InvalidProtocolBufferException;
24+
import com.google.protobuf.UnknownFieldSet;
2425
import java.security.SignatureException;
2526
import java.util.ArrayList;
2627
import java.util.Arrays;
@@ -328,6 +329,26 @@ public boolean hasWitnessSignature() {
328329
return !getInstance().getBlockHeader().getWitnessSignature().isEmpty();
329330
}
330331

332+
public boolean sanitize() {
333+
boolean blockHasUnknown = !this.block.getUnknownFields().asMap().isEmpty();
334+
boolean headerHasUnknown = !this.block.getBlockHeader().getUnknownFields().asMap().isEmpty();
335+
if (!blockHasUnknown && !headerHasUnknown) {
336+
return false;
337+
}
338+
UnknownFieldSet empty = UnknownFieldSet.getDefaultInstance();
339+
Block.Builder builder = this.block.toBuilder();
340+
if (blockHasUnknown) {
341+
builder.setUnknownFields(empty);
342+
}
343+
if (headerHasUnknown) {
344+
builder.setBlockHeader(this.block.getBlockHeader().toBuilder()
345+
.setUnknownFields(empty)
346+
.build());
347+
}
348+
this.block = builder.build();
349+
return true;
350+
}
351+
331352
@Override
332353
public String toString() {
333354
StringBuilder toStringBuff = new StringBuilder();

‎chainbase/src/main/java/org/tron/core/capsule/TransactionCapsule.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import com.google.protobuf.GeneratedMessageV3;
2929
import com.google.protobuf.Internal;
3030
import com.google.protobuf.InvalidProtocolBufferException;
31+
import com.google.protobuf.UnknownFieldSet;
3132
import java.io.IOException;
3233
import java.security.SignatureException;
3334
import java.util.ArrayList;
@@ -494,6 +495,16 @@ public static boolean validateSignature(Transaction transaction,
494495
return false;
495496
}
496497

498+
public boolean sanitize() {
499+
if (this.transaction.getUnknownFields().asMap().isEmpty()) {
500+
return false;
501+
}
502+
this.transaction = this.transaction.toBuilder()
503+
.setUnknownFields(UnknownFieldSet.getDefaultInstance())
504+
.build();
505+
return true;
506+
}
507+
497508
public void resetResult() {
498509
if (this.getInstance().getRetCount() > 0) {
499510
this.transaction = this.getInstance().toBuilder().clearRet().build();

‎framework/src/main/java/org/tron/core/Wallet.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -556,9 +556,9 @@ public GrpcAPI.Return broadcastTransaction(Transaction signedTransaction) {
556556
if (trx.getInstance().getRawData().getContractCount() == 0) {
557557
throw new ContractValidateException(ActuatorConstant.CONTRACT_NOT_EXIST);
558558
}
559-
TransactionMessage message = new TransactionMessage(trx.getInstance().toByteArray());
560559
trx.checkExpiration(chainBaseManager.getNextBlockSlotTime());
561560
dbManager.pushTransaction(trx);
561+
TransactionMessage message = new TransactionMessage(trx.getInstance().toByteArray());
562562
int num = tronNetService.fastBroadcastTransaction(message);
563563
if (num == 0 && minEffectiveConnection != 0) {
564564
return builder.setResult(false).setCode(response_code.NOT_ENOUGH_EFFECTIVE_CONNECTION)

‎framework/src/main/java/org/tron/core/db/Manager.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1534,6 +1534,9 @@ public TransactionInfo processTransaction(final TransactionCapsule trxCap, Block
15341534
String.format(" %s transaction signature validate failed", txId));
15351535
}
15361536

1537+
if (!trxCap.isInBlock()) {
1538+
trxCap.sanitize();
1539+
}
15371540
TransactionTrace trace = new TransactionTrace(trxCap, StoreFactory.getInstance(),
15381541
new RuntimeImpl());
15391542
trxCap.setTrxTrace(trace);

‎framework/src/main/java/org/tron/core/net/message/adv/BlockMessage.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,12 @@ public BlockMessage(BlockCapsule block) {
2828
this.block = block;
2929
}
3030

31+
public void sanitize() {
32+
if (this.block.sanitize()) {
33+
this.data = this.block.getData();
34+
}
35+
}
36+
3137
public BlockId getBlockId() {
3238
return getBlockCapsule().getBlockId();
3339
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,8 @@ public void processMessage(PeerConnection peer, TronMessage msg) throws P2pExcep
7777
check(peer, blockMessage);
7878
}
7979

80+
blockMessage.sanitize();
81+
8082
if (peer.getSyncBlockRequested().containsKey(blockId)) {
8183
peer.getSyncBlockRequested().remove(blockId);
8284
peer.getSyncBlockInProcess().add(blockId);

‎framework/src/test/java/org/tron/core/net/message/adv/SanitizeUnknownFieldsTest.java‎

Lines changed: 200 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,200 @@
1+
package org.tron.core.net.message.adv;
2+
3+
import static org.junit.Assert.assertArrayEquals;
4+
import static org.junit.Assert.assertEquals;
5+
import static org.junit.Assert.assertFalse;
6+
import static org.junit.Assert.assertNotEquals;
7+
import static org.junit.Assert.assertSame;
8+
import static org.junit.Assert.assertTrue;
9+
10+
import com.google.protobuf.ByteString;
11+
import com.google.protobuf.UnknownFieldSet;
12+
import org.junit.BeforeClass;
13+
import org.junit.Test;
14+
import org.mockito.Mockito;
15+
import org.tron.common.overlay.message.Message;
16+
import org.tron.core.capsule.BlockCapsule;
17+
import org.tron.core.capsule.TransactionCapsule;
18+
import org.tron.core.store.DynamicPropertiesStore;
19+
import org.tron.protos.Protocol.Block;
20+
import org.tron.protos.Protocol.BlockHeader;
21+
import org.tron.protos.Protocol.Transaction;
22+
23+
/**
24+
* Verifies the {@code sanitize()} helpers on {@link BlockCapsule},
25+
* {@link TransactionCapsule} and {@link BlockMessage}: they strip outer
26+
* unknown protobuf fields while leaving every consensus-hashed / signed
27+
* region byte-identical.
28+
*/
29+
public class SanitizeUnknownFieldsTest {
30+
31+
private static final UnknownFieldSet PADDING = UnknownFieldSet.newBuilder()
32+
.addField(99999, UnknownFieldSet.Field.newBuilder()
33+
.addLengthDelimited(ByteString.copyFrom(new byte[1024]))
34+
.build())
35+
.build();
36+
37+
@BeforeClass
38+
public static void setUp() {
39+
// BlockMessage(byte[]) calls Message.isFilter() which dereferences the
40+
// static DynamicPropertiesStore. The mock's primitive-long getter returns
41+
// 0L by default, so isFilter() returns false.
42+
Message.setDynamicPropertiesStore(Mockito.mock(DynamicPropertiesStore.class));
43+
}
44+
45+
private static BlockHeader.raw sampleRawHeader() {
46+
return BlockHeader.raw.newBuilder()
47+
.setNumber(100)
48+
.setTimestamp(123456789L)
49+
.build();
50+
}
51+
52+
private static Block sampleBlock() {
53+
return Block.newBuilder()
54+
.setBlockHeader(BlockHeader.newBuilder().setRawData(sampleRawHeader()).build())
55+
.build();
56+
}
57+
58+
private static Transaction sampleTransaction() {
59+
return Transaction.newBuilder()
60+
.setRawData(Transaction.raw.newBuilder().setTimestamp(123456789L).build())
61+
.build();
62+
}
63+
64+
// ---- BlockCapsule.sanitize ----
65+
66+
@Test
67+
public void blockCapsuleSanitizeStripsBlockLevelUnknownFields() {
68+
Block padded = sampleBlock().toBuilder().setUnknownFields(PADDING).build();
69+
BlockCapsule capsule = new BlockCapsule(padded);
70+
long originalSize = capsule.getData().length;
71+
72+
assertTrue("sanitize() should report it mutated the capsule", capsule.sanitize());
73+
74+
assertTrue("Block-level unknown fields should be stripped",
75+
capsule.getInstance().getUnknownFields().asMap().isEmpty());
76+
assertTrue("Sanitized capsule bytes should shrink",
77+
capsule.getData().length < originalSize);
78+
}
79+
80+
@Test
81+
public void blockCapsuleSanitizeStripsBlockHeaderOuterUnknownFields() {
82+
BlockHeader paddedHeader = BlockHeader.newBuilder()
83+
.setRawData(sampleRawHeader())
84+
.setUnknownFields(PADDING)
85+
.build();
86+
Block padded = Block.newBuilder().setBlockHeader(paddedHeader).build();
87+
BlockCapsule capsule = new BlockCapsule(padded);
88+
long originalSize = capsule.getData().length;
89+
90+
assertTrue("sanitize() should report it mutated the capsule", capsule.sanitize());
91+
92+
assertTrue("BlockHeader outer unknown fields should be stripped",
93+
capsule.getInstance().getBlockHeader().getUnknownFields().asMap().isEmpty());
94+
assertTrue(capsule.getData().length < originalSize);
95+
}
96+
97+
@Test
98+
public void blockCapsuleSanitizePreservesBlockHeaderRawData() {
99+
Block clean = sampleBlock();
100+
Block padded = clean.toBuilder().setUnknownFields(PADDING).build();
101+
BlockCapsule capsule = new BlockCapsule(padded);
102+
103+
capsule.sanitize();
104+
105+
assertEquals("BlockHeader.raw_data must be byte-identical so block hash matches",
106+
clean.getBlockHeader().getRawData(),
107+
capsule.getInstance().getBlockHeader().getRawData());
108+
}
109+
110+
@Test
111+
public void blockCapsuleSanitizeIsNoOpOnCleanBlock() {
112+
Block clean = sampleBlock();
113+
BlockCapsule capsule = new BlockCapsule(clean);
114+
Block beforeInstance = capsule.getInstance();
115+
byte[] beforeData = capsule.getData();
116+
117+
assertFalse("sanitize() should report no-op on a clean block", capsule.sanitize());
118+
119+
assertSame("Underlying Block reference should not be rebuilt",
120+
beforeInstance, capsule.getInstance());
121+
assertArrayEquals("Clean block should pass through unchanged", beforeData, capsule.getData());
122+
}
123+
124+
// ---- TransactionCapsule.sanitize ----
125+
126+
@Test
127+
public void transactionCapsuleSanitizeStripsTopLevelUnknownFields() {
128+
Transaction padded = sampleTransaction().toBuilder().setUnknownFields(PADDING).build();
129+
TransactionCapsule capsule = new TransactionCapsule(padded);
130+
long originalSize = capsule.getData().length;
131+
132+
assertTrue("sanitize() should report it mutated the capsule", capsule.sanitize());
133+
134+
assertTrue("Transaction-level unknown fields should be stripped",
135+
capsule.getInstance().getUnknownFields().asMap().isEmpty());
136+
assertTrue(capsule.getData().length < originalSize);
137+
}
138+
139+
@Test
140+
public void transactionCapsuleSanitizePreservesTransactionId() {
141+
Transaction clean = sampleTransaction();
142+
Transaction padded = clean.toBuilder().setUnknownFields(PADDING).build();
143+
TransactionCapsule cleanCapsule = new TransactionCapsule(clean);
144+
TransactionCapsule paddedCapsule = new TransactionCapsule(padded);
145+
146+
paddedCapsule.sanitize();
147+
148+
assertEquals("Padding outside raw_data must not change the transaction id",
149+
cleanCapsule.getTransactionId(),
150+
paddedCapsule.getTransactionId());
151+
}
152+
153+
@Test
154+
public void transactionCapsuleSanitizeIsNoOpOnCleanTransaction() {
155+
Transaction clean = sampleTransaction();
156+
TransactionCapsule capsule = new TransactionCapsule(clean);
157+
Transaction beforeInstance = capsule.getInstance();
158+
byte[] beforeData = capsule.getData();
159+
160+
assertFalse("sanitize() should report no-op on a clean transaction", capsule.sanitize());
161+
162+
assertSame("Underlying Transaction reference should not be rebuilt",
163+
beforeInstance, capsule.getInstance());
164+
assertArrayEquals(beforeData, capsule.getData());
165+
}
166+
167+
// ---- BlockMessage.sanitize ----
168+
169+
@Test
170+
public void blockMessageSanitizeUpdatesBothCapsuleAndWireBytes() throws Exception {
171+
Block padded = sampleBlock().toBuilder().setUnknownFields(PADDING).build();
172+
byte[] paddedBytes = padded.toByteArray();
173+
BlockMessage msg = new BlockMessage(paddedBytes);
174+
assertArrayEquals("Constructor should not sanitize on its own",
175+
paddedBytes, msg.getData());
176+
177+
msg.sanitize();
178+
179+
assertTrue("BlockCapsule should be sanitized",
180+
msg.getBlockCapsule().getInstance().getUnknownFields().asMap().isEmpty());
181+
assertTrue("msg.data should also be rewritten to canonical bytes",
182+
msg.getData().length < paddedBytes.length);
183+
assertArrayEquals("msg.data should equal capsule.getData() after sanitize",
184+
msg.getBlockCapsule().getData(), msg.getData());
185+
assertNotEquals("msg.data should no longer match the padded wire bytes",
186+
paddedBytes.length, msg.getData().length);
187+
}
188+
189+
@Test
190+
public void blockMessageSanitizeSkipsDataRewriteOnCleanBlock() throws Exception {
191+
byte[] cleanBytes = sampleBlock().toByteArray();
192+
BlockMessage msg = new BlockMessage(cleanBytes);
193+
byte[] before = msg.getData();
194+
195+
msg.sanitize();
196+
197+
assertSame("msg.data should not be rewritten on the no-op path",
198+
before, msg.getData());
199+
}
200+
}

0 commit comments

Comments
 (0)