White flag

src/main/java/com/iota/iri/service/ledger/LedgerService.java

@@ -101,10 +101,11 @@ public interface LedgerService {
      * @param visitedTransactions a set of transaction hashes that shall be considered to be visited already
      * @param startTransaction the transaction that marks the start of the dag traversal and that has its approvees
      *                        examined
+     * @param allowGenesisReference Allows for confirmation of a transaction that references genesis
      * @return a map of the balance changes (addresses associated to their balance) or {@code null} if the balance could
      *         not be generated due to inconsistencies
      * @throws LedgerException if anything unexpected happens while generating the balance changes
      */
-    Map<Hash, Long> generateBalanceDiff(Set<Hash> visitedTransactions, Hash startTransaction, int milestoneIndex)
+    Map<Hash, Long> generateBalanceDiff(Set<Hash> visitedTransactions, Hash startTransaction, int milestoneIndex, boolean allowGenesisReference)
             throws LedgerException;
 }

src/main/java/com/iota/iri/service/ledger/impl/LedgerServiceImpl.java

@@ -132,7 +132,7 @@ public boolean isBalanceDiffConsistent(Set<Hash> approvedHashes, Map<Hash, Long>
         }
         Set<Hash> visitedHashes = new HashSet<>(approvedHashes);
         Map<Hash, Long> currentState = generateBalanceDiff(visitedHashes, tip,
-                snapshotProvider.getLatestSnapshot().getIndex());
+                snapshotProvider.getLatestSnapshot().getIndex(), true);
         if (currentState == null) {
             return false;
         }
@@ -151,7 +151,7 @@ public boolean isBalanceDiffConsistent(Set<Hash> approvedHashes, Map<Hash, Long>
     }
 
     @Override
-    public Map<Hash, Long> generateBalanceDiff(Set<Hash> visitedTransactions, Hash startTransaction, int milestoneIndex)
+    public Map<Hash, Long> generateBalanceDiff(Set<Hash> visitedTransactions, Hash startTransaction, int milestoneIndex, boolean allowGenesisReference)
             throws LedgerException {
 
         Map<Hash, Long> state = new HashMap<>();
@@ -175,6 +175,18 @@ public Map<Hash, Long> generateBalanceDiff(Set<Hash> visitedTransactions, Hash s
 
                 boolean isEmptyTrunk = trunkTransactionViewModel.getType() == TransactionViewModel.PREFILLED_SLOT;
                 boolean isEmptyBranch = branchTransactionViewModel.getType() == TransactionViewModel.PREFILLED_SLOT;
+
+                //Don't confirm non-solid txs.
+                Hash genesisHash = Hash.NULL_HASH;
+                if(isEmptyTrunk  && !allowGenesisReference){
+                    return null;
+                 }
+                //proceed only if empty trunk or branch is genesis
+                if((isEmptyTrunk && !trunkTransactionViewModel.getHash().equals(genesisHash))||
+                (isEmptyBranch && !branchTransactionViewModel.getHash().equals(genesisHash))){
+                    return null;
+                }
+
                 boolean isApprovedTrunk = (trunkTransactionViewModel.snapshotIndex() > 0) && (trunkTransactionViewModel.snapshotIndex() != milestoneIndex);
                 boolean isApprovedBranch = (branchTransactionViewModel.snapshotIndex() > 0) && (branchTransactionViewModel.snapshotIndex() != milestoneIndex);
                 boolean isLeafTrunk = isEmptyTrunk  || visitedTransactions.contains(trunkTransactionViewModel.getHash()) || isApprovedTrunk;
@@ -286,7 +298,7 @@ private boolean generateStateDiff(MilestoneViewModel milestone) throws LedgerExc
                 try {
                     Hash tail = transactionViewModel.getHash();
                     Map<Hash, Long> balanceChanges = generateBalanceDiff(new HashSet<>(), tail,
-                            snapshotProvider.getLatestSnapshot().getIndex());
+                            snapshotProvider.getLatestSnapshot().getIndex(), true);
                     successfullyProcessed = balanceChanges != null;
                     if (successfullyProcessed) {
                         milestoneService.updateMilestoneIndexOfMilestoneTransactions(milestone.getHash(),

src/test/java/com/iota/iri/service/ledger/impl/LedgerServiceImplTest.java

@@ -8,7 +8,10 @@
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 
+import com.iota.iri.model.Hash;
+import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -71,7 +74,19 @@ public void generateBalanceDiffWithPersistsSpentAddresses() throws Exception {
         int milestoneIndex = 1;
         when(milestoneService.isTransactionConfirmed(tailTx, milestoneIndex)).thenReturn(false);
         when(snapshotProvider.getInitialSnapshot().getSolidEntryPoints()).thenReturn(Collections.emptyMap());
-        ledgerService.generateBalanceDiff(new HashSet<>(), tailTx.getHash(), milestoneIndex);
+        ledgerService.generateBalanceDiff(new HashSet<>(), tailTx.getHash(), milestoneIndex, true);
         verify(spentAddressesService, times(1)).persistValidatedSpentAddressesAsync(eq(bundle));
     }
+
+    @Test
+    public void generateBalanceDiffWithGenesisReference() throws Exception {
+        List<TransactionViewModel> bundle = TangleMockUtils.mockValidBundle(tangle, bundleValidator, 1,
+                "A", "Z");
+        TransactionViewModel tailTx = bundle.get(0);
+        int milestoneIndex = 1;
+        when(milestoneService.isTransactionConfirmed(tailTx, milestoneIndex)).thenReturn(false);
+        when(snapshotProvider.getInitialSnapshot().getSolidEntryPoints()).thenReturn(Collections.emptyMap());
+        Map<Hash, Long> diffMap = ledgerService.generateBalanceDiff(new HashSet<>(), tailTx.getHash(), milestoneIndex, false);
+        Assert.assertNull("Diff map should be null because genesis trunk reference is not allowed", diffMap);
+    }
 }