ceremonyclient/go-libp2p/p2p/host/peerstore/test/keybook_suite.go
2023-08-20 23:07:43 -05:00

338 lines
7.2 KiB
Go

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()
}
}
}