content: /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyten ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *    http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  * KIND, either express or implied.  See the License for the
  * specific language governing permissions and limitations
  * under the License.
  */
  package org.apache.cassandra.utils;
  
  import java.math.BigInteger;
  import java.util.*;
  
  import com.google.common.collect.AbstractIterator;
  
  import org.junit.Before;
  import org.junit.BeforeClass;
  import org.junit.Test;
  
  import org.apache.cassandra.config.DatabaseDescriptor;
  import org.apache.cassandra.dht.*;
  import org.apache.cassandra.dht.RandomPartitioner.BigIntegerToken;
  import org.apache.cassandra.io.util.DataInputBuffer;
  import org.apache.cassandra.io.util.DataOutputBuffer;
  import org.apache.cassandra.net.MessagingService;
  import org.apache.cassandra.service.StorageService;
  import org.apache.cassandra.utils.MerkleTree.RowHash;
  import org.apache.cassandra.utils.MerkleTree.TreeRange;
  import org.apache.cassandra.utils.MerkleTrees.TreeRangeIterator;
  
  import static org.apache.cassandra.utils.MerkleTreeTest.digest;
  import static org.junit.Assert.*;
  
  public class MerkleTreesTest
  {
      private static final byte[] DUMMY = digest("dummy");
  
      /**
       * If a test assumes that the tree is 8 units wide, then it should set this value
       * to 8.
       */
      public static BigInteger TOKEN_SCALE = new BigInteger("8");
  
      protected static final IPartitioner partitioner = RandomPartitioner.instance;
      protected MerkleTrees mts;
  
      private Range<Token> fullRange()
      {
          return new Range<>(partitioner.getMinimumToken(), partitioner.getMinimumToken());
      }
  
      @BeforeClass
      public static void setUp()
      {
          DatabaseDescriptor.daemonInitialization();
          StorageService.instance.setPartitionerUnsafe(partitioner);
      }
      @Before
      public void clear()
      {
          TOKEN_SCALE = new BigInteger("8");
          mts = new MerkleTrees(partitioner);
          mts.addMerkleTree(Integer.MAX_VALUE, fullRange());
      }
  
      public static void assertHashEquals(final byte[] left, final byte[] right)
      {
          assertHashEquals("", left, right);
      }
  
      public static void assertHashEquals(String message, final byte[] left, final byte[] right)
      {
          String lstring = left == null ? "null" : Hex.bytesToHex(left);
          String rstring = right == null ? "null" : Hex.bytesToHex(right);
          assertEquals(message, lstring, rstring);
      }
  
      /**
       * The value returned by this method is affected by TOKEN_SCALE: setting TOKEN_SCALE
       * to 8 means that passing -1 through 8 for this method will return values mapped
       * between -1 and Token.MAX_VALUE.
       */
      public static Token tok(int i)
      {
          if (i == -1)
              return new BigIntegerToken(new BigInteger("-1"));
          BigInteger bint = RandomPartitioner.MAXIMUM.divide(TOKEN_SCALE).multiply(new BigInteger(""+i));
          return new BigIntegerToken(bint);
      }
  
      @Test
      public void testIntersectingRanges()
      {
          mts = new MerkleTrees(partitioner);
  
          boolean failure = true;
          mts.addMerkleTree(1, new Range<>(tok(1), tok(3)));
  
          try
          {
              mts.addMerkleTree(1, new Range<>(tok(2), tok(4)));
          }
          catch (AssertionError e)
          {
              failure = false;
          }
  
          assertFalse(failure);
      }
  
      @Test
      public void testSplit()
      {
          // split the range  (zero, zero] into:
          //  (zero,four], (four,six], (six,seven] and (seven, zero]
          mts.split(tok(4));
          mts.split(tok(6));
          mts.split(tok(7));
  
          assertEquals(4, mts.size());
          assertEquals(new Range<>(tok(7), tok(-1)), mts.get(tok(-1)));
          assertEquals(new Range<>(tok(-1), tok(4)), mts.get(tok(3)));
          assertEquals(new Range<>(tok(-1), tok(4)), mts.get(tok(4)));
          assertEquals(new Range<>(tok(4), tok(6)), mts.get(tok(6)));
          assertEquals(new Range<>(tok(6), tok(7)), mts.get(tok(7)));
  
          // check depths
          assertEquals((byte) 1, mts.get(tok(4)).depth);
          assertEquals((byte) 2, mts.get(tok(6)).depth);
          assertEquals((byte) 3, mts.get(tok(7)).depth);
          assertEquals((byte) 3, mts.get(tok(-1)).depth);
  
          try
          {
              mts.split(tok(-1));
              fail("Shouldn't be able to split outside the initial range.");
          }
          catch (AssertionError e)
          {
              // pass
          }
      }
  
      @Test
      public void testSplitLimitDepth()
      {
          mts = new MerkleTrees(partitioner);
  
          mts.addMerkleTree(Integer.MAX_VALUE, (byte) 2, fullRange());
  
          assertTrue(mts.split(tok(4)));
          assertTrue(mts.split(tok(2)));
          assertEquals(3, mts.size());
  
          // should fail to split below hashdepth
          assertFalse(mts.split(tok(1)));
          assertEquals(3, mts.size());
          assertEquals(new Range<>(tok(4), tok(-1)), mts.get(tok(-1)));
          assertEquals(new Range<>(tok(-1), tok(2)), mts.get(tok(2)));
          assertEquals(new Range<>(tok(2), tok(4)), mts.get(tok(4)));
      }
  
      @Test
      public void testSplitLimitSize()
      {
          mts = new MerkleTrees(partitioner);
  
          mts.addMerkleTree(2, fullRange());
  
          assertTrue(mts.split(tok(4)));
          assertEquals(2, mts.size());
  
          // should fail to split above maxsize
          assertFalse(mts.split(tok(2)));
          assertEquals(2, mts.size());
          assertEquals(new Range<>(tok(4), tok(-1)), mts.get(tok(-1)));
          assertEquals(new Range<>(tok(-1), tok(4)), mts.get(tok(4)));
      }
  
      @Test
      public void testInvalids()
      {
          Iterator<TreeRange> ranges;
  
          // (zero, zero]
          ranges = mts.rangeIterator();
          assertEquals(new Range<>(tok(-1), tok(-1)), ranges.next());
          assertFalse(ranges.hasNext());
  
          // all invalid
          mts.split(tok(4));
          mts.split(tok(2));
          mts.split(tok(6));
          mts.split(tok(3));
          mts.split(tok(5));
          ranges = mts.rangeIterator();
          assertEquals(new Range<>(tok(6), tok(-1)), ranges.next());
          assertEquals(new Range<>(tok(-1), tok(2)), ranges.next());
          assertEquals(new Range<>(tok(2), tok(3)), ranges.next());
          assertEquals(new Range<>(tok(3), tok(4)), ranges.next());
          assertEquals(new Range<>(tok(4), tok(5)), ranges.next());
          assertEquals(new Range<>(tok(5), tok(6)), ranges.next());
          assertEquals(new Range<>(tok(6), tok(-1)), ranges.next());
          assertFalse(ranges.hasNext());
      }
  
  
      @Test
      public void testHashFull()
      {
          byte[] val = DUMMY;
          Range<Token> range = new Range<>(tok(-1), tok(-1));
  
          // (zero, zero]
          assertNull(mts.hash(range));
  
          // validate the range
          mts.get(tok(-1)).hash(val);
  
          assertHashEquals(val, mts.hash(range));
      }
  
      @Test
      public void testHashPartial()
      {
          byte[] val = DUMMY;
          byte[] leftval = hashed(val, 1, 1);
          byte[] partialval = hashed(val, 1);
          Range<Token> left = new Range<>(tok(-1), tok(4));
          Range<Token> partial = new Range<>(tok(2), tok(4));
          Range<Token> right = new Range<>(tok(4), tok(-1));
          Range<Token> linvalid = new Range<>(tok(1), tok(4));
          Range<Token> rinvalid = new Range<>(tok(4), tok(6));
  
          // (zero,two] (two,four] (four, zero]
          mts.split(tok(4));
          mts.split(tok(2));
  
          // validate the range
          mts.get(tok(2)).hash(val);
          mts.get(tok(4)).hash(val);
          mts.get(tok(-1)).hash(val);
  
          assertHashEquals(leftval, mts.hash(left));
          assertHashEquals(partialval, mts.hash(partial));
          assertHashEquals(val, mts.hash(right));
          assertNull(mts.hash(linvalid));
          assertNull(mts.hash(rinvalid));
      }
  
      @Test
      public void testHashInner()
      {
          byte[] val = DUMMY;
          byte[] lchildval = hashed(val, 3, 3, 2);
          byte[] rchildval = hashed(val, 2, 2);
          byte[] fullval = hashed(val, 3, 3, 2, 2, 2);
          Range<Token> full = new Range<>(tok(-1), tok(-1));
          Range<Token> lchild = new Range<>(tok(-1), tok(4));
          Range<Token> rchild = new Range<>(tok(4), tok(-1));
          Range<Token> invalid = new Range<>(tok(1), tok(-1));
  
          // (zero,one] (one, two] (two,four] (four, six] (six, zero]
          mts.split(tok(4));
          mts.split(tok(2));
          mts.split(tok(6));
          mts.split(tok(1));
  
          // validate the range
          mts.get(tok(1)).hash(val);
          mts.get(tok(2)).hash(val);
          mts.get(tok(4)).hash(val);
          mts.get(tok(6)).hash(val);
          mts.get(tok(-1)).hash(val);
  
          assertHashEquals(fullval, mts.hash(full));
          assertHashEquals(lchildval, mts.hash(lchild));
          assertHashEquals(rchildval, mts.hash(rchild));
          assertNull(mts.hash(invalid));
      }
  
      @Test
      public void testHashDegenerate()
      {
          TOKEN_SCALE = new BigInteger("32");
  
          byte[] val = DUMMY;
          byte[] childfullval = hashed(val, 5, 5, 4);
          byte[] fullval = hashed(val, 5, 5, 4, 3, 2, 1);
          Range<Token> childfull = new Range<>(tok(-1), tok(4));
          Range<Token> full = new Range<>(tok(-1), tok(-1));
          Range<Token> invalid = new Range<>(tok(4), tok(-1));
  
          mts.split(tok(16));
          mts.split(tok(8));
          mts.split(tok(4));
          mts.split(tok(2));
          mts.split(tok(1));
  
          // validate the range
          mts.get(tok(1)).hash(val);
          mts.get(tok(2)).hash(val);
          mts.get(tok(4)).hash(val);
          mts.get(tok(8)).hash(val);
          mts.get(tok(16)).hash(val);
          mts.get(tok(-1)).hash(val);
  
          assertHashEquals(fullval, mts.hash(full));
          assertHashEquals(childfullval, mts.hash(childfull));
          assertNull(mts.hash(invalid));
      }
  
      @Test
      public void testHashRandom()
      {
          int max = 1000000;
          TOKEN_SCALE = new BigInteger("" + max);
  
          mts = new MerkleTrees(partitioner);
          mts.addMerkleTree(32, fullRange());
  
          Random random = new Random();
          while (true)
          {
              if (!mts.split(tok(random.nextInt(max))))
                  break;
          }
  
          // validate the tree
          TreeRangeIterator ranges = mts.rangeIterator();
          for (TreeRange range : ranges)
              range.addHash(new RowHash(range.right, new byte[0], 0));
  
          assert mts.hash(new Range<>(tok(-1), tok(-1))) != null : "Could not hash tree " + mts;
      }
  
      /**
       * Generate two trees with different splits, but containing the same keys, and
       * check that they compare equally.
       *
       * The set of keys used in this test is: #{2,4,6,8,12,14,0}
       */
      @Test
      public void testValidateTree()
      {
          TOKEN_SCALE = new BigInteger("16"); // this test needs slightly more resolution
  
          Range<Token> full = new Range<>(tok(-1), tok(-1));
          Iterator<TreeRange> ranges;
          MerkleTrees mts2 = new MerkleTrees(partitioner);
          mts2.addMerkleTree(Integer.MAX_VALUE, fullRange());
  
          mts.split(tok(8));
          mts.split(tok(4));
          mts.split(tok(12));
          mts.split(tok(6));
          mts.split(tok(10));
  
          int seed = 123456789;
  
          Random random1 = new Random(seed);
          ranges = mts.rangeIterator();
          ranges.next().addAll(new HIterator(random1, 2, 4)); // (-1,4]: depth 2
          ranges.next().addAll(new HIterator(random1, 6)); // (4,6]
          ranges.next().addAll(new HIterator(random1, 8)); // (6,8]
          ranges.next().addAll(new HIterator(/*empty*/ new int[0])); // (8,10]
          ranges.next().addAll(new HIterator(random1, 12)); // (10,12]
          ranges.next().addAll(new HIterator(random1, 14, -1)); // (12,-1]: depth 2
  
  
          mts2.split(tok(8));
          mts2.split(tok(4));
          mts2.split(tok(12));
          mts2.split(tok(2));
          mts2.split(tok(10));
          mts2.split(tok(9));
          mts2.split(tok(11));
  
          Random random2 = new Random(seed);
          ranges = mts2.rangeIterator();
          ranges.next().addAll(new HIterator(random2, 2)); // (-1,2]
          ranges.next().addAll(new HIterator(random2, 4)); // (2,4]
          ranges.next().addAll(new HIterator(random2, 6, 8)); // (4,8]: depth 2
          ranges.next().addAll(new HIterator(/*empty*/ new int[0])); // (8,9]
          ranges.next().addAll(new HIterator(/*empty*/ new int[0])); // (9,10]
          ranges.next().addAll(new HIterator(/*empty*/ new int[0])); // (10,11]: depth 4
          ranges.next().addAll(new HIterator(random2, 12)); // (11,12]: depth 4
          ranges.next().addAll(new HIterator(random2, 14, -1)); // (12,-1]: depth 2
  
          byte[] mthash = mts.hash(full);
          byte[] mt2hash = mts2.hash(full);
          assertHashEquals("Tree hashes did not match: " + mts + " && " + mts2, mthash, mt2hash);
      }
  
      @Test
      public void testSerialization() throws Exception
      {
          Range<Token> first = new Range<>(tok(3), tok(4));
  
          Collection<Range<Token>> ranges = new ArrayList<>();
  
          ranges.add(first);
          ranges.add(new Range<Token>(tok(5), tok(2)));
  
          mts = new MerkleTrees(partitioner);
          mts.addMerkleTrees(256, ranges);
  
          // populate and validate the tree
          mts.init();
          for (TreeRange range : mts.rangeIterator())
              range.addAll(new HIterator(range.right));
  
          byte[] initialhash = mts.hash(first);
  
          long serializedSize = MerkleTrees.serializer.serializedSize(mts, MessagingService.current_version);
          DataOutputBuffer out = new DataOutputBuffer();
          MerkleTrees.serializer.serialize(mts, out, MessagingService.current_version);
          byte[] serialized = out.toByteArray();
  
          assertEquals(serializedSize, serialized.length);
  
          DataInputBuffer in = new DataInputBuffer(serialized);
          MerkleTrees restored = MerkleTrees.serializer.deserialize(in, MessagingService.current_version);
  
          assertHashEquals(initialhash, restored.hash(first));
      }
  
      @Test
      public void testDifference()
      {
          int maxsize = 16;
          mts = new MerkleTrees(partitioner);
          mts.addMerkleTree(32, fullRange());
  
          MerkleTrees mts2 = new MerkleTrees(partitioner);
          mts2.addMerkleTree(32, fullRange());
  
          mts.init();
          mts2.init();
  
          int seed = 123456789;
          // add dummy hashes to both trees
          Random random1 = new Random(seed);
          for (TreeRange range : mts.rangeIterator())
              range.addAll(new HIterator(random1, range.right));
  
          Random random2 = new Random(seed);
          for (TreeRange range : mts2.rangeIterator())
              range.addAll(new HIterator(random2, range.right));
  
          TreeRange leftmost = null;
          TreeRange middle = null;
  
          mts.maxsize(fullRange(), maxsize + 2); // give some room for splitting
  
          // split the leftmost
          Iterator<TreeRange> ranges = mts.rangeIterator();
          leftmost = ranges.next();
          mts.split(leftmost.right);
  
          // set the hashes for the leaf of the created split
          middle = mts.get(leftmost.right);
          middle.hash(digest("arbitrary!"));
          mts.get(partitioner.midpoint(leftmost.left, leftmost.right)).hash(digest("even more arbitrary!"));
  
          // trees should disagree for (leftmost.left, middle.right]
          List<Range<Token>> diffs = MerkleTrees.difference(mts, mts2);
          assertEquals(diffs + " contains wrong number of differences:", 1, diffs.size());
          assertTrue(diffs.contains(new Range<>(leftmost.left, middle.right)));
      }
  
      /**
       * Return the root hash of a binary tree with leaves at the given depths
       * and with the given hash val in each leaf.
       */
      byte[] hashed(byte[] val, Integer... depths)
      {
          ArrayDeque<Integer> dstack = new ArrayDeque<Integer>();
          ArrayDeque<byte[]> hstack = new ArrayDeque<byte[]>();
          Iterator<Integer> depthiter = Arrays.asList(depths).iterator();
          if (depthiter.hasNext())
          {
              dstack.push(depthiter.next());
              hstack.push(val);
          }
          while (depthiter.hasNext())
          {
              Integer depth = depthiter.next();
              byte[] hash = val;
              while (depth.equals(dstack.peek()))
              {
                  // consume the stack
                  hash = MerkleTree.xor(hstack.pop(), hash);
                  depth = dstack.pop()-1;
              }
              dstack.push(depth);
              hstack.push(hash);
          }
          assert hstack.size() == 1;
          return hstack.pop();
      }
  
      public static class HIterator extends AbstractIterator<RowHash>
      {
          private final Random random;
          private final Iterator<Token> tokens;
  
          HIterator(int... tokens)
          {
              this(new Random(), tokens);
          }
  
          HIterator(Random random, int... tokens)
          {
              List<Token> tlist = new ArrayList<>(tokens.length);
              for (int token : tokens)
                  tlist.add(tok(token));
              this.tokens = tlist.iterator();
              this.random = random;
          }
  
          public HIterator(Token... tokens)
          {
              this(new Random(), tokens);
          }
  
          HIterator(Random random, Token... tokens)
          {
              this(random, Arrays.asList(tokens).iterator());
          }
  
          private HIterator(Random random, Iterator<Token> tokens)
          {
              this.random = random;
              this.tokens = tokens;
          }
  
          public RowHash computeNext()
          {
              if (tokens.hasNext())
              {
                  byte[] digest = new byte[32];
                  random.nextBytes(digest);
                  return new RowHash(tokens.next(), digest, 12345L);
              }
              return endOfData();
          }
      }
  }
  