github.com/aquanetwork/aquachain@v1.7.8/trie/proof_test.go (about)

     1  // Copyright 2015 The aquachain Authors
     2  // This file is part of the aquachain library.
     3  //
     4  // The aquachain library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // The aquachain library is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the aquachain library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package trie
    18  
    19  import (
    20  	"bytes"
    21  	crand "crypto/rand"
    22  	mrand "math/rand"
    23  	"testing"
    24  	"time"
    25  
    26  	"gitlab.com/aquachain/aquachain/aquadb"
    27  	"gitlab.com/aquachain/aquachain/common"
    28  	"gitlab.com/aquachain/aquachain/crypto"
    29  )
    30  
    31  func init() {
    32  	mrand.Seed(time.Now().Unix())
    33  }
    34  
    35  func TestProof(t *testing.T) {
    36  	trie, vals := randomTrie(500)
    37  	root := trie.Hash()
    38  	for _, kv := range vals {
    39  		proofs := aquadb.NewMemDatabase()
    40  		if trie.Prove(kv.k, 0, proofs) != nil {
    41  			t.Fatalf("missing key %x while constructing proof", kv.k)
    42  		}
    43  		val, err, _ := VerifyProof(root, kv.k, proofs)
    44  		if err != nil {
    45  			t.Fatalf("VerifyProof error for key %x: %v\nraw proof: %v", kv.k, err, proofs)
    46  		}
    47  		if !bytes.Equal(val, kv.v) {
    48  			t.Fatalf("VerifyProof returned wrong value for key %x: got %x, want %x", kv.k, val, kv.v)
    49  		}
    50  	}
    51  }
    52  
    53  func TestOneElementProof(t *testing.T) {
    54  	trie := new(Trie)
    55  	updateString(trie, "k", "v")
    56  	proofs := aquadb.NewMemDatabase()
    57  	trie.Prove([]byte("k"), 0, proofs)
    58  	if len(proofs.Keys()) != 1 {
    59  		t.Error("proof should have one element")
    60  	}
    61  	val, err, _ := VerifyProof(trie.Hash(), []byte("k"), proofs)
    62  	if err != nil {
    63  		t.Fatalf("VerifyProof error: %v\nproof hashes: %v", err, proofs.Keys())
    64  	}
    65  	if !bytes.Equal(val, []byte("v")) {
    66  		t.Fatalf("VerifyProof returned wrong value: got %x, want 'k'", val)
    67  	}
    68  }
    69  
    70  func TestVerifyBadProof(t *testing.T) {
    71  	trie, vals := randomTrie(800)
    72  	root := trie.Hash()
    73  	for _, kv := range vals {
    74  		proofs := aquadb.NewMemDatabase()
    75  		trie.Prove(kv.k, 0, proofs)
    76  		if len(proofs.Keys()) == 0 {
    77  			t.Fatal("zero length proof")
    78  		}
    79  		keys := proofs.Keys()
    80  		key := keys[mrand.Intn(len(keys))]
    81  		node, _ := proofs.Get(key)
    82  		proofs.Delete(key)
    83  		mutateByte(node)
    84  		proofs.Put(crypto.Keccak256(node), node)
    85  		if _, err, _ := VerifyProof(root, kv.k, proofs); err == nil {
    86  			t.Fatalf("expected proof to fail for key %x", kv.k)
    87  		}
    88  	}
    89  }
    90  
    91  // mutateByte changes one byte in b.
    92  func mutateByte(b []byte) {
    93  	for r := mrand.Intn(len(b)); ; {
    94  		new := byte(mrand.Intn(255))
    95  		if new != b[r] {
    96  			b[r] = new
    97  			break
    98  		}
    99  	}
   100  }
   101  
   102  func BenchmarkProve(b *testing.B) {
   103  	trie, vals := randomTrie(100)
   104  	var keys []string
   105  	for k := range vals {
   106  		keys = append(keys, k)
   107  	}
   108  
   109  	b.ResetTimer()
   110  	for i := 0; i < b.N; i++ {
   111  		kv := vals[keys[i%len(keys)]]
   112  		proofs := aquadb.NewMemDatabase()
   113  		if trie.Prove(kv.k, 0, proofs); len(proofs.Keys()) == 0 {
   114  			b.Fatalf("zero length proof for %x", kv.k)
   115  		}
   116  	}
   117  }
   118  
   119  func BenchmarkVerifyProof(b *testing.B) {
   120  	trie, vals := randomTrie(100)
   121  	root := trie.Hash()
   122  	var keys []string
   123  	var proofs []*aquadb.MemDatabase
   124  	for k := range vals {
   125  		keys = append(keys, k)
   126  		proof := aquadb.NewMemDatabase()
   127  		trie.Prove([]byte(k), 0, proof)
   128  		proofs = append(proofs, proof)
   129  	}
   130  
   131  	b.ResetTimer()
   132  	for i := 0; i < b.N; i++ {
   133  		im := i % len(keys)
   134  		if _, err, _ := VerifyProof(root, []byte(keys[im]), proofs[im]); err != nil {
   135  			b.Fatalf("key %x: %v", keys[im], err)
   136  		}
   137  	}
   138  }
   139  
   140  func randomTrie(n int) (*Trie, map[string]*kv) {
   141  	trie := new(Trie)
   142  	vals := make(map[string]*kv)
   143  	for i := byte(0); i < 100; i++ {
   144  		value := &kv{common.LeftPadBytes([]byte{i}, 32), []byte{i}, false}
   145  		value2 := &kv{common.LeftPadBytes([]byte{i + 10}, 32), []byte{i}, false}
   146  		trie.Update(value.k, value.v)
   147  		trie.Update(value2.k, value2.v)
   148  		vals[string(value.k)] = value
   149  		vals[string(value2.k)] = value2
   150  	}
   151  	for i := 0; i < n; i++ {
   152  		value := &kv{randBytes(32), randBytes(20), false}
   153  		trie.Update(value.k, value.v)
   154  		vals[string(value.k)] = value
   155  	}
   156  	return trie, vals
   157  }
   158  
   159  func randBytes(n int) []byte {
   160  	r := make([]byte, n)
   161  	crand.Read(r)
   162  	return r
   163  }