package test import ( "sort" "testing" ic "github.com/libp2p/go-libp2p/core/crypto" "github.com/libp2p/go-libp2p/core/peer" pstore "github.com/libp2p/go-libp2p/core/peerstore" pt "github.com/libp2p/go-libp2p/core/test" "github.com/stretchr/testify/require" ) var keyBookSuite = map[string]func(kb pstore.KeyBook) func(*testing.T){ "AddGetPrivKey": testKeybookPrivKey, "AddGetPubKey": testKeyBookPubKey, "PeersWithKeys": testKeyBookPeers, "PubKeyAddedOnRetrieve": testInlinedPubKeyAddedOnRetrieve, "Delete": testKeyBookDelete, } type KeyBookFactory func() (pstore.KeyBook, func()) func TestKeyBook(t *testing.T, factory KeyBookFactory) { for name, test := range keyBookSuite { // Create a new peerstore. kb, closeFunc := factory() // Run the test. t.Run(name, test(kb)) // Cleanup. if closeFunc != nil { closeFunc() } } } func testKeybookPrivKey(kb pstore.KeyBook) func(t *testing.T) { return func(t *testing.T) { if peers := kb.PeersWithKeys(); len(peers) > 0 { t.Error("expected peers to be empty on init") } priv, _, err := pt.RandTestKeyPair(ic.RSA, 2048) if err != nil { t.Fatal(err) } id, err := peer.IDFromPrivateKey(priv) if err != nil { t.Error(err) } if res := kb.PrivKey(id); res != nil { t.Error("retrieving private key should have failed") } err = kb.AddPrivKey(id, priv) if err != nil { t.Error(err) } if res := kb.PrivKey(id); !priv.Equals(res) { t.Error("retrieved private key did not match stored private key") } if peers := kb.PeersWithKeys(); len(peers) != 1 || peers[0] != id { t.Error("list of peers did not include test peer") } } } func testKeyBookPubKey(kb pstore.KeyBook) func(t *testing.T) { return func(t *testing.T) { if peers := kb.PeersWithKeys(); len(peers) > 0 { t.Error("expected peers to be empty on init") } _, pub, err := pt.RandTestKeyPair(ic.RSA, 2048) if err != nil { t.Fatal(err) } id, err := peer.IDFromPublicKey(pub) if err != nil { t.Fatal(err) } if res := kb.PubKey(id); res != nil { t.Error("retrieving public key should have failed") } err = kb.AddPubKey(id, pub) if err != nil { t.Error(err) } if res := kb.PubKey(id); !pub.Equals(res) { t.Error("retrieved public key did not match stored public key") } if peers := kb.PeersWithKeys(); len(peers) != 1 || peers[0] != id { t.Error("list of peers did not include test peer") } } } func testKeyBookPeers(kb pstore.KeyBook) func(t *testing.T) { return func(t *testing.T) { if peers := kb.PeersWithKeys(); len(peers) > 0 { t.Error("expected peers to be empty on init") } var peers peer.IDSlice for i := 0; i < 10; i++ { // Add a public key. _, pub, err := pt.RandTestKeyPair(ic.RSA, 2048) if err != nil { t.Fatal(err) } p1, err := peer.IDFromPublicKey(pub) if err != nil { t.Fatal(err) } kb.AddPubKey(p1, pub) // Add a private key. priv, _, err := pt.RandTestKeyPair(ic.RSA, 2048) if err != nil { t.Fatal(err) } p2, err := peer.IDFromPrivateKey(priv) if err != nil { t.Fatal(err) } kb.AddPrivKey(p2, priv) peers = append(peers, []peer.ID{p1, p2}...) } kbPeers := kb.PeersWithKeys() sort.Sort(kbPeers) sort.Sort(peers) for i, p := range kbPeers { if p != peers[i] { t.Errorf("mismatch of peer at index %d", i) } } } } func testInlinedPubKeyAddedOnRetrieve(kb pstore.KeyBook) func(t *testing.T) { return func(t *testing.T) { t.Skip("key inlining disabled for now: see libp2p/specs#111") if peers := kb.PeersWithKeys(); len(peers) > 0 { t.Error("expected peers to be empty on init") } // Key small enough for inlining. _, pub, err := ic.GenerateKeyPair(ic.Ed25519, 256) if err != nil { t.Error(err) } id, err := peer.IDFromPublicKey(pub) if err != nil { t.Error(err) } pubKey := kb.PubKey(id) if !pubKey.Equals(pub) { t.Error("mismatch between original public key and keybook-calculated one") } } } func testKeyBookDelete(kb pstore.KeyBook) func(t *testing.T) { return func(t *testing.T) { // don't use an ed25519 key here, otherwise the key book might try to derive the pubkey from the peer ID priv, pub, err := ic.GenerateKeyPair(ic.RSA, 2048) require.NoError(t, err) p, err := peer.IDFromPublicKey(pub) require.NoError(t, err) require.NoError(t, kb.AddPubKey(p, pub)) require.NoError(t, kb.AddPrivKey(p, priv)) require.NotNil(t, kb.PrivKey(p)) require.NotNil(t, kb.PubKey(p)) kb.RemovePeer(p) require.Nil(t, kb.PrivKey(p)) require.Nil(t, kb.PubKey(p)) } } var keybookBenchmarkSuite = map[string]func(kb pstore.KeyBook) func(*testing.B){ "PubKey": benchmarkPubKey, "AddPubKey": benchmarkAddPubKey, "PrivKey": benchmarkPrivKey, "AddPrivKey": benchmarkAddPrivKey, "PeersWithKeys": benchmarkPeersWithKeys, } func BenchmarkKeyBook(b *testing.B, factory KeyBookFactory) { ordernames := make([]string, 0, len(keybookBenchmarkSuite)) for name := range keybookBenchmarkSuite { ordernames = append(ordernames, name) } sort.Strings(ordernames) for _, name := range ordernames { bench := keybookBenchmarkSuite[name] kb, closeFunc := factory() b.Run(name, bench(kb)) if closeFunc != nil { closeFunc() } } } func benchmarkPubKey(kb pstore.KeyBook) func(*testing.B) { return func(b *testing.B) { _, pub, err := pt.RandTestKeyPair(ic.RSA, 2048) if err != nil { b.Fatal(err) } id, err := peer.IDFromPublicKey(pub) if err != nil { b.Fatal(err) } err = kb.AddPubKey(id, pub) if err != nil { b.Fatal(err) } b.ResetTimer() for i := 0; i < b.N; i++ { kb.PubKey(id) } } } func benchmarkAddPubKey(kb pstore.KeyBook) func(*testing.B) { return func(b *testing.B) { _, pub, err := pt.RandTestKeyPair(ic.RSA, 2048) if err != nil { b.Fatal(err) } id, err := peer.IDFromPublicKey(pub) if err != nil { b.Fatal(err) } b.ResetTimer() for i := 0; i < b.N; i++ { kb.AddPubKey(id, pub) } } } func benchmarkPrivKey(kb pstore.KeyBook) func(*testing.B) { return func(b *testing.B) { priv, _, err := pt.RandTestKeyPair(ic.RSA, 2048) if err != nil { b.Fatal(err) } id, err := peer.IDFromPrivateKey(priv) if err != nil { b.Fatal(err) } err = kb.AddPrivKey(id, priv) if err != nil { b.Fatal(err) } b.ResetTimer() for i := 0; i < b.N; i++ { kb.PrivKey(id) } } } func benchmarkAddPrivKey(kb pstore.KeyBook) func(*testing.B) { return func(b *testing.B) { priv, _, err := pt.RandTestKeyPair(ic.RSA, 2048) if err != nil { b.Fatal(err) } id, err := peer.IDFromPrivateKey(priv) if err != nil { b.Fatal(err) } b.ResetTimer() for i := 0; i < b.N; i++ { kb.AddPrivKey(id, priv) } } } func benchmarkPeersWithKeys(kb pstore.KeyBook) func(*testing.B) { return func(b *testing.B) { for i := 0; i < 10; i++ { priv, pub, err := pt.RandTestKeyPair(ic.RSA, 2048) if err != nil { b.Fatal(err) } id, err := peer.IDFromPublicKey(pub) if err != nil { b.Fatal(err) } err = kb.AddPubKey(id, pub) if err != nil { b.Fatal(err) } err = kb.AddPrivKey(id, priv) if err != nil { b.Fatal(err) } } b.ResetTimer() for i := 0; i < b.N; i++ { kb.PeersWithKeys() } } }