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 slice: &[u8] = buffer.into();
+                {
+                    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 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..]
+                            );
 
-                // Update TLS session
-                self.session.borrow_mut().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,78 +396,113 @@ 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();
+                    
+                    let mut tls_repr_vec: Vec<(&[u8], TlsRepr)> = Vec::new();
+                    let mut bytes = &buffer[..buffer_size];
 
-        // Take the TLS representation out of the vector,
-        // Process as a queue
-        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);
+                    // 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, ())
+                        };
+                    }
 
-            // 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
-                let mut app_data = repr.payload.take().unwrap();
-                let mut associated_data = [0; 5];
-                associated_data[0] = repr.content_type.into();
-                NetworkEndian::write_u16(
-                    &mut associated_data[1..3],
-                    repr.version.into()
-                );
-                NetworkEndian::write_u16(
-                    &mut associated_data[3..5],
-                    repr.length
-                );
-                {
-                    let mut session = self.session.borrow_mut();
-                    session.decrypt_in_place_detached(
-                        &associated_data,
-                        &mut app_data
-                    ).unwrap();
-                    session.increment_server_sequence_number();
-                }
+                    // 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);
 
-                // Discard last 16 bytes (auth tag)
-                let inner_plaintext = &app_data[..app_data.len()-16];
-                let (inner_content_type, _) = get_content_type_inner_plaintext(
-                    inner_plaintext
-                );
-                if inner_content_type != TlsContentType::Handshake {
-                    // Silently ignore non-handshakes
-                    continue;
-                }
-                let (_, mut inner_handshakes) = complete(
-                    parse_inner_plaintext_for_handshake
-                )(inner_plaintext).unwrap();
+                        // Process record base on content type
+                        log::info!("Record type: {:?}", repr.content_type);
 
-                // Sequentially process all handshakes
-                let num_of_handshakes = inner_handshakes.len();
-                for _ in 0..num_of_handshakes {
-                    let (handshake_slice, handshake_repr) = inner_handshakes.remove(0);
-                    self.process(
-                        handshake_slice,
-                        TlsRepr {
-                            content_type: TlsContentType::Handshake,
-                            version: repr.version,
-                            length: u16::try_from(handshake_repr.length).unwrap() + 4,
-                            payload: None,
-                            handshake: Some(handshake_repr)
+                        if repr.content_type == TlsContentType::ApplicationData {
+                            log::info!("Found application data");
+                            // Take the payload out of TLS Record and decrypt
+                            let mut app_data = repr.payload.take().unwrap();
+                            let mut associated_data = [0; 5];
+                            associated_data[0] = repr.content_type.into();
+                            NetworkEndian::write_u16(
+                                &mut associated_data[1..3],
+                                repr.version.into()
+                            );
+                            NetworkEndian::write_u16(
+                                &mut associated_data[3..5],
+                                repr.length
+                            );
+                            {
+                                let mut session = self.session.borrow_mut();
+                                session.decrypt_in_place_detached(
+                                    &associated_data,
+                                    &mut app_data
+                                ).unwrap();
+                                session.increment_server_sequence_number();
+                            }
+
+                            // Discard last 16 bytes (auth tag)
+                            let inner_plaintext = &app_data[..app_data.len()-16];
+                            let (inner_content_type, _) = get_content_type_inner_plaintext(
+                                inner_plaintext
+                            );
+                            if inner_content_type != TlsContentType::Handshake {
+                                // Silently ignore non-handshakes
+                                continue;
+                            }
+                            let (_, mut inner_handshakes) = complete(
+                                parse_inner_plaintext_for_handshake
+                            )(inner_plaintext).unwrap();
+
+                            // Sequentially process all handshakes
+                            let num_of_handshakes = inner_handshakes.len();
+                            for _ in 0..num_of_handshakes {
+                                let (handshake_slice, handshake_repr) = inner_handshakes.remove(0);
+                                if self.process(
+                                    handshake_slice,
+                                    TlsRepr {
+                                        content_type: TlsContentType::Handshake,
+                                        version: repr.version,
+                                        length: u16::try_from(handshake_repr.length).unwrap() + 4,
+                                        payload: None,
+                                        handshake: Some(handshake_repr)
+                                    }
+                                ).is_err() {
+                                    return (buffer_size, ())
+                                }
+                            }
                         }
-                    )?;
+                        
+                        else {
+                            if self.process(repr_slice, repr).is_err() {
+                                return (buffer_size, ())
+                            }
+                            log::info!("Processed record");
+                        }
+                    }
+                    (buffer_size, ())
                 }
-            }
-
-            
-            else {
-                self.process(repr_slice, repr)?;
-                log::info!("Processed record");
-            }
+            )?;
         }
 
         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