forked from libp2p/go-libp2p
-
Notifications
You must be signed in to change notification settings - Fork 0
/
keybook.go
97 lines (83 loc) · 1.82 KB
/
keybook.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
package pstoremem
import (
"errors"
"sync"
ic "github.com/JonyBepary/go-libp2p-pq/core/crypto"
"github.com/JonyBepary/go-libp2p-pq/core/peer"
pstore "github.com/JonyBepary/go-libp2p-pq/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()
}