Skip to content

Commit 0d97d4c

Browse files
halibobo1205halibobo1205
committed
feat(net): strip unknown protobuf fields on inbound messages
1 parent 2c50400 commit 0d97d4c

4 files changed

Lines changed: 265 additions & 6 deletions

File tree

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

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,24 @@
11
package org.tron.core.net.message.adv;
22

3+
import com.google.protobuf.UnknownFieldSet;
34
import org.tron.common.overlay.message.Message;
45
import org.tron.common.utils.Sha256Hash;
56
import org.tron.core.capsule.BlockCapsule;
67
import org.tron.core.capsule.BlockCapsule.BlockId;
78
import org.tron.core.capsule.TransactionCapsule;
89
import org.tron.core.net.message.MessageTypes;
910
import org.tron.core.net.message.TronMessage;
11+
import org.tron.protos.Protocol.Block;
1012

1113
public class BlockMessage extends TronMessage {
1214

1315
private BlockCapsule block;
1416

1517
public BlockMessage(byte[] data) throws Exception {
16-
super(data);
1718
this.type = MessageTypes.BLOCK.asByte();
18-
this.block = new BlockCapsule(getCodedInputStream(data));
19+
Block parsed = Block.parseFrom(getCodedInputStream(data));
20+
this.block = new BlockCapsule(sanitize(parsed));
21+
this.data = this.block.getData();
1922
if (Message.isFilter()) {
2023
Message.compareBytes(data, block.getInstance().toByteArray());
2124
TransactionCapsule.validContractProto(block.getInstance().getTransactionsList());
@@ -28,6 +31,25 @@ public BlockMessage(BlockCapsule block) {
2831
this.block = block;
2932
}
3033

34+
private static Block sanitize(Block block) {
35+
boolean blockHasUnknown = !block.getUnknownFields().asMap().isEmpty();
36+
boolean headerHasUnknown = !block.getBlockHeader().getUnknownFields().asMap().isEmpty();
37+
if (!blockHasUnknown && !headerHasUnknown) {
38+
return block;
39+
}
40+
UnknownFieldSet empty = UnknownFieldSet.getDefaultInstance();
41+
Block.Builder builder = block.toBuilder();
42+
if (blockHasUnknown) {
43+
builder.setUnknownFields(empty);
44+
}
45+
if (headerHasUnknown) {
46+
builder.setBlockHeader(block.getBlockHeader().toBuilder()
47+
.setUnknownFields(empty)
48+
.build());
49+
}
50+
return builder.build();
51+
}
52+
3153
public BlockId getBlockId() {
3254
return getBlockCapsule().getBlockId();
3355
}

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

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package org.tron.core.net.message.adv;
22

3+
import com.google.protobuf.UnknownFieldSet;
34
import org.tron.common.overlay.message.Message;
45
import org.tron.common.utils.Sha256Hash;
56
import org.tron.core.capsule.TransactionCapsule;
@@ -12,8 +13,10 @@ public class TransactionMessage extends TronMessage {
1213
private TransactionCapsule transactionCapsule;
1314

1415
public TransactionMessage(byte[] data) throws Exception {
15-
super(data);
16-
this.transactionCapsule = new TransactionCapsule(getCodedInputStream(data));
16+
Transaction parsed = Transaction.parseFrom(getCodedInputStream(data));
17+
Transaction sanitized = sanitize(parsed);
18+
this.transactionCapsule = new TransactionCapsule(sanitized);
19+
this.data = this.transactionCapsule.getData();
1720
this.type = MessageTypes.TRX.asByte();
1821
if (Message.isFilter()) {
1922
compareBytes(data, transactionCapsule.getInstance().toByteArray());
@@ -28,6 +31,15 @@ public TransactionMessage(Transaction trx) {
2831
this.data = trx.toByteArray();
2932
}
3033

34+
static Transaction sanitize(Transaction transaction) {
35+
if (transaction.getUnknownFields().asMap().isEmpty()) {
36+
return transaction;
37+
}
38+
return transaction.toBuilder()
39+
.setUnknownFields(UnknownFieldSet.getDefaultInstance())
40+
.build();
41+
}
42+
3143
@Override
3244
public String toString() {
3345
return new StringBuilder().append(super.toString())

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

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package org.tron.core.net.message.adv;
22

3+
import com.google.protobuf.UnknownFieldSet;
34
import java.util.List;
45
import org.tron.core.capsule.TransactionCapsule;
56
import org.tron.core.net.message.MessageTypes;
@@ -20,15 +21,41 @@ public TransactionsMessage(List<Transaction> trxs) {
2021
}
2122

2223
public TransactionsMessage(byte[] data) throws Exception {
23-
super(data);
2424
this.type = MessageTypes.TRXS.asByte();
25-
this.transactions = Protocol.Transactions.parseFrom(getCodedInputStream(data));
25+
Protocol.Transactions parsed = Protocol.Transactions.parseFrom(getCodedInputStream(data));
26+
this.transactions = sanitize(parsed);
27+
this.data = this.transactions.toByteArray();
2628
if (isFilter()) {
2729
compareBytes(data, transactions.toByteArray());
2830
TransactionCapsule.validContractProto(transactions.getTransactionsList());
2931
}
3032
}
3133

34+
private static Protocol.Transactions sanitize(Protocol.Transactions transactions) {
35+
boolean wrapperHasUnknown = !transactions.getUnknownFields().asMap().isEmpty();
36+
boolean anyTxHasUnknown = false;
37+
for (Transaction tx : transactions.getTransactionsList()) {
38+
if (!tx.getUnknownFields().asMap().isEmpty()) {
39+
anyTxHasUnknown = true;
40+
break;
41+
}
42+
}
43+
if (!wrapperHasUnknown && !anyTxHasUnknown) {
44+
return transactions;
45+
}
46+
Protocol.Transactions.Builder builder = transactions.toBuilder();
47+
if (wrapperHasUnknown) {
48+
builder.setUnknownFields(UnknownFieldSet.getDefaultInstance());
49+
}
50+
if (anyTxHasUnknown) {
51+
builder.clearTransactions();
52+
for (Transaction tx : transactions.getTransactionsList()) {
53+
builder.addTransactions(TransactionMessage.sanitize(tx));
54+
}
55+
}
56+
return builder.build();
57+
}
58+
3259
public Protocol.Transactions getTransactions() {
3360
return transactions;
3461
}

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

Lines changed: 198 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,198 @@
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.assertTrue;
6+
7+
import com.google.protobuf.ByteString;
8+
import com.google.protobuf.UnknownFieldSet;
9+
import org.junit.BeforeClass;
10+
import org.junit.Test;
11+
import org.mockito.Mockito;
12+
import org.tron.common.overlay.message.Message;
13+
import org.tron.core.store.DynamicPropertiesStore;
14+
import org.tron.protos.Protocol.Block;
15+
import org.tron.protos.Protocol.BlockHeader;
16+
import org.tron.protos.Protocol.Transaction;
17+
import org.tron.protos.Protocol.Transactions;
18+
19+
/**
20+
* Verifies that the three P2P message constructors that take raw bytes
21+
* (BlockMessage, TransactionMessage, TransactionsMessage) strip unknown
22+
* protobuf fields at positions not covered by any consensus hash, while
23+
* leaving the signed regions byte-identical so transaction id and block hash
24+
* remain stable.
25+
*/
26+
public class SanitizeUnknownFieldsTest {
27+
28+
/**
29+
* 1 KiB length-delimited unknown field at tag 99999, large enough that any
30+
* sanitize-stripped message is noticeably shorter than the padded input.
31+
*/
32+
private static final UnknownFieldSet PADDING = UnknownFieldSet.newBuilder()
33+
.addField(99999, UnknownFieldSet.Field.newBuilder()
34+
.addLengthDelimited(ByteString.copyFrom(new byte[1024]))
35+
.build())
36+
.build();
37+
38+
@BeforeClass
39+
public static void setUp() {
40+
// Mock dynamicPropertiesStore so Message.isFilter() returns false
41+
// (mock's primitive-long getter returns 0L by default).
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+
// ---- BlockMessage ----
65+
66+
@Test
67+
public void testBlockMessageStripsBlockLevelUnknownFields() throws Exception {
68+
Block padded = sampleBlock().toBuilder().setUnknownFields(PADDING).build();
69+
byte[] paddedBytes = padded.toByteArray();
70+
71+
BlockMessage msg = new BlockMessage(paddedBytes);
72+
73+
assertTrue("Block-level unknown fields should be stripped",
74+
msg.getBlockCapsule().getInstance().getUnknownFields().asMap().isEmpty());
75+
assertTrue("msg.data should be canonical (no padding)",
76+
msg.getData().length < paddedBytes.length);
77+
}
78+
79+
@Test
80+
public void testBlockMessageStripsBlockHeaderOuterUnknownFields() throws Exception {
81+
BlockHeader paddedHeader = BlockHeader.newBuilder()
82+
.setRawData(sampleRawHeader())
83+
.setUnknownFields(PADDING)
84+
.build();
85+
Block padded = Block.newBuilder().setBlockHeader(paddedHeader).build();
86+
byte[] paddedBytes = padded.toByteArray();
87+
88+
BlockMessage msg = new BlockMessage(paddedBytes);
89+
90+
assertTrue("BlockHeader outer unknown fields should be stripped",
91+
msg.getBlockCapsule().getInstance().getBlockHeader()
92+
.getUnknownFields().asMap().isEmpty());
93+
assertTrue(msg.getData().length < paddedBytes.length);
94+
}
95+
96+
@Test
97+
public void testBlockMessagePreservesBlockHeaderRawData() throws Exception {
98+
Block clean = sampleBlock();
99+
Block padded = clean.toBuilder().setUnknownFields(PADDING).build();
100+
101+
BlockMessage msg = new BlockMessage(padded.toByteArray());
102+
103+
assertEquals("BlockHeader.raw_data must be byte-identical so block hash matches",
104+
clean.getBlockHeader().getRawData(),
105+
msg.getBlockCapsule().getInstance().getBlockHeader().getRawData());
106+
}
107+
108+
@Test
109+
public void testBlockMessageCleanBlockPassesThroughUnchanged() throws Exception {
110+
Block clean = sampleBlock();
111+
byte[] cleanBytes = clean.toByteArray();
112+
113+
BlockMessage msg = new BlockMessage(cleanBytes);
114+
115+
assertArrayEquals("Clean block bytes should pass through unchanged",
116+
cleanBytes, msg.getData());
117+
}
118+
119+
// ---- TransactionMessage ----
120+
121+
@Test
122+
public void testTransactionMessageStripsTopLevelUnknownFields() throws Exception {
123+
Transaction padded = sampleTransaction().toBuilder().setUnknownFields(PADDING).build();
124+
byte[] paddedBytes = padded.toByteArray();
125+
126+
TransactionMessage msg = new TransactionMessage(paddedBytes);
127+
128+
assertTrue("Transaction-level unknown fields should be stripped",
129+
msg.getTransactionCapsule().getInstance().getUnknownFields().asMap().isEmpty());
130+
assertTrue(msg.getData().length < paddedBytes.length);
131+
}
132+
133+
@Test
134+
public void testTransactionMessagePreservesTransactionId() throws Exception {
135+
Transaction clean = sampleTransaction();
136+
Transaction padded = clean.toBuilder().setUnknownFields(PADDING).build();
137+
138+
TransactionMessage cleanMsg = new TransactionMessage(clean.toByteArray());
139+
TransactionMessage paddedMsg = new TransactionMessage(padded.toByteArray());
140+
141+
assertEquals("Padding outside raw_data must not change tx id",
142+
cleanMsg.getTransactionCapsule().getTransactionId(),
143+
paddedMsg.getTransactionCapsule().getTransactionId());
144+
}
145+
146+
@Test
147+
public void testTransactionMessageCleanTransactionPassesThroughUnchanged() throws Exception {
148+
Transaction clean = sampleTransaction();
149+
byte[] cleanBytes = clean.toByteArray();
150+
151+
TransactionMessage msg = new TransactionMessage(cleanBytes);
152+
153+
assertArrayEquals(cleanBytes, msg.getData());
154+
}
155+
156+
// ---- TransactionsMessage ----
157+
158+
@Test
159+
public void testTransactionsMessageStripsWrapperUnknownFields() throws Exception {
160+
Transactions padded = Transactions.newBuilder()
161+
.addTransactions(sampleTransaction())
162+
.setUnknownFields(PADDING)
163+
.build();
164+
byte[] paddedBytes = padded.toByteArray();
165+
166+
TransactionsMessage msg = new TransactionsMessage(paddedBytes);
167+
168+
assertTrue("Wrapper unknown fields should be stripped",
169+
msg.getTransactions().getUnknownFields().asMap().isEmpty());
170+
assertTrue(msg.getData().length < paddedBytes.length);
171+
}
172+
173+
@Test
174+
public void testTransactionsMessageStripsNestedTxUnknownFields() throws Exception {
175+
Transaction paddedTx = sampleTransaction().toBuilder().setUnknownFields(PADDING).build();
176+
Transactions wrapper = Transactions.newBuilder().addTransactions(paddedTx).build();
177+
byte[] paddedBytes = wrapper.toByteArray();
178+
179+
TransactionsMessage msg = new TransactionsMessage(paddedBytes);
180+
181+
assertTrue("Nested tx unknown fields should be stripped",
182+
msg.getTransactions().getTransactions(0)
183+
.getUnknownFields().asMap().isEmpty());
184+
assertTrue(msg.getData().length < paddedBytes.length);
185+
}
186+
187+
@Test
188+
public void testTransactionsMessageCleanWrapperPassesThroughUnchanged() throws Exception {
189+
Transactions clean = Transactions.newBuilder()
190+
.addTransactions(sampleTransaction())
191+
.build();
192+
byte[] cleanBytes = clean.toByteArray();
193+
194+
TransactionsMessage msg = new TransactionsMessage(cleanBytes);
195+
196+
assertArrayEquals(cleanBytes, msg.getData());
197+
}
198+
}

0 commit comments

Comments
 (0)