Merge pull request #6501 from yanghang8612/optimize_selfdestruct_restriction

func(vm): optimize selfdestruct restriction

Author: lvs0075

actuator/src/main/java/org/tron/core/vm/PrecompiledContracts.java

@@ -16,6 +16,7 @@
 import static org.tron.common.utils.ByteUtil.parseWord;
 import static org.tron.common.utils.ByteUtil.stripLeadingZeroes;
 import static org.tron.core.config.Parameter.ChainConstant.TRX_PRECISION;
+import static org.tron.core.vm.VMConstant.SIG_LENGTH;
 
 import com.google.protobuf.ByteString;
 
@@ -202,7 +203,7 @@ public class PrecompiledContracts {
   public static PrecompiledContract getOptimizedContractForConstant(PrecompiledContract contract) {
     try {
       Constructor<?> constructor = contract.getClass().getDeclaredConstructor();
-      return  (PrecompiledContracts.PrecompiledContract) constructor.newInstance();
+      return (PrecompiledContracts.PrecompiledContract) constructor.newInstance();
     } catch (Exception e) {
       throw new RuntimeException(e);
     }
@@ -395,6 +396,20 @@ private static byte[][] extractBytesArray(DataWord[] words, int offset, byte[] d
     return bytesArray;
   }
 
+  private static byte[][] extractSigArray(DataWord[] words, int offset, byte[] data) {
+    if (offset > words.length - 1) {
+      return new byte[0][];
+    }
+    int len = words[offset].intValueSafe();
+    byte[][] bytesArray = new byte[len][];
+    for (int i = 0; i < len; i++) {
+      int bytesOffset = words[offset + i + 1].intValueSafe() / WORD_SIZE;
+      bytesArray[i] = extractBytes(data, (bytesOffset + offset + 2) * WORD_SIZE,
+          SIG_LENGTH);
+    }
+    return bytesArray;
+  }
+
   private static byte[] extractBytes(byte[] data, int offset, int len) {
     return Arrays.copyOfRange(data, offset, offset + len);
   }
@@ -936,8 +951,15 @@ public Pair<Boolean, byte[]> execute(byte[] rawData) {
       byte[] hash = Sha256Hash.hash(CommonParameter
           .getInstance().isECKeyCryptoEngine(), combine);
 
-      byte[][] signatures = extractBytesArray(
-          words, words[3].intValueSafe() / WORD_SIZE, rawData);
+      if (VMConfig.allowTvmSelfdestructRestriction()) {
+        int sigArraySize = words[words[3].intValueSafe() / WORD_SIZE].intValueSafe();
+        if (sigArraySize > MAX_SIZE) {
+          return Pair.of(true, DATA_FALSE);
+        }
+      }
+      byte[][] signatures = VMConfig.allowTvmSelfdestructRestriction() ?
+          extractSigArray(words, words[3].intValueSafe() / WORD_SIZE, rawData) :
+          extractBytesArray(words, words[3].intValueSafe() / WORD_SIZE, rawData);
 
       if (signatures.length == 0 || signatures.length > MAX_SIZE) {
         return Pair.of(true, DATA_FALSE);
@@ -1021,8 +1043,18 @@ private Pair<Boolean, byte[]> doExecute(byte[] data)
         throws InterruptedException, ExecutionException {
       DataWord[] words = DataWord.parseArray(data);
       byte[] hash = words[0].getData();
-      byte[][] signatures = extractBytesArray(
-          words, words[1].intValueSafe() / WORD_SIZE, data);
+
+      if (VMConfig.allowTvmSelfdestructRestriction()) {
+        int sigArraySize = words[words[1].intValueSafe() / WORD_SIZE].intValueSafe();
+        int addrArraySize = words[words[2].intValueSafe() / WORD_SIZE].intValueSafe();
+        if (sigArraySize > MAX_SIZE || addrArraySize > MAX_SIZE) {
+          return Pair.of(true, DATA_FALSE);
+        }
+      }
+
+      byte[][] signatures = VMConfig.allowTvmSelfdestructRestriction() ?
+          extractSigArray(words, words[1].intValueSafe() / WORD_SIZE, data) :
+          extractBytesArray(words, words[1].intValueSafe() / WORD_SIZE, data);
       byte[][] addresses = extractBytes32Array(
           words, words[2].intValueSafe() / WORD_SIZE);
       int cnt = signatures.length;

actuator/src/main/java/org/tron/core/vm/VMConstant.java

@@ -4,6 +4,7 @@ public class VMConstant {
 
   public static final int CONTRACT_NAME_LENGTH = 32;
   public static final int MIN_TOKEN_ID = 1_000_000;
+  public static final int SIG_LENGTH = 65;
 
   // Numbers
   public static final int ONE_HUNDRED = 100;

actuator/src/main/java/org/tron/core/vm/program/Memory.java

@@ -98,7 +98,7 @@ public void write(int address, byte[] data, int dataSize, boolean limited) {
 
   public void extendAndWrite(int address, int allocSize, byte[] data) {
     extend(address, allocSize);
-    write(address, data, data.length, false);
+    write(address, data, allocSize, false);
   }
 
   public void extend(int address, int size) {

actuator/src/main/java/org/tron/core/vm/program/Program.java

@@ -1746,7 +1746,11 @@ public void callToPrecompiledAddress(MessageCall msg,
         }
       }
 
-      this.memorySave(msg.getOutDataOffs().intValue(), out.getRight());
+      if (VMConfig.allowTvmSelfdestructRestriction()) {
+        this.memorySave(msg.getOutDataOffs().intValueSafe(), msg.getOutDataSize().intValueSafe(), out.getRight());
+      } else {
+        this.memorySave(msg.getOutDataOffs().intValue(), out.getRight());
+      }
     }
   }
 

framework/src/test/java/org/tron/common/runtime/vm/BatchValidateSignContractTest.java

@@ -15,6 +15,7 @@
 import org.tron.core.db.TransactionTrace;
 import org.tron.core.vm.PrecompiledContracts;
 import org.tron.core.vm.PrecompiledContracts.BatchValidateSign;
+import org.tron.core.vm.config.VMConfig;
 
 
 @Slf4j
@@ -74,6 +75,13 @@ public void staticCallTest() {
     ret = validateMultiSign(hash, signatures, addresses);
     Assert.assertEquals(ret.getValue().length, 32);
     Assert.assertArrayEquals(ret.getValue(), new byte[32]);
+
+    //after optimized
+    VMConfig.initAllowTvmSelfdestructRestriction(1);
+    ret = validateMultiSign(hash, signatures, addresses);
+    Assert.assertEquals(ret.getValue().length, 32);
+    Assert.assertArrayEquals(ret.getValue(), new byte[32]);
+    VMConfig.initAllowTvmSelfdestructRestriction(0);
     System.gc(); // force triggering full gc to avoid timeout for next test
   }
 

framework/src/test/java/org/tron/common/runtime/vm/OperationsTest.java

@@ -16,22 +16,26 @@
 import org.junit.Ignore;
 import org.junit.Test;
 import org.mockito.Mockito;
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.util.StringUtils;
 import org.tron.common.BaseTest;
 import org.tron.common.parameter.CommonParameter;
 import org.tron.common.runtime.InternalTransaction;
 import org.tron.common.utils.DecodeUtil;
 import org.tron.core.Constant;
+import org.tron.core.Wallet;
 import org.tron.core.capsule.AccountCapsule;
 import org.tron.core.config.args.Args;
 import org.tron.core.exception.ContractValidateException;
 import org.tron.core.store.StoreFactory;
 import org.tron.core.vm.EnergyCost;
 import org.tron.core.vm.JumpTable;
+import org.tron.core.vm.MessageCall;
 import org.tron.core.vm.Op;
 import org.tron.core.vm.Operation;
 import org.tron.core.vm.OperationActions;
 import org.tron.core.vm.OperationRegistry;
+import org.tron.core.vm.PrecompiledContracts;
 import org.tron.core.vm.VM;
 import org.tron.core.vm.config.ConfigLoader;
 import org.tron.core.vm.config.VMConfig;
@@ -46,6 +50,8 @@ public class OperationsTest extends BaseTest {
   private ProgramInvokeMockImpl invoke;
   private Program program;
   private final JumpTable jumpTable = OperationRegistry.getTable();
+  @Autowired
+  private Wallet wallet;
 
   @BeforeClass
   public static void init() {
@@ -755,6 +761,30 @@ public void testPushDupSwapAndLogOperations() throws ContractValidateException {
     Assert.assertEquals(2158, program.getResult().getEnergyUsed());
   }
 
+  @Test
+  public void testCallOperations() throws ContractValidateException {
+    invoke = new ProgramInvokeMockImpl();
+    Protocol.Transaction trx = Protocol.Transaction.getDefaultInstance();
+    InternalTransaction interTrx =
+        new InternalTransaction(trx, InternalTransaction.TrxType.TRX_UNKNOWN_TYPE);
+
+    byte prePrefixByte = DecodeUtil.addressPreFixByte;
+    DecodeUtil.addressPreFixByte = Constant.ADD_PRE_FIX_BYTE_MAINNET;
+    VMConfig.initAllowTvmSelfdestructRestriction(1);
+
+    program = new Program(new byte[0], new byte[0], invoke, interTrx);
+    MessageCall messageCall = new MessageCall(
+        Op.CALL, new DataWord(10000),
+        DataWord.ZERO(), DataWord.ZERO(),
+        DataWord.ZERO(), DataWord.ZERO(),
+        DataWord.ZERO(), DataWord.ZERO(),
+        DataWord.ZERO(), false);
+    program.callToPrecompiledAddress(messageCall, new PrecompiledContracts.ECRecover());
+
+    DecodeUtil.addressPreFixByte = prePrefixByte;
+    VMConfig.initAllowTvmSelfdestructRestriction(0);
+  }
+
   @Test
   public void testOtherOperations() throws ContractValidateException {
     invoke = new ProgramInvokeMockImpl();

framework/src/test/java/org/tron/common/runtime/vm/ValidateMultiSignContractTest.java

@@ -25,6 +25,7 @@
 import org.tron.core.config.args.Args;
 import org.tron.core.store.StoreFactory;
 import org.tron.core.vm.PrecompiledContracts.ValidateMultiSign;
+import org.tron.core.vm.config.VMConfig;
 import org.tron.core.vm.repository.Repository;
 import org.tron.core.vm.repository.RepositoryImpl;
 import org.tron.protos.Protocol;
@@ -123,6 +124,13 @@ public void testDifferentCase() {
         validateMultiSign(StringUtil.encode58Check(key.getAddress()), permissionId, data, signs)
             .getValue(), DataWord.ONE().getData());
 
+    //after optimized
+    VMConfig.initAllowTvmSelfdestructRestriction(1);
+    Assert.assertArrayEquals(
+        validateMultiSign(StringUtil.encode58Check(key.getAddress()), permissionId, data, signs)
+            .getValue(), DataWord.ONE().getData());
+    VMConfig.initAllowTvmSelfdestructRestriction(0);
+
     //weight not enough
     signs = new ArrayList<>();
     signs.add(Hex.toHexString(key1.sign(toSign).toByteArray()));