implement correct dropping of Initial keys

internal/handshake/aead.go

@@ -79,3 +79,69 @@ func (o *longHeaderOpener) Open(dst, src []byte, pn protocol.PacketNumber, ad []
 func (o *longHeaderOpener) DecryptHeader(sample []byte, firstByte *byte, pnBytes []byte) {
 	o.headerProtector.DecryptHeader(sample, firstByte, pnBytes)
 }
+
+type handshakeSealer struct {
+	LongHeaderSealer
+
+	dropInitialKeys func()
+	dropped         bool
+}
+
+func newHandshakeSealer(
+	aead cipher.AEAD,
+	headerProtector headerProtector,
+	dropInitialKeys func(),
+	perspective protocol.Perspective,
+) LongHeaderSealer {
+	sealer := newLongHeaderSealer(aead, headerProtector)
+	// The client drops Initial keys when sending the first Handshake packet.
+	if perspective == protocol.PerspectiveServer {
+		return sealer
+	}
+	return &handshakeSealer{
+		LongHeaderSealer: sealer,
+		dropInitialKeys:  dropInitialKeys,
+	}
+}
+
+func (s *handshakeSealer) Seal(dst, src []byte, pn protocol.PacketNumber, ad []byte) []byte {
+	data := s.LongHeaderSealer.Seal(dst, src, pn, ad)
+	if !s.dropped {
+		s.dropInitialKeys()
+		s.dropped = true
+	}
+	return data
+}
+
+type handshakeOpener struct {
+	LongHeaderOpener
+
+	dropInitialKeys func()
+	dropped         bool
+}
+
+func newHandshakeOpener(
+	aead cipher.AEAD,
+	headerProtector headerProtector,
+	dropInitialKeys func(),
+	perspective protocol.Perspective,
+) LongHeaderOpener {
+	opener := newLongHeaderOpener(aead, headerProtector)
+	// The server drops Initial keys when first successfully processing a Handshake packet.
+	if perspective == protocol.PerspectiveClient {
+		return opener
+	}
+	return &handshakeOpener{
+		LongHeaderOpener: opener,
+		dropInitialKeys:  dropInitialKeys,
+	}
+}
+
+func (o *handshakeOpener) Open(dst, src []byte, pn protocol.PacketNumber, ad []byte) ([]byte, error) {
+	dec, err := o.LongHeaderOpener.Open(dst, src, pn, ad)
+	if err == nil && !o.dropped {
+		o.dropInitialKeys()
+		o.dropped = true
+	}
+	return dec, err
+}

internal/handshake/aead_test.go

@@ -6,11 +6,13 @@ import (
 	"crypto/rand"
 	"fmt"
 
+	"github.com/lucas-clemente/quic-go/internal/protocol"
+
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
 )
 
-var _ = Describe("AEAD", func() {
+var _ = Describe("Long Header AEAD", func() {
 	for i := range cipherSuites {
 		cs := cipherSuites[i]
 
@@ -25,8 +27,8 @@ var _ = Describe("AEAD", func() {
 				aead, err := cipher.NewGCM(block)
 				Expect(err).ToNot(HaveOccurred())
 
-				return newLongHeaderSealer(aead, newHeaderProtector(cs, key, true)),
+				return newLongHeaderSealer(aead, newHeaderProtector(cs, hpKey, true)),
-					newLongHeaderOpener(aead, newHeaderProtector(cs, key, true))
+					newLongHeaderOpener(aead, newHeaderProtector(cs, hpKey, true))
 			}
 
 			Context("message encryption", func() {
@@ -91,3 +93,70 @@ var _ = Describe("AEAD", func() {
 		})
 	}
 })
+
+var _ = Describe("Long Header AEAD", func() {
+	var (
+		dropped chan struct{} // use a chan because closing it twice will panic
+		aead    cipher.AEAD
+		hp      headerProtector
+	)
+	dropCb := func() { close(dropped) }
+	msg := []byte("Lorem ipsum dolor sit amet.")
+	ad := []byte("Donec in velit neque.")
+
+	BeforeEach(func() {
+		dropped = make(chan struct{})
+		key := make([]byte, 16)
+		hpKey := make([]byte, 16)
+		rand.Read(key)
+		rand.Read(hpKey)
+		block, err := aes.NewCipher(key)
+		Expect(err).ToNot(HaveOccurred())
+		aead, err = cipher.NewGCM(block)
+		Expect(err).ToNot(HaveOccurred())
+		hp = newHeaderProtector(cipherSuites[0], hpKey, true)
+	})
+
+	Context("for the server", func() {
+		It("drops keys when first successfully processing a Handshake packet", func() {
+			serverOpener := newHandshakeOpener(aead, hp, dropCb, protocol.PerspectiveServer)
+			// first try to open an invalid message
+			_, err := serverOpener.Open(nil, []byte("invalid"), 0, []byte("invalid"))
+			Expect(err).To(HaveOccurred())
+			Expect(dropped).ToNot(BeClosed())
+			// then open a valid message
+			enc := newLongHeaderSealer(aead, hp).Seal(nil, msg, 10, ad)
+			_, err = serverOpener.Open(nil, enc, 10, ad)
+			Expect(err).ToNot(HaveOccurred())
+			Expect(dropped).To(BeClosed())
+			// now open the same message again to make sure the callback is only called once
+			_, err = serverOpener.Open(nil, enc, 10, ad)
+			Expect(err).ToNot(HaveOccurred())
+		})
+
+		It("doesn't drop keys when sealing a Handshake packet", func() {
+			serverSealer := newHandshakeSealer(aead, hp, dropCb, protocol.PerspectiveServer)
+			serverSealer.Seal(nil, msg, 1, ad)
+			Expect(dropped).ToNot(BeClosed())
+		})
+	})
+
+	Context("for the client", func() {
+		It("drops keys when first sealing a Handshake packet", func() {
+			clientSealer := newHandshakeSealer(aead, hp, dropCb, protocol.PerspectiveClient)
+			// seal the first message
+			clientSealer.Seal(nil, msg, 1, ad)
+			Expect(dropped).To(BeClosed())
+			// seal another message to make sure the callback is only called once
+			clientSealer.Seal(nil, msg, 2, ad)
+		})
+
+		It("doesn't drop keys when processing a Handshake packet", func() {
+			enc := newLongHeaderSealer(aead, hp).Seal(nil, msg, 42, ad)
+			clientOpener := newHandshakeOpener(aead, hp, dropCb, protocol.PerspectiveClient)
+			_, err := clientOpener.Open(nil, enc, 42, ad)
+			Expect(err).ToNot(HaveOccurred())
+			Expect(dropped).ToNot(BeClosed())
+		})
+	})
+})

internal/handshake/crypto_setup.go

@@ -221,14 +221,6 @@ func (h *cryptoSetup) ChangeConnectionID(id protocol.ConnectionID) {
 
 func (h *cryptoSetup) SetLargest1RTTAcked(pn protocol.PacketNumber) {
 	h.aead.SetLargestAcked(pn)
-	// drop initial keys
-	// TODO: do this earlier
-	if h.initialOpener != nil {
-		h.initialOpener = nil
-		h.initialSealer = nil
-		h.runner.DropKeys(protocol.EncryptionInitial)
-		h.logger.Debugf("Dropping Initial keys.")
-	}
 	// drop handshake keys
 	if h.handshakeOpener != nil {
 		h.handshakeOpener = nil
@@ -483,9 +475,11 @@ func (h *cryptoSetup) SetReadKey(encLevel qtls.EncryptionLevel, suite *qtls.Ciph
 	switch encLevel {
 	case qtls.EncryptionHandshake:
 		h.readEncLevel = protocol.EncryptionHandshake
-		h.handshakeOpener = newLongHeaderOpener(
+		h.handshakeOpener = newHandshakeOpener(
 			createAEAD(suite, trafficSecret),
 			newHeaderProtector(suite, trafficSecret, true),
+			h.dropInitialKeys,
+			h.perspective,
 		)
 		h.logger.Debugf("Installed Handshake Read keys (using %s)", cipherSuiteName(suite.ID))
 	case qtls.EncryptionApplication:
@@ -505,9 +499,11 @@ func (h *cryptoSetup) SetWriteKey(encLevel qtls.EncryptionLevel, suite *qtls.Cip
 	switch encLevel {
 	case qtls.EncryptionHandshake:
 		h.writeEncLevel = protocol.EncryptionHandshake
-		h.handshakeSealer = newLongHeaderSealer(
+		h.handshakeSealer = newHandshakeSealer(
 			createAEAD(suite, trafficSecret),
 			newHeaderProtector(suite, trafficSecret, true),
+			h.dropInitialKeys,
+			h.perspective,
 		)
 		h.logger.Debugf("Installed Handshake Write keys (using %s)", cipherSuiteName(suite.ID))
 	case qtls.EncryptionApplication:
@@ -552,6 +548,16 @@ func (h *cryptoSetup) SendAlert(alert uint8) {
 	h.alertChan <- alert
 }
 
+// used a callback in the handshakeSealer and handshakeOpener
+func (h *cryptoSetup) dropInitialKeys() {
+	h.mutex.Lock()
+	h.initialOpener = nil
+	h.initialSealer = nil
+	h.mutex.Unlock()
+	h.runner.DropKeys(protocol.EncryptionInitial)
+	h.logger.Debugf("Dropping Initial keys.")
+}
+
 func (h *cryptoSetup) GetInitialSealer() (LongHeaderSealer, error) {
 	h.mutex.Lock()
 	defer h.mutex.Unlock()