device: immediately rekey all peers after changing device private key

Reported-by: Derrick Pallas <derrick@pallas.us>

2 files changed, 25 insertions, 0 deletions

diff --git a/device/device.go b/device/device.go
index a583fa9..ab5e4b0 100644
--- a/device/device.go
+++ b/device/device.go
@@ -207,6 +207,10 @@ func (device *Device) SetPrivateKey(sk NoisePrivateKey) error {
 	device.staticIdentity.Lock()
 	defer device.staticIdentity.Unlock()
 
+	if sk.Equals(device.staticIdentity.privateKey) {
+		return nil
+	}
+
 	device.peers.Lock()
 	defer device.peers.Unlock()
 
@@ -246,6 +250,8 @@ func (device *Device) SetPrivateKey(sk NoisePrivateKey) error {
 
 		if isZero(handshake.precomputedStaticStatic[:]) {
 			unsafeRemovePeer(device, peer, key)
+		} else {
+			peer.ExpireCurrentKeypairs()
 		}
 	}
 
diff --git a/device/peer.go b/device/peer.go
index 4e7f2da..256e4f5 100644
--- a/device/peer.go
+++ b/device/peer.go
@@ -232,6 +232,25 @@ func (peer *Peer) ZeroAndFlushAll() {
 	peer.FlushNonceQueue()
 }
 
+func (peer *Peer) ExpireCurrentKeypairs() {
+	handshake := &peer.handshake
+	handshake.mutex.Lock()
+	peer.device.indexTable.Delete(handshake.localIndex)
+	handshake.Clear()
+	handshake.mutex.Unlock()
+	peer.handshake.lastSentHandshake = time.Now().Add(-(RekeyTimeout + time.Second))
+
+	keypairs := &peer.keypairs
+	keypairs.Lock()
+	if keypairs.current != nil {
+		keypairs.current.sendNonce = RejectAfterMessages
+	}
+	if keypairs.next != nil {
+		keypairs.next.sendNonce = RejectAfterMessages
+	}
+	keypairs.Unlock()
+}
+
 func (peer *Peer) Stop() {
 
 	// prevent simultaneous start/stop operations