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

98 lines
1.8 KiB
Go

package pstoremem
import (
"errors"
"sync"
ic "github.com/libp2p/go-libp2p/core/crypto"
"github.com/libp2p/go-libp2p/core/peer"
pstore "github.com/libp2p/go-libp2p/core/peerstore"
)
type memoryKeyBook struct {
sync.RWMutex // same lock. wont happen a ton.
pks map[peer.ID]ic.PubKey
sks map[peer.ID]ic.PrivKey
}
var _ pstore.KeyBook = (*memoryKeyBook)(nil)
func NewKeyBook() *memoryKeyBook {
return &memoryKeyBook{
pks: map[peer.ID]ic.PubKey{},
sks: map[peer.ID]ic.PrivKey{},
}
}
func (mkb *memoryKeyBook) PeersWithKeys() peer.IDSlice {
mkb.RLock()
ps := make(peer.IDSlice, 0, len(mkb.pks)+len(mkb.sks))
for p := range mkb.pks {
ps = append(ps, p)
}
for p := range mkb.sks {
if _, found := mkb.pks[p]; !found {
ps = append(ps, p)
}
}
mkb.RUnlock()
return ps
}
func (mkb *memoryKeyBook) PubKey(p peer.ID) ic.PubKey {
mkb.RLock()
pk := mkb.pks[p]
mkb.RUnlock()
if pk != nil {
return pk
}
pk, err := p.ExtractPublicKey()
if err == nil {
mkb.Lock()
mkb.pks[p] = pk
mkb.Unlock()
}
return pk
}
func (mkb *memoryKeyBook) AddPubKey(p peer.ID, pk ic.PubKey) error {
// check it's correct first
if !p.MatchesPublicKey(pk) {
return errors.New("ID does not match PublicKey")
}
mkb.Lock()
mkb.pks[p] = pk
mkb.Unlock()
return nil
}
func (mkb *memoryKeyBook) PrivKey(p peer.ID) ic.PrivKey {
mkb.RLock()
defer mkb.RUnlock()
return mkb.sks[p]
}
func (mkb *memoryKeyBook) AddPrivKey(p peer.ID, sk ic.PrivKey) error {
if sk == nil {
return errors.New("sk is nil (PrivKey)")
}
// check it's correct first
if !p.MatchesPrivateKey(sk) {
return errors.New("ID does not match PrivateKey")
}
mkb.Lock()
mkb.sks[p] = sk
mkb.Unlock()
return nil
}
func (mkb *memoryKeyBook) RemovePeer(p peer.ID) {
mkb.Lock()
delete(mkb.sks, p)
delete(mkb.pks, p)
mkb.Unlock()
}