diff --git a/src/session.rs b/src/session.rs index 68b6eb6..3e999e4 100644 --- a/src/session.rs +++ b/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; 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, @@ -1143,6 +1163,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 +1644,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)] diff --git a/src/set.rs b/src/set.rs index 79a1d81..89a0c80 100644 --- a/src/set.rs +++ b/src/set.rs @@ -68,11 +68,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"); } } diff --git a/src/tls.rs b/src/tls.rs index 16b2213..a169bbc 100644 --- a/src/tls.rs +++ b/src/tls.rs @@ -117,31 +117,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 { - // Check TCP socket - { - let mut tcp_socket = sockets.get::(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::(self.tcp_handle); + let mut 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 => { @@ -906,6 +928,15 @@ impl<'s> TlsSocket<'s> { return Ok(0); } + 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); + } + // TODO: Use `recv` to receive instead // Issue with using recv slice: // Encrypted application data can cramp together into a TCP Segment @@ -947,7 +978,6 @@ 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]); @@ -991,6 +1021,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. + if session.get_tls_state() != TlsState::CONNECTED { + return Ok(0); + } + // Sending order: // 1. Associated data/ TLS Record layer // 2. Encrypted { Payload (data) | Content type: Application Data }