From 9327784e99641961eb5264e2dbaeffec25ff3fac Mon Sep 17 00:00:00 2001
From: Dirkjan Ochtman <dirkjan@ochtman.nl>
Date: Wed, 5 Sep 2018 22:00:00 +0200
Subject: [PATCH] Move transmit_handshake() into Connection type

---
 src/connection.rs | 22 ++++++++++++++++++++++
 src/endpoint.rs   | 42 +++++++++++-------------------------------
 2 files changed, 33 insertions(+), 31 deletions(-)

diff --git a/src/connection.rs b/src/connection.rs
index dc656b61f..b9a292bf2 100644
--- a/src/connection.rs
+++ b/src/connection.rs
@@ -731,6 +731,28 @@ impl Connection {
         }
     }
 
+    pub fn transmit_handshake(&mut self, messages: &[u8]) {
+        let offset = {
+            let ss = self
+                .streams
+                .get_mut(&StreamId(0))
+                .unwrap()
+                .send_mut()
+                .unwrap();
+            let x = ss.offset;
+            ss.offset += messages.len() as u64;
+            ss.bytes_in_flight += messages.len() as u64;
+            x
+        };
+        self.handshake_pending.stream.push_back(frame::Stream {
+            id: StreamId(0),
+            fin: false,
+            offset,
+            data: messages.into(),
+        });
+        self.awaiting_handshake = true;
+    }
+
     pub fn transmit(&mut self, stream: StreamId, data: Bytes) {
         let ss = self.streams.get_mut(&stream).unwrap().send_mut().unwrap();
         assert_eq!(ss.state, stream::SendState::Ready);
diff --git a/src/endpoint.rs b/src/endpoint.rs
index 69866b73e..5916098d0 100644
--- a/src/endpoint.rs
+++ b/src/endpoint.rs
@@ -629,7 +629,7 @@ impl Endpoint {
             Err(HandshakeError::WouldBlock(tls)) => tls,
             Err(e) => panic!("unexpected TLS error: {}", e),
         };
-        self.transmit_handshake(conn, &tls.get_mut().take_outgoing());
+        self.connections[conn.0].transmit_handshake(&tls.get_mut().take_outgoing());
         self.connections[conn.0].state = Some(State::Handshake(state::Handshake {
             tls,
             clienthello_packet: None,
@@ -842,7 +842,7 @@ impl Endpoint {
                     self.connection_ids_initial.insert(dest_id, conn);
                     self.connections[conn.0].zero_rtt_crypto = zero_rtt_crypto;
                     self.connections[conn.0].on_packet_authenticated(now, packet_number as u64);
-                    self.transmit_handshake(conn, &tls.get_mut().take_outgoing());
+                    self.connections[conn.0].transmit_handshake(&tls.get_mut().take_outgoing());
                     self.connections[conn.0].state = Some(State::Handshake(state::Handshake {
                         tls,
                         clienthello_packet: None,
@@ -1429,12 +1429,16 @@ impl Endpoint {
                                         local_id,
                                         remote_id,
                                         remote,
-                                        self.ctx.initial_packet_number.sample(&mut self.ctx.rng).into(),
+                                        self.ctx
+                                            .initial_packet_number
+                                            .sample(&mut self.ctx.rng)
+                                            .into(),
                                         Side::Client,
                                         &self.ctx.config,
                                     );
                                     // Send updated ClientHello
-                                    self.transmit_handshake(conn, &tls.get_mut().take_outgoing());
+                                    self.connections[conn.0]
+                                        .transmit_handshake(&tls.get_mut().take_outgoing());
                                     // Prepare to receive Handshake packets that start stream 0 from offset 0
                                     tls.get_mut().reset_read();
                                     State::Handshake(state::Handshake {
@@ -1619,7 +1623,8 @@ impl Endpoint {
                                 );
                                 self.connections[conn.0].handshake_cleanup(&self.ctx.config);
                                 if self.connections[conn.0].side == Side::Client {
-                                    self.transmit_handshake(conn, &tls.get_mut().take_outgoing());
+                                    self.connections[conn.0]
+                                        .transmit_handshake(&tls.get_mut().take_outgoing());
                                 } else {
                                     self.connections[conn.0].transmit(
                                         StreamId(0),
@@ -1673,7 +1678,7 @@ impl Endpoint {
                                 {
                                     let response = tls.get_mut().take_outgoing();
                                     if !response.is_empty() {
-                                        self.transmit_handshake(conn, &response);
+                                        self.connections[conn.0].transmit_handshake(&response);
                                     }
                                 }
                                 State::Handshake(state::Handshake {
@@ -2122,31 +2127,6 @@ impl Endpoint {
         }
     }
 
-    fn transmit_handshake(&mut self, conn: ConnectionHandle, messages: &[u8]) {
-        let offset = {
-            let ss = self.connections[conn.0]
-                .streams
-                .get_mut(&StreamId(0))
-                .unwrap()
-                .send_mut()
-                .unwrap();
-            let x = ss.offset;
-            ss.offset += messages.len() as u64;
-            ss.bytes_in_flight += messages.len() as u64;
-            x
-        };
-        self.connections[conn.0]
-            .handshake_pending
-            .stream
-            .push_back(frame::Stream {
-                id: StreamId(0),
-                fin: false,
-                offset,
-                data: messages.into(),
-            });
-        self.connections[conn.0].awaiting_handshake = true;
-    }
-
     fn on_ack_received(&mut self, now: u64, conn: ConnectionHandle, ack: frame::Ack) {
         trace!(self.ctx.log, "got ack"; "ranges" => ?ack.iter().collect::<Vec<_>>());
         let was_blocked = self.connections[conn.0].blocked();