tls: reduce redundant array alloc

src/lib.rs

@@ -46,7 +46,6 @@ use net::iface::EthernetInterface;
 use net::time::Instant;
 use net::phy::Device;
 
-use crate::tls::TlsSocket;
 use crate::set::TlsSocketSet;
 
 // One-call function for polling all sockets within socket set

src/session.rs

@@ -10,6 +10,7 @@ use generic_array::GenericArray;
 use byteorder::{ByteOrder, NetworkEndian};
 use rsa::{RSAPublicKey, PublicKey, PaddingScheme, Hash as RSAHash};
 use hmac::{ Hmac, Mac, NewMac };
+use smoltcp::wire::IpEndpoint;
 
 use crate::tls::TlsState;
 use crate::tls_packet::CipherSuite;
@@ -25,6 +26,12 @@ type Aes128Ccm = Ccm<Aes128, U16, U12>;
 pub(crate) struct Session<'a> {
     state: TlsState,
     role: TlsRole,
+    // Local and remote endpoints of the socket
+    // TCP socket does store these 2 information and it is gettable
+    // However, upon invokation of `reset`, these endpoints are wiped out
+    // Should TLS socket requires TCP socket to restart, we need these info
+    local_endpoint: IpEndpoint,
+    remote_endpoint: IpEndpoint,
     // Session ID for this session
     session_id: Option<[u8; 32]>,
     // Changed cipher spec
