1package memberlist
2
3import (
4	"bytes"
5	"fmt"
6	"sync"
7)
8
9type Keyring struct {
10	// Keys stores the key data used during encryption and decryption. It is
11	// ordered in such a way where the first key (index 0) is the primary key,
12	// which is used for encrypting messages, and is the first key tried during
13	// message decryption.
14	keys [][]byte
15
16	// The keyring lock is used while performing IO operations on the keyring.
17	l sync.Mutex
18}
19
20// Init allocates substructures
21func (k *Keyring) init() {
22	k.keys = make([][]byte, 0)
23}
24
25// NewKeyring constructs a new container for a set of encryption keys. The
26// keyring contains all key data used internally by memberlist.
27//
28// While creating a new keyring, you must do one of:
29//   - Omit keys and primary key, effectively disabling encryption
30//   - Pass a set of keys plus the primary key
31//   - Pass only a primary key
32//
33// If only a primary key is passed, then it will be automatically added to the
34// keyring. If creating a keyring with multiple keys, one key must be designated
35// primary by passing it as the primaryKey. If the primaryKey does not exist in
36// the list of secondary keys, it will be automatically added at position 0.
37//
38// A key should be either 16, 24, or 32 bytes to select AES-128,
39// AES-192, or AES-256.
40func NewKeyring(keys [][]byte, primaryKey []byte) (*Keyring, error) {
41	keyring := &Keyring{}
42	keyring.init()
43
44	if len(keys) > 0 || len(primaryKey) > 0 {
45		if len(primaryKey) == 0 {
46			return nil, fmt.Errorf("Empty primary key not allowed")
47		}
48		if err := keyring.AddKey(primaryKey); err != nil {
49			return nil, err
50		}
51		for _, key := range keys {
52			if err := keyring.AddKey(key); err != nil {
53				return nil, err
54			}
55		}
56	}
57
58	return keyring, nil
59}
60
61// AddKey will install a new key on the ring. Adding a key to the ring will make
62// it available for use in decryption. If the key already exists on the ring,
63// this function will just return noop.
64//
65// key should be either 16, 24, or 32 bytes to select AES-128,
66// AES-192, or AES-256.
67func (k *Keyring) AddKey(key []byte) error {
68	if l := len(key); l != 16 && l != 24 && l != 32 {
69		return fmt.Errorf("key size must be 16, 24 or 32 bytes")
70	}
71
72	// No-op if key is already installed
73	for _, installedKey := range k.keys {
74		if bytes.Equal(installedKey, key) {
75			return nil
76		}
77	}
78
79	keys := append(k.keys, key)
80	primaryKey := k.GetPrimaryKey()
81	if primaryKey == nil {
82		primaryKey = key
83	}
84	k.installKeys(keys, primaryKey)
85	return nil
86}
87
88// UseKey changes the key used to encrypt messages. This is the only key used to
89// encrypt messages, so peers should know this key before this method is called.
90func (k *Keyring) UseKey(key []byte) error {
91	for _, installedKey := range k.keys {
92		if bytes.Equal(key, installedKey) {
93			k.installKeys(k.keys, key)
94			return nil
95		}
96	}
97	return fmt.Errorf("Requested key is not in the keyring")
98}
99
100// RemoveKey drops a key from the keyring. This will return an error if the key
101// requested for removal is currently at position 0 (primary key).
102func (k *Keyring) RemoveKey(key []byte) error {
103	if bytes.Equal(key, k.keys[0]) {
104		return fmt.Errorf("Removing the primary key is not allowed")
105	}
106	for i, installedKey := range k.keys {
107		if bytes.Equal(key, installedKey) {
108			keys := append(k.keys[:i], k.keys[i+1:]...)
109			k.installKeys(keys, k.keys[0])
110		}
111	}
112	return nil
113}
114
115// installKeys will take out a lock on the keyring, and replace the keys with a
116// new set of keys. The key indicated by primaryKey will be installed as the new
117// primary key.
118func (k *Keyring) installKeys(keys [][]byte, primaryKey []byte) {
119	k.l.Lock()
120	defer k.l.Unlock()
121
122	newKeys := [][]byte{primaryKey}
123	for _, key := range keys {
124		if !bytes.Equal(key, primaryKey) {
125			newKeys = append(newKeys, key)
126		}
127	}
128	k.keys = newKeys
129}
130
131// GetKeys returns the current set of keys on the ring.
132func (k *Keyring) GetKeys() [][]byte {
133	k.l.Lock()
134	defer k.l.Unlock()
135
136	return k.keys
137}
138
139// GetPrimaryKey returns the key on the ring at position 0. This is the key used
140// for encrypting messages, and is the first key tried for decrypting messages.
141func (k *Keyring) GetPrimaryKey() (key []byte) {
142	k.l.Lock()
143	defer k.l.Unlock()
144
145	if len(k.keys) > 0 {
146		key = k.keys[0]
147	}
148	return
149}
150