From 09ac124ae64ef27fd9dc8a79d5c7366ed59ad1bf Mon Sep 17 00:00:00 2001
From: Benjamin Saunders <ben.e.saunders@gmail.com>
Date: Sun, 3 May 2020 16:35:15 -0700
Subject: [PATCH] Pass header/packet keys together when sensible

---
 quinn-proto/src/connection/mod.rs | 16 ++++++--------
 quinn-proto/src/endpoint.rs       | 36 ++++++++-----------------------
 2 files changed, 16 insertions(+), 36 deletions(-)

diff --git a/quinn-proto/src/connection/mod.rs b/quinn-proto/src/connection/mod.rs
index e8ae941d4..7246a9fde 100644
--- a/quinn-proto/src/connection/mod.rs
+++ b/quinn-proto/src/connection/mod.rs
@@ -2796,17 +2796,15 @@ where
     }
 }
 
-pub fn initial_close<K, H, R>(
-    crypto: &K,
-    header_crypto: &H,
+pub fn initial_close<S, R>(
+    crypto: &Keys<S>,
     remote_id: &ConnectionId,
     local_id: &ConnectionId,
     packet_number: u8,
     reason: R,
 ) -> Box<[u8]>
 where
-    K: crypto::PacketKey,
-    H: crypto::HeaderKey,
+    S: crypto::Session,
     R: Into<Close>,
 {
     let number = PacketNumber::U8(packet_number);
@@ -2819,13 +2817,13 @@ where
 
     let mut buf = Vec::<u8>::new();
     let partial_encode = header.encode(&mut buf);
-    let max_len = MIN_MTU as usize - partial_encode.header_len - crypto.tag_len();
+    let max_len = MIN_MTU as usize - partial_encode.header_len - crypto.packet.local.tag_len();
     reason.into().encode(&mut buf, max_len);
-    buf.resize(buf.len() + crypto.tag_len(), 0);
+    buf.resize(buf.len() + crypto.packet.local.tag_len(), 0);
     partial_encode.finish(
         &mut buf,
-        header_crypto,
-        Some((u64::from(packet_number), crypto)),
+        &crypto.header.local,
+        Some((u64::from(packet_number), &crypto.packet.local)),
     );
     buf.into()
 }
diff --git a/quinn-proto/src/endpoint.rs b/quinn-proto/src/endpoint.rs
index 52e4e4ca4..e7d4a7e5f 100644
--- a/quinn-proto/src/endpoint.rs
+++ b/quinn-proto/src/endpoint.rs
@@ -19,7 +19,7 @@ use crate::{
     config::{ClientConfig, ConfigError, EndpointConfig, ServerConfig},
     connection::{initial_close, Connection, ConnectionError},
     crypto::{
-        self, ClientConfig as ClientCryptoConfig, KeyPair, PacketKey,
+        self, ClientConfig as ClientCryptoConfig, Keys, PacketKey,
         ServerConfig as ServerCryptoConfig,
     },
     packet::{Header, Packet, PacketDecodeError, PartialDecode},
@@ -258,15 +258,7 @@ where
             let crypto = S::initial_keys(&dst_cid, Side::Server);
             return match first_decode.finish(Some(&crypto.header.remote)) {
                 Ok(packet) => self
-                    .handle_first_packet(
-                        now,
-                        remote,
-                        ecn,
-                        packet,
-                        remaining,
-                        &crypto.packet,
-                        &crypto.header,
-                    )
+                    .handle_first_packet(now, remote, ecn, packet, remaining, &crypto)
                     .map(|(ch, conn)| (ch, DatagramEvent::NewConnection(conn))),
                 Err(e) => {
                     trace!("unable to decode initial packet: {}", e);
@@ -457,8 +449,7 @@ where
         ecn: Option<EcnCodepoint>,
         mut packet: Packet,
         rest: Option<BytesMut>,
-        crypto: &KeyPair<S::PacketKey>,
-        header_crypto: &KeyPair<S::HeaderKey>,
+        crypto: &Keys<S>,
     ) -> Option<(ConnectionHandle, Connection<S>)> {
         let (src_cid, dst_cid, token, packet_number) = match packet.header {
             Header::Initial {
@@ -472,6 +463,7 @@ where
         let packet_number = packet_number.expand(0);
 
         if crypto
+            .packet
             .remote
             .decrypt(
                 packet_number as u64,
@@ -502,8 +494,7 @@ where
                 destination: remote,
                 ecn: None,
                 contents: initial_close(
-                    &crypto.local,
-                    &header_crypto.local,
+                    crypto,
                     &src_cid,
                     &temp_loc_cid,
                     0,
@@ -524,8 +515,7 @@ where
                 destination: remote,
                 ecn: None,
                 contents: initial_close(
-                    &crypto.local,
-                    &header_crypto.local,
+                    crypto,
                     &src_cid,
                     &temp_loc_cid,
                     0,
@@ -552,7 +542,7 @@ where
                 let encode = header.encode(&mut buf);
                 buf.put_slice(&token);
                 buf.extend_from_slice(&S::retry_tag(&dst_cid, &buf));
-                encode.finish::<S::PacketKey, S::HeaderKey>(&mut buf, &header_crypto.local, None);
+                encode.finish::<S::PacketKey, S::HeaderKey>(&mut buf, &crypto.header.local, None);
 
                 self.transmits.push_back(Transmit {
                     destination: remote,
@@ -578,8 +568,7 @@ where
                         destination: remote,
                         ecn: None,
                         contents: initial_close(
-                            &crypto.local,
-                            &header_crypto.local,
+                            crypto,
                             &src_cid,
                             &temp_loc_cid,
                             0,
@@ -620,14 +609,7 @@ where
                     self.transmits.push_back(Transmit {
                         destination: remote,
                         ecn: None,
-                        contents: initial_close(
-                            &crypto.local,
-                            &header_crypto.local,
-                            &src_cid,
-                            &temp_loc_cid,
-                            0,
-                            e,
-                        ),
+                        contents: initial_close(crypto, &src_cid, &temp_loc_cid, 0, e),
                     });
                 }
                 None