From 55b88be009014676c1a18e7264fcac4f6f6c939a Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Wed, 12 Jun 2019 22:45:31 +0800
Subject: [PATCH] check that the peer doesn't update keys too quickly

---
 internal/handshake/updatable_aead.go      | 12 +++++++-----
 internal/handshake/updatable_aead_test.go | 20 +++++++++++++++++++-
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/internal/handshake/updatable_aead.go b/internal/handshake/updatable_aead.go
index 153c1eb55..ee36c0d59 100644
--- a/internal/handshake/updatable_aead.go
+++ b/internal/handshake/updatable_aead.go
@@ -112,12 +112,14 @@ func (a *updatableAEAD) Open(dst, src []byte, pn protocol.PacketNumber, kp proto
 		// try opening the packet with the next key phase
 		dec, err := a.nextRcvAEAD.Open(dst, a.nonceBuf, src, ad)
 		if err != nil {
-			err = ErrDecryptionFailed
-		} else {
-			// if opening succeeds, roll over to the next key phase
-			a.rollKeys()
-			a.firstRcvdWithCurrentKey = pn
+			return nil, ErrDecryptionFailed
 		}
+		// Opening succeeded. Check if the peer was allowed to update.
+		if a.firstSentWithCurrentKey == protocol.InvalidPacketNumber {
+			return nil, qerr.Error(qerr.ProtocolViolation, "keys updated too quickly")
+		}
+		a.rollKeys()
+		a.firstRcvdWithCurrentKey = pn
 		return dec, err
 	}
 	// The AEAD we're using here will be the qtls.aeadAESGCM13.
diff --git a/internal/handshake/updatable_aead_test.go b/internal/handshake/updatable_aead_test.go
index 7737cb04d..36ce8fc45 100644
--- a/internal/handshake/updatable_aead_test.go
+++ b/internal/handshake/updatable_aead_test.go
@@ -111,13 +111,17 @@ var _ = Describe("Updatable AEAD", func() {
 			})
 
 			It("updates the keys when receiving a packet with the next key phase", func() {
+				// receive the first packet at key phase zero
 				encrypted0 := client.Seal(nil, msg, 0x42, ad)
 				decrypted, err := server.Open(nil, encrypted0, 0x42, protocol.KeyPhaseZero, ad)
 				Expect(err).ToNot(HaveOccurred())
 				Expect(decrypted).To(Equal(msg))
+				// send one packet at key phase zero
+				Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseZero))
+				_ = server.Seal(nil, msg, 0x1, ad)
+				// now received a message at key phase one
 				client.rollKeys()
 				encrypted1 := client.Seal(nil, msg, 0x43, ad)
-				Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseZero))
 				decrypted, err = server.Open(nil, encrypted1, 0x43, protocol.KeyPhaseOne, ad)
 				Expect(err).ToNot(HaveOccurred())
 				Expect(decrypted).To(Equal(msg))
@@ -130,6 +134,8 @@ var _ = Describe("Updatable AEAD", func() {
 				// receive the first packet with key phase 0
 				_, err := server.Open(nil, encrypted01, 0x42, protocol.KeyPhaseZero, ad)
 				Expect(err).ToNot(HaveOccurred())
+				// send one packet at key phase zero
+				_ = server.Seal(nil, msg, 0x1, ad)
 				// now receive a packet with key phase 1
 				client.rollKeys()
 				encrypted1 := client.Seal(nil, msg, 0x44, ad)
@@ -150,6 +156,18 @@ var _ = Describe("Updatable AEAD", func() {
 				_, err := server.Open(nil, encrypted, 0x1337, protocol.KeyPhaseOne, ad)
 				Expect(err).To(MatchError("PROTOCOL_VIOLATION: wrong initial keyphase"))
 			})
+
+			It("errors when the peer updates keys too frequently", func() {
+				// receive the first packet at key phase zero
+				encrypted0 := client.Seal(nil, msg, 0x42, ad)
+				_, err := server.Open(nil, encrypted0, 0x42, protocol.KeyPhaseZero, ad)
+				Expect(err).ToNot(HaveOccurred())
+				// now receive a packet at key phase one, before having sent any packets
+				client.rollKeys()
+				encrypted1 := client.Seal(nil, msg, 0x42, ad)
+				_, err = server.Open(nil, encrypted1, 0x42, protocol.KeyPhaseOne, ad)
+				Expect(err).To(MatchError("PROTOCOL_VIOLATION: keys updated too quickly"))
+			})
 		})
 	})
 })