|
1 | 1 | package org.tron.core.capsule; |
2 | 2 |
|
| 3 | +import static org.mockito.Mockito.mock; |
| 4 | +import static org.mockito.Mockito.when; |
| 5 | + |
3 | 6 | import com.google.protobuf.ByteString; |
4 | 7 | import java.io.IOException; |
5 | 8 | import java.util.ArrayList; |
|
21 | 24 | import org.tron.core.config.args.Args; |
22 | 25 | import org.tron.core.exception.BadBlockException; |
23 | 26 | import org.tron.core.exception.BadItemException; |
| 27 | +import org.tron.core.exception.ValidateSignatureException; |
| 28 | +import org.tron.core.store.AccountStore; |
| 29 | +import org.tron.core.store.DynamicPropertiesStore; |
| 30 | +import org.tron.protos.Protocol.Block; |
| 31 | +import org.tron.protos.Protocol.BlockHeader; |
24 | 32 | import org.tron.protos.Protocol.Transaction.Contract.ContractType; |
25 | 33 | import org.tron.protos.contract.BalanceContract.TransferContract; |
26 | 34 |
|
@@ -180,6 +188,95 @@ public void testGetTimeStamp() { |
180 | 188 | Assert.assertEquals(1234L, blockCapsule0.getTimeStamp()); |
181 | 189 | } |
182 | 190 |
|
| 191 | + /** |
| 192 | + * Pin the contract that switchFork's signature recheck relies on: |
| 193 | + * when the recovered signer address does not match the witness address, |
| 194 | + * validateSignature returns false (no exception). switchFork uses the |
| 195 | + * boolean return to decide whether to throw, so this contract is what |
| 196 | + * makes the fix work for "wrong signer" attacks. |
| 197 | + */ |
| 198 | + @Test |
| 199 | + public void testValidateSignatureReturnsFalseWhenSignerMismatch() throws Exception { |
| 200 | + String signerKey = PublicMethod.getRandomPrivateKey(); |
| 201 | + String witnessKey = PublicMethod.getRandomPrivateKey(); |
| 202 | + byte[] witnessAddress = PublicMethod.getAddressByteByPrivateKey(witnessKey); |
| 203 | + |
| 204 | + BlockCapsule block = new BlockCapsule(2, |
| 205 | + Sha256Hash.wrap(ByteString.copyFrom(ByteArray.fromHexString( |
| 206 | + "9938a342238077182498b464ac0292229938a342238077182498b464ac029222"))), |
| 207 | + 4321, |
| 208 | + ByteString.copyFrom(witnessAddress)); |
| 209 | + block.sign(ByteArray.fromHexString(signerKey)); |
| 210 | + |
| 211 | + DynamicPropertiesStore dps = mock(DynamicPropertiesStore.class); |
| 212 | + when(dps.getAllowMultiSign()).thenReturn(0L); |
| 213 | + AccountStore accountStore = mock(AccountStore.class); |
| 214 | + |
| 215 | + Assert.assertFalse(block.validateSignature(dps, accountStore)); |
| 216 | + } |
| 217 | + |
| 218 | + /** |
| 219 | + * Same key path under the happy case: when signer == witness, validateSignature |
| 220 | + * returns true. Guards against any future refactor that accidentally inverts |
| 221 | + * the comparison or strips the witness check. |
| 222 | + */ |
| 223 | + @Test |
| 224 | + public void testValidateSignatureReturnsTrueWhenSignerMatches() throws Exception { |
| 225 | + String key = PublicMethod.getRandomPrivateKey(); |
| 226 | + byte[] witnessAddress = PublicMethod.getAddressByteByPrivateKey(key); |
| 227 | + |
| 228 | + BlockCapsule block = new BlockCapsule(3, |
| 229 | + Sha256Hash.wrap(ByteString.copyFrom(ByteArray.fromHexString( |
| 230 | + "9938a342238077182498b464ac0292229938a342238077182498b464ac029222"))), |
| 231 | + 5678, |
| 232 | + ByteString.copyFrom(witnessAddress)); |
| 233 | + block.sign(ByteArray.fromHexString(key)); |
| 234 | + |
| 235 | + DynamicPropertiesStore dps = mock(DynamicPropertiesStore.class); |
| 236 | + when(dps.getAllowMultiSign()).thenReturn(0L); |
| 237 | + AccountStore accountStore = mock(AccountStore.class); |
| 238 | + |
| 239 | + Assert.assertTrue(block.validateSignature(dps, accountStore)); |
| 240 | + } |
| 241 | + |
| 242 | + /** |
| 243 | + * The other failure mode switchFork must handle: signature bytes are |
| 244 | + * malformed (cannot recover a public key). validateSignature wraps the |
| 245 | + * underlying SignatureException as ValidateSignatureException, which the |
| 246 | + * existing catch block in switchFork already handles. |
| 247 | + */ |
| 248 | + @Test(expected = ValidateSignatureException.class) |
| 249 | + public void testValidateSignatureThrowsForMalformedSignature() throws Exception { |
| 250 | + byte[] witnessAddress = PublicMethod.getAddressByteByPrivateKey( |
| 251 | + PublicMethod.getRandomPrivateKey()); |
| 252 | + |
| 253 | + // 65-byte signature with valid length but garbage content — passes Rsv parsing |
| 254 | + // but fails ECDSA recovery, surfacing SignatureException → ValidateSignatureException. |
| 255 | + byte[] garbageSigBytes = new byte[65]; |
| 256 | + Arrays.fill(garbageSigBytes, (byte) 0xAB); |
| 257 | + ByteString garbageSig = ByteString.copyFrom(garbageSigBytes); |
| 258 | + |
| 259 | + BlockHeader.raw rawData = BlockHeader.raw.newBuilder() |
| 260 | + .setNumber(4) |
| 261 | + .setTimestamp(1111) |
| 262 | + .setParentHash(ByteString.copyFrom(ByteArray.fromHexString( |
| 263 | + "9938a342238077182498b464ac0292229938a342238077182498b464ac029222"))) |
| 264 | + .setWitnessAddress(ByteString.copyFrom(witnessAddress)) |
| 265 | + .build(); |
| 266 | + BlockHeader header = BlockHeader.newBuilder() |
| 267 | + .setRawData(rawData) |
| 268 | + .setWitnessSignature(garbageSig) |
| 269 | + .build(); |
| 270 | + Block proto = Block.newBuilder().setBlockHeader(header).build(); |
| 271 | + BlockCapsule block = new BlockCapsule(proto); |
| 272 | + |
| 273 | + DynamicPropertiesStore dps = mock(DynamicPropertiesStore.class); |
| 274 | + when(dps.getAllowMultiSign()).thenReturn(0L); |
| 275 | + AccountStore accountStore = mock(AccountStore.class); |
| 276 | + |
| 277 | + block.validateSignature(dps, accountStore); |
| 278 | + } |
| 279 | + |
183 | 280 | @Test |
184 | 281 | public void testConcurrentToString() throws InterruptedException { |
185 | 282 | List<Thread> threadList = new ArrayList<>(); |
|
0 commit comments