From e07835b954d6c8653b488e82c167b09cdf594573 Mon Sep 17 00:00:00 2001
From: Dirkjan Ochtman <dirkjan@ochtman.nl>
Date: Wed, 31 Mar 2021 10:50:06 +0200
Subject: [PATCH] crypto: return Option from next_1rtt_keys()

A crypto session cannot return 1-RTT keys before handshaking is finished.
Currently, rustls contains unwrap() calls to deal with this, but hoisting
out the unwrap to the `Connection` level here seems more sensible.

We can adjust the call into rustls once we update rustls for other reasons.
---
 quinn-proto/src/connection/mod.rs | 12 ++++++++++--
 quinn-proto/src/crypto.rs         |  2 +-
 quinn-proto/src/crypto/rustls.rs  |  6 +++---
 3 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/quinn-proto/src/connection/mod.rs b/quinn-proto/src/connection/mod.rs
index eca602c77..c7d94c852 100644
--- a/quinn-proto/src/connection/mod.rs
+++ b/quinn-proto/src/connection/mod.rs
@@ -1591,8 +1591,13 @@ where
         trace!("{:?} keys ready", space);
         if space == SpaceId::Data {
             // Precompute the first key update
-            self.next_crypto = Some(self.crypto.next_1rtt_keys());
+            self.next_crypto = Some(
+                self.crypto
+                    .next_1rtt_keys()
+                    .expect("handshake should be complete"),
+            );
         }
+
         self.spaces[space].crypto = Some(crypto);
         debug_assert!(space as usize > self.highest_space as usize);
         self.highest_space = space;
@@ -2788,7 +2793,10 @@ where
         // Generate keys for the key phase after the one we're switching to, store them in
         // `next_crypto`, make the contents of `next_crypto` current, and move the current keys into
         // `prev_crypto`.
-        let new = self.crypto.next_1rtt_keys();
+        let new = self
+            .crypto
+            .next_1rtt_keys()
+            .expect("only called for `Data` packets");
         let old = mem::replace(
             &mut self.spaces[SpaceId::Data]
                 .crypto
diff --git a/quinn-proto/src/crypto.rs b/quinn-proto/src/crypto.rs
index 8bee8a231..1a84fd71b 100644
--- a/quinn-proto/src/crypto.rs
+++ b/quinn-proto/src/crypto.rs
@@ -94,7 +94,7 @@ pub trait Session: Send + Sized {
     fn write_handshake(&mut self, buf: &mut Vec<u8>) -> Option<Keys<Self>>;
 
     /// Compute keys for the next key update
-    fn next_1rtt_keys(&mut self) -> KeyPair<Self::PacketKey>;
+    fn next_1rtt_keys(&mut self) -> Option<KeyPair<Self::PacketKey>>;
 
     /// Generate the integrity tag for a retry packet
     fn retry_tag(orig_dst_cid: &ConnectionId, packet: &[u8]) -> [u8; 16];
diff --git a/quinn-proto/src/crypto/rustls.rs b/quinn-proto/src/crypto/rustls.rs
index ee900d205..654be751e 100644
--- a/quinn-proto/src/crypto/rustls.rs
+++ b/quinn-proto/src/crypto/rustls.rs
@@ -170,12 +170,12 @@ impl crypto::Session for TlsSession {
         })
     }
 
-    fn next_1rtt_keys(&mut self) -> KeyPair<Self::PacketKey> {
+    fn next_1rtt_keys(&mut self) -> Option<KeyPair<Self::PacketKey>> {
         let keys = (**self).next_1rtt_keys();
-        KeyPair {
+        Some(KeyPair {
             local: keys.local,
             remote: keys.remote,
-        }
+        })
     }
 
     fn retry_tag(orig_dst_cid: &ConnectionId, packet: &[u8]) -> [u8; 16] {