@@ -84,6 +91,8 @@ impl<'a> Session<'a> {
         Self {
             state: TlsState::START,
             role,
+            local_endpoint: IpEndpoint::default(),
+            remote_endpoint: IpEndpoint::default(),
             session_id: None,
             changed_cipher_spec: false,
             latest_secret: None,
@@ -110,6 +119,17 @@ impl<'a> Session<'a> {
         }
     }
 
+    pub(crate) fn connect(
+        &mut self,
+        remote_endpoint: IpEndpoint,
+        local_endpoint: IpEndpoint
+    ) {
+        self.role = TlsRole::Client;
+        self.state = TlsState::START;
+        self.local_endpoint = local_endpoint;
+        self.remote_endpoint = remote_endpoint;
+    }
+
     // State transition from START to WAIT_SH
     pub(crate) fn client_update_for_ch(
         &mut self,
@@ -129,8 +149,6 @@ impl<'a> Session<'a> {
     }
 
     // State transition from WAIT_SH to WAIT_EE
-    // TODO: Memory allocation
-    // It current dumps too much memory onto the stack on invocation
     pub(crate) fn client_update_for_sh(
         &mut self,
         cipher_suite: CipherSuite,
@@ -1143,6 +1161,14 @@ impl<'a> Session<'a> {
         self.state
     }
 
+    pub(crate) fn get_local_endpoint(&self) -> IpEndpoint {
+        self.local_endpoint
+    }
+
+    pub(crate) fn get_remote_endpoint(&self) -> IpEndpoint {
+        self.remote_endpoint
+    }
+
     pub(crate) fn has_completed_handshake(&self) -> bool {
         self.state == TlsState::CONNECTED
     }
@@ -1616,10 +1642,6 @@ impl<'a> Session<'a> {
     pub(crate) fn get_session_role(&self) -> TlsRole {
         self.role
     }
-
-    pub(crate) fn becomes_client(&mut self) {
-        self.role = TlsRole::Client;
-    }
 }
 
 #[derive(Debug, PartialEq, Eq, Clone, Copy)]

src/set.rs

@@ -2,21 +2,7 @@ use smoltcp as net;
 
 use managed::ManagedSlice;
 use crate::tls::TlsSocket;
-use net::socket::SocketSetItem;
 use net::socket::SocketSet;
-use net::socket::SocketHandle;
-use net::socket::Socket;
-use net::socket::TcpSocket;
-use net::socket::AnySocket;
-use net::socket::SocketRef;
-use net::iface::EthernetInterface;
-use net::time::Instant;
-use net::phy::Device;
-
-use core::convert::From;
-use core::cell::RefCell;
-
-use alloc::vec::Vec;
 
 pub struct TlsSocketSet<'a> {
     tls_sockets: ManagedSlice<'a, Option<TlsSocket<'a>>>
@@ -68,11 +54,9 @@ impl<'a> TlsSocketSet<'a> {
     {
         for socket in self.tls_sockets.iter_mut() {
             if socket.is_some() {
-                log::info!("Found TLS");
                 socket.as_mut()
                     .unwrap()
                     .update_handshake(sockets)?;
-                log::info!("Updated TLS");
             }
         }
 

src/tls.rs

@@ -3,13 +3,9 @@ use smoltcp::socket::TcpState;
 use smoltcp::socket::SocketHandle;
 use smoltcp::socket::SocketSet;
 use smoltcp::socket::TcpSocketBuffer;
-use smoltcp::socket::SocketRef;
 use smoltcp::wire::IpEndpoint;
 use smoltcp::Result;
 use smoltcp::Error;
-use smoltcp::iface::EthernetInterface;
-use smoltcp::time::Instant;
-use smoltcp::phy::Device;
 
 use byteorder::{ByteOrder, NetworkEndian};
 use generic_array::GenericArray;
@@ -18,7 +14,6 @@ use core::convert::TryFrom;
 use core::convert::TryInto;
 use core::cell::RefCell;
 
-use rand_core::{RngCore, CryptoRng};
 use p256::{EncodedPoint, ecdh::EphemeralSecret};
 use ccm::consts::*;
 
@@ -117,31 +112,53 @@ impl<'s> TlsSocket<'s> {
 
         // Permit TLS handshake as well
         let mut session = self.session.borrow_mut();
-        session.becomes_client();
+        session.connect(
+            tcp_socket.remote_endpoint(),
+            tcp_socket.local_endpoint()
+        );
         Ok(())
     }
 
     pub fn update_handshake(&mut self, sockets: &mut SocketSet) -> Result<bool> {
-        // Check TCP socket
-        {
-            let mut tcp_socket = sockets.get::<TcpSocket>(self.tcp_handle);
-            tcp_socket.set_keep_alive(Some(smoltcp::time::Duration::from_millis(1000)));
-            if tcp_socket.state() != TcpState::Established {
-                log::info!("TCP not established");
-                return Ok(false);
-            }
-        }
-        // Check TLS session state
-        {
-            let role = self.session.borrow().get_session_role();
-            if role != crate::session::TlsRole::Client {
-                return Ok(true);
-            }
-        }
         // Handle TLS handshake through TLS states
         let tls_state = {
             self.session.borrow().get_tls_state()
         };
+
+        // Check TCP socket/ TLS session
+        {
+            let mut tcp_socket = sockets.get::<TcpSocket>(self.tcp_handle);
+            let tls_socket = self.session.borrow();
+
+            // Check if it should connect to client or not
+            if tls_socket.get_session_role() != crate::session::TlsRole::Client {
+                // Return true for no need to do anymore handshake
+                return Ok(true);
+            }
+
+            // Skip handshake processing if it is already completed
+            // However, redo TCP handshake if TLS socket is trying to connect and
+            // TCP socket is not connected
+            if tcp_socket.state() != TcpState::Established {
+                if tls_state == TlsState::START {
+                    // Restart TCP handshake is it is closed for some reason
+                    if !tcp_socket.is_open() {
+                        tcp_socket.connect(
+                            tls_socket.get_remote_endpoint(),
+                            tls_socket.get_local_endpoint()
+                        )?;                        
+                    }
+                    return Ok(false);
+                } else {
+                    // Do nothing, either handshake failed or the socket closed
+                    // after finishing the handshake
+                    return Ok(false);                    
+                }
+
+            }
+        }
+
+        // Handle TLS handshake through TLS states
         match tls_state {
             // Initiate TLS handshake
             TlsState::START => {
@@ -156,22 +173,34 @@ impl<'s> TlsSocket<'s> {
                 let repr = TlsRepr::new()
                     .client_hello(&ecdh_secret, &x25519_secret, random, session_id.clone());
 
-                // Update hash function with client hello handshake
-                let mut array = [0; 512];
-                let mut buffer = TlsBuffer::new(&mut array);
-                buffer.enqueue_tls_repr(repr)?;
+                {
+                    let mut tcp_socket = sockets.get::<TcpSocket>(self.tcp_handle);
+                    tcp_socket.send(
+                        |data| {
+                            // Enqueue tls representation without extra allocation
+                            let mut buffer = TlsBuffer::new(data);
+                            if buffer.enqueue_tls_repr(repr).is_err() {
+                                return (0, ())
+                            }
                             let slice: &[u8] = buffer.into();
 
-                // Send the packet
-                self.send_tls_slice(sockets, slice)?;
-
-                // Update TLS session
-                self.session.borrow_mut().client_update_for_ch(
+                            // Update the session
+                            // No sequence number calculation in CH
+                            // because there is no encryption
+                            // Still, data needs to be hashed
+                            let mut session = self.session.borrow_mut();
+                            session.client_update_for_ch(
                                 ecdh_secret,
                                 x25519_secret,
                                 session_id,
                                 &slice[5..]
                             );
+
+                            // Finally send the data
+                            (slice.len(), ())
+                        }
+                    )?;
+                }
             },
 
             // TLS Client wait for Server Hello
@@ -367,20 +396,48 @@ impl<'s> TlsSocket<'s> {
             }
         }
 
+
+
         // Read for TLS packet
         // Proposition: Decouple all data from TLS record layer before processing
         //      Recouple a brand new TLS record wrapper
-        let mut array: [u8; 2048] = [0; 2048];
-        let mut tls_repr_vec = self.recv_tls_repr(sockets, &mut array)?;
+        // Use recv to avoid buffer allocation
+        {
+            let mut tcp_socket = sockets.get::<TcpSocket>(self.tcp_handle);
+            tcp_socket.recv(
+                |buffer| {
+                    let buffer_size = buffer.len();
                     
-        // Take the TLS representation out of the vector,
-        // Process as a queue
+                    let mut tls_repr_vec: Vec<(&[u8], TlsRepr)> = Vec::new();
+                    let mut bytes = &buffer[..buffer_size];
+
+                    // Sequentially push reprs into vec
+                    loop {
+                        match parse_tls_repr(bytes) {
+                            Ok((rest, (repr_slice, repr))) => {
+                                tls_repr_vec.push(
+                                    (repr_slice, repr)
+                                );
+                                if rest.len() == 0 {
+                                    break;
+                                } else {
+                                    bytes = rest;
+                                }
+                            },
+                            // Dequeue everything and abort processing if it is malformed
+                            _ => return (buffer_size, ())
+                        };
+                    }
+
+                    // Sequencially process the representations in vector
+                    // Decrypt and split the handshake if necessary
                     let tls_repr_vec_size = tls_repr_vec.len();
                     for _index in 0..tls_repr_vec_size {
                         let (repr_slice, mut repr) = tls_repr_vec.remove(0);
 
                         // Process record base on content type
                         log::info!("Record type: {:?}", repr.content_type);
+
                         if repr.content_type == TlsContentType::ApplicationData {
                             log::info!("Found application data");
                             // Take the payload out of TLS Record and decrypt
@@ -421,7 +478,7 @@ impl<'s> TlsSocket<'s> {
                             let num_of_handshakes = inner_handshakes.len();
                             for _ in 0..num_of_handshakes {
                                 let (handshake_slice, handshake_repr) = inner_handshakes.remove(0);
-                    self.process(
+                                if self.process(
                                     handshake_slice,
                                     TlsRepr {
                                         content_type: TlsContentType::Handshake,
@@ -430,16 +487,23 @@ impl<'s> TlsSocket<'s> {
                                         payload: None,
                                         handshake: Some(handshake_repr)
                                     }
-                    )?;
+                                ).is_err() {
+                                    return (buffer_size, ())
+                                }
                             }
                         }
-
                         
                         else {
-                self.process(repr_slice, repr)?;
+                            if self.process(repr_slice, repr).is_err() {
+                                return (buffer_size, ())
+                            }
                             log::info!("Processed record");
                         }
                     }
+                    (buffer_size, ())
+                }
+            )?;
+        }
 
         Ok(self.session.borrow().has_completed_handshake())
     }
@@ -804,28 +868,6 @@ impl<'s> TlsSocket<'s> {
         Ok(())
     }
 
-    // Generic inner send method for buffer IO, through TCP socket
-    // Usage: Push a slice representation of ONE TLS packet
-    // This function will only increment sequence number by 1
-    // Repeatedly call this function if sending multiple TLS packets is needed
-    fn send_tls_slice(&self, sockets: &mut SocketSet, slice: &[u8]) -> Result<()> {
-        let mut tcp_socket = sockets.get::<TcpSocket>(self.tcp_handle);
-        if !tcp_socket.can_send() {
-            return Err(Error::Illegal);
-        }
-        let buffer_size = slice.len();
-        tcp_socket.send_slice(slice)
-            .and_then(
-                |size| if size == buffer_size {
-                    Ok(())
-                } else {
-                    Err(Error::Truncated)
-                }
-            )?;
-        self.session.borrow_mut().increment_client_sequence_number();
-        Ok(())
-    }
-
     // Send method for TLS Handshake that needs to be encrypted.
     // Does the following things:
     // 1. Encryption
@@ -872,46 +914,21 @@ impl<'s> TlsSocket<'s> {
         Ok(())
     }
 
-    // Generic inner recv method, through TCP socket
-    // A TCP packet can contain multiple TLS records (including 0)
-    // Therefore, sequence nubmer incrementation is not completed here
-    fn recv_tls_repr<'a>(&'a self, sockets: &mut SocketSet, byte_array: &'a mut [u8]) -> Result<Vec<(&[u8], TlsRepr)>> {
-        let mut tcp_socket = sockets.get::<TcpSocket>(self.tcp_handle);
-        if !tcp_socket.can_recv() {
-            return Ok(Vec::new());
-        }
-        let array_size = tcp_socket.recv_slice(byte_array)?;
-        let mut vec: Vec<(&[u8], TlsRepr)> = Vec::new();
-        let mut bytes: &[u8] = &byte_array[..array_size];
-        loop {
-            match parse_tls_repr(bytes) {
-                Ok((rest, (repr_slice, repr))) => {
-                    vec.push(
-                        (repr_slice, repr)
-                    );
-                    if rest.len() == 0 {
-                        return Ok(vec);
-                    } else {
-                        bytes = rest;
-                    }
-                },
-                _ => return Err(Error::Unrecognized),
-            };
-        }
-    }
-
     pub fn recv_slice(&self, sockets: &mut SocketSet, data: &mut [u8]) -> Result<usize> {
         let mut tcp_socket = sockets.get::<TcpSocket>(self.tcp_handle);
         if !tcp_socket.can_recv() {
             return Ok(0);
         }
 
-        // TODO: Use `recv` to receive instead
-        // Issue with using recv slice:
-        // Encrypted application data can cramp together into a TCP Segment
-        // Dequeuing all bytes from the buffer immediately can cause
-        // 1. Incorrect decryption, hence throwing error, and
-        // 2. sequence number to go out of sync forever
+        let mut session = self.session.borrow_mut();
+
+        // If the handshake is not completed, do not pull bytes out of the buffer
+        // through TlsSocket.recv_slice()
+        // Handshake recv should be through TCPSocket directly.
+        if session.get_tls_state() != TlsState::CONNECTED {
+            return Ok(0);
+        }
+
         let (recv_slice_size, acceptable) = tcp_socket.recv(
             |buffer| {
                 // Read the size of the TLS record beforehand
@@ -936,8 +953,6 @@ impl<'s> TlsSocket<'s> {
             return Ok(0);
         }
 
-        // let recv_slice_size = tcp_socket.recv_slice(data)?;
-
         // Encrypted data need a TLS record wrapper (5 bytes)
         // Authentication tag (16 bytes, for all supported AEADs)
         // Content type byte (1 byte)
@@ -947,10 +962,8 @@ impl<'s> TlsSocket<'s> {
         }
         
         // Get Associated Data
-        let mut session = self.session.borrow_mut();
         let mut associated_data: [u8; 5] = [0; 5];
         associated_data.clone_from_slice(&data[..5]);
-        // log::info!("Received encrypted appdata: {:?}", &data[..recv_slice_size]);
 
         // Dump association data (TLS Record wrapper)
         // Only decrypt application data
@@ -968,7 +981,7 @@ impl<'s> TlsSocket<'s> {
 
         // If it is not application data, handle it internally
         if content_type != TlsContentType::ApplicationData {
-            // TODO:: Implement key update
+            // TODO: Implement key update here, as it could be a key update
             log::info!("Other decrypted: {:?}", &data[..(recv_slice_size-16)]);
             return Ok(0);
         }
@@ -991,6 +1004,14 @@ impl<'s> TlsSocket<'s> {
     }
 
     pub fn send_slice(&self, sockets: &mut SocketSet, data: &[u8]) -> Result<()> {
+        // If the handshake is not completed, do not push bytes onto the buffer
+        // through TlsSocket.send_slice()
+        // Handshake send should be through TCPSocket directly.
+        let mut session = self.session.borrow_mut();
+        if session.get_tls_state() != TlsState::CONNECTED {
+            return Ok(());
+        }
+
         // Sending order:
         // 1. Associated data/ TLS Record layer
         // 2. Encrypted { Payload (data) | Content type: Application Data }
@@ -1010,7 +1031,6 @@ impl<'s> TlsSocket<'s> {
         let mut vec: HeaplessVec<u8, U1024> = HeaplessVec::from_slice(data).unwrap();
         vec.push(0x17).unwrap();                // Content type
         
-        let mut session = self.session.borrow_mut();
         let tag = session.encrypt_application_data_in_place_detached(
             &associated_data,
             &mut vec