github.com/waltonchain/waltonchain_gwtc_src@v1.1.4-0.20201225072101-8a298c95a819/trie/proof_test.go (about)

     1  // Copyright 2015 The go-ethereum Authors
     2  // This file is part of the go-ethereum library.
     3  //
     4  // The go-wtc 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 go-wtc 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 go-ethereum 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  	"github.com/wtc/go-wtc/common"
    27  	"github.com/wtc/go-wtc/rlp"
    28  )
    29  
    30  func init() {
    31  	mrand.Seed(time.Now().Unix())
    32  }
    33  
    34  func TestProof(t *testing.T) {
    35  	trie, vals := randomTrie(500)
    36  	root := trie.Hash()
    37  	for _, kv := range vals {
    38  		proof := trie.Prove(kv.k)
    39  		if proof == nil {
    40  			t.Fatalf("missing key %x while constructing proof", kv.k)
    41  		}
    42  		val, err := VerifyProof(root, kv.k, proof)
    43  		if err != nil {
    44  			t.Fatalf("VerifyProof error for key %x: %v\nraw proof: %x", kv.k, err, proof)
    45  		}
    46  		if !bytes.Equal(val, kv.v) {
    47  			t.Fatalf("VerifyProof returned wrong value for key %x: got %x, want %x", kv.k, val, kv.v)
    48  		}
    49  	}
    50  }
    51  
    52  func TestOneElementProof(t *testing.T) {
    53  	trie := new(Trie)
    54  	updateString(trie, "k", "v")
    55  	proof := trie.Prove([]byte("k"))
    56  	if proof == nil {
    57  		t.Fatal("nil proof")
    58  	}
    59  	if len(proof) != 1 {
    60  		t.Error("proof should have one element")
    61  	}
    62  	val, err := VerifyProof(trie.Hash(), []byte("k"), proof)
    63  	if err != nil {
    64  		t.Fatalf("VerifyProof error: %v\nraw proof: %x", err, proof)
    65  	}
    66  	if !bytes.Equal(val, []byte("v")) {
    67  		t.Fatalf("VerifyProof returned wrong value: got %x, want 'k'", val)
    68  	}
    69  }
    70  
    71  func TestVerifyBadProof(t *testing.T) {
    72  	trie, vals := randomTrie(800)
    73  	root := trie.Hash()
    74  	for _, kv := range vals {
    75  		proof := trie.Prove(kv.k)
    76  		if proof == nil {
    77  			t.Fatal("nil proof")
    78  		}
    79  		mutateByte(proof[mrand.Intn(len(proof))])
    80  		if _, err := VerifyProof(root, kv.k, proof); err == nil {
    81  			t.Fatalf("expected proof to fail for key %x", kv.k)
    82  		}
    83  	}
    84  }
    85  
    86  // mutateByte changes one byte in b.
    87  func mutateByte(b []byte) {
    88  	for r := mrand.Intn(len(b)); ; {
    89  		new := byte(mrand.Intn(255))
    90  		if new != b[r] {
    91  			b[r] = new
    92  			break
    93  		}
    94  	}
    95  }
    96  
    97  func BenchmarkProve(b *testing.B) {
    98  	trie, vals := randomTrie(100)
    99  	var keys []string
   100  	for k := range vals {
   101  		keys = append(keys, k)
   102  	}
   103  
   104  	b.ResetTimer()
   105  	for i := 0; i < b.N; i++ {
   106  		kv := vals[keys[i%len(keys)]]
   107  		if trie.Prove(kv.k) == nil {
   108  			b.Fatalf("nil proof for %x", kv.k)
   109  		}
   110  	}
   111  }
   112  
   113  func BenchmarkVerifyProof(b *testing.B) {
   114  	trie, vals := randomTrie(100)
   115  	root := trie.Hash()
   116  	var keys []string
   117  	var proofs [][]rlp.RawValue
   118  	for k := range vals {
   119  		keys = append(keys, k)
   120  		proofs = append(proofs, trie.Prove([]byte(k)))
   121  	}
   122  
   123  	b.ResetTimer()
   124  	for i := 0; i < b.N; i++ {
   125  		im := i % len(keys)
   126  		if _, err := VerifyProof(root, []byte(keys[im]), proofs[im]); err != nil {
   127  			b.Fatalf("key %x: %v", keys[im], err)
   128  		}
   129  	}
   130  }
   131  
   132  func randomTrie(n int) (*Trie, map[string]*kv) {
   133  	trie := new(Trie)
   134  	vals := make(map[string]*kv)
   135  	for i := byte(0); i < 100; i++ {
   136  		value := &kv{common.LeftPadBytes([]byte{i}, 32), []byte{i}, false}
   137  		value2 := &kv{common.LeftPadBytes([]byte{i + 10}, 32), []byte{i}, false}
   138  		trie.Update(value.k, value.v)
   139  		trie.Update(value2.k, value2.v)
   140  		vals[string(value.k)] = value
   141  		vals[string(value2.k)] = value2
   142  	}
   143  	for i := 0; i < n; i++ {
   144  		value := &kv{randBytes(32), randBytes(20), false}
   145  		trie.Update(value.k, value.v)
   146  		vals[string(value.k)] = value
   147  	}
   148  	return trie, vals
   149  }
   150  
   151  func randBytes(n int) []byte {
   152  	r := make([]byte, n)
   153  	crand.Read(r)
   154  	return r
   155  }