cert: rewrite key handling

Cargo.toml

@@ -26,7 +26,7 @@ default-features = true
 features = [ "heapless" ]
 
 [dependencies.smoltcp]
-version = "0.6.0"
+git = "https://github.com/smoltcp-rs/smoltcp.git"
 default-features = false
 features = ["ethernet", "proto-ipv4", "proto-ipv6", "socket-tcp", "alloc"]
 
@@ -61,11 +61,6 @@ version = "5.1.2"
 default-features = false
 features = []
 
-[dependencies.arraydeque]
-version = "0.4.5"
-default-features = false
-features = [ "use_generic_array" ]
-
 [dependencies.simple_logger]
 version = "1.11.0"
 optional = true

src/certificate.rs

@@ -1,8 +1,11 @@
 use num_enum::IntoPrimitive;
 use num_enum::TryFromPrimitive;
 
+use generic_array::GenericArray;
+
 use crate::parse::parse_asn1_der_rsa_public_key;
 use crate::Error as TlsError;
+use crate::session::CertificatePublicKey;
 
 use sha1::{Sha1, Digest};
 use rsa::{PublicKey, RSAPublicKey, PaddingScheme, BigUint, Hash};
@@ -182,10 +185,112 @@ impl<'a> Certificate<'a> {
             BigUint::from_bytes_be(exponent)
         ).map_err(|_| ())
     }
+
+    // General return public key method
+    // TODO: Replace return_rsa_public_key() with this method
+    // Things to change: session.rs: client_update_for_wait_cert_cr
+    pub(crate) fn get_cert_public_key(&self) -> Result<CertificatePublicKey, ()> {
+        let public_key_info = &self.tbs_certificate.subject_public_key_info;
+        let algorithm_identifier = &public_key_info.algorithm;
+
+        // 3 possibilities: RSA_ENCRYPTION, ID_EC_PUBLIC_KEY, and EdDSA25519
+        match algorithm_identifier.algorithm {
+            oid::RSA_ENCRYPTION => {
+                let (_, (modulus, exponent)) = parse_asn1_der_rsa_public_key(
+                    self.tbs_certificate.subject_public_key_info.subject_public_key
+                ).map_err(|_| ())?;
+        
+                let public_key = RSAPublicKey::new(
+                    BigUint::from_bytes_be(modulus),
+                    BigUint::from_bytes_be(exponent)
+                ).map_err(|_| ())?;
+                Ok(
+                    CertificatePublicKey::RSA {
+                        cert_rsa_public_key: public_key
+                    }
+                )
+            },
+            oid::ID_EC_PUBLIC_KEY => {
+                // Check the type of EC, only support secp256r1
+                // Will definitely NOT support custom curve
+                if algorithm_identifier.parameters != oid::PRIME256V1 {
+                    return Err(());
+                }
+                let p256_verify_key = p256::ecdsa::VerifyKey::from_encoded_point(
+                    &p256::EncodedPoint::from_untagged_bytes(
+                        GenericArray::from_slice(
+                            &public_key_info.subject_public_key[1..]
+                        )
+                    )
+                ).map_err(|_| ())?;
+                Ok(
+                    CertificatePublicKey::ECDSA_SECP256R1_SHA256 {
+                        cert_verify_key: p256_verify_key
+                    }
+                )
+            },
+            oid::ID_EDDSA_25519 => {
+                let ed25519_public_key = ed25519_dalek::PublicKey::from_bytes(
+                    public_key_info.subject_public_key
+                ).map_err(|_| ())?;
+                Ok(
+                    CertificatePublicKey::ED25519 {
+                        cert_eddsa_key: ed25519_public_key
+                    }
+                )
+            },
+            _ => Err(())
+        }
+    }
 }
 
+// TODO: Centralize OID, another OID module can be found in `parse.rs`
 mod oid {
-    // ECDSA signature algorithms
-    pub const SHA1_WITH_RSA_ENCRYPTION: &'static [u8] = &[0x2A, 0x86, 0x48, 0x86, 0xF7, 0x0D, 0x01, 0x01, 0x05];
+    // RSA public key
+    pub const RSA_ENCRYPTION: &'static [u8] =
+        &[0x2A, 0x86, 0x48, 0x86, 0xF7, 0x0D, 0x01, 0x01, 0x01];
+        
+    // EC public key for secp256r1
+    pub const ID_EC_PUBLIC_KEY: &'static [u8] = 
+        &[0x2A, 0x86, 0x48, 0xCE, 0x3D, 0x02, 0x01];
+    pub const PRIME256V1: &'static [u8] =
+        &[0x2A, 0x86, 0x48, 0xCE, 0x3D, 0x03, 0x01, 0x07];
+
+    // EDDSA25519 public key, signature algorithm
+    pub const ID_EDDSA_25519: &'static [u8] =
+        &[0x2B, 0x65, 0x70];
+    
+    // Supported Signature Algorithm (RFC 4055, RFC 3279)
+    // PKCS #1 v1.5
+    pub const SHA1_WITH_RSA_ENCRYPTION: &'static [u8] =
+        &[0x2A, 0x86, 0x48, 0x86, 0xF7, 0x0D, 0x01, 0x01, 0x05];
+    pub const SHA224_WITH_RSA_ENCRYPTION: &'static [u8] =
+        &[0x2A, 0x86, 0x48, 0x86, 0xF7, 0x0D, 0x01, 0x01, 0x0E];
+    pub const SHA256_WITH_RSA_ENCRYPTION: &'static [u8] =
+        &[0x2A, 0x86, 0x48, 0x86, 0xF7, 0x0D, 0x01, 0x01, 0x0B];
+    pub const SHA384_WITH_RSA_ENCRYPTION: &'static [u8] =
+        &[0x2A, 0x86, 0x48, 0x86, 0xF7, 0x0D, 0x01, 0x01, 0x0C];
+    pub const SHA512_WITH_RSA_ENCRYPTION: &'static [u8] =
+        &[0x2A, 0x86, 0x48, 0x86, 0xF7, 0x0D, 0x01, 0x01, 0x0D];
+    
+    // RSASSA_PSS
+    pub const ID_RSASSA_PSS: &'static [u8] =
+        &[0x2A, 0x86, 0x48, 0x86, 0xF7, 0x0D, 0x01, 0x01, 0x0A];
+    
+    // RSAES_OAEP
+    pub const ID_RSAES_OAEP: &'static [u8] =
+        &[0x2A, 0x86, 0x48, 0x86, 0xF7, 0x0D, 0x01, 0x01, 0x07];
+
+    // ECDSA signature algorithms, from OID repo
+    pub const ECDSA_WITH_SHA1: &'static [u8] =
+        &[0x2A, 0x86, 0x48, 0xCE, 0x3D, 0x04, 0x01];
+    pub const ECDSA_WITH_SHA224: &'static [u8] =
+        &[0x2A, 0x86, 0x48, 0xCE, 0x3D, 0x04, 0x03, 0x01];
+    pub const ECDSA_WITH_SHA256: &'static [u8] =
+        &[0x2A, 0x86, 0x48, 0xCE, 0x3D, 0x04, 0x03, 0x02];
+    pub const ECDSA_WITH_SHA384: &'static [u8] =
+        &[0x2A, 0x86, 0x48, 0xCE, 0x3D, 0x04, 0x03, 0x03];
+    pub const ECDSA_WITH_SHA512: &'static [u8] =
+        &[0x2A, 0x86, 0x48, 0xCE, 0x3D, 0x04, 0x03, 0x04];
 }
 

src/parse.rs

@@ -31,7 +31,8 @@ use core::convert::TryInto;
 
 use alloc::vec::Vec;
 
-pub(crate) fn parse_tls_repr(bytes: &[u8]) -> IResult<&[u8], TlsRepr> {
+// Return handshake/payload slice and TLS Record
+pub(crate) fn parse_tls_repr(bytes: &[u8]) -> IResult<&[u8], (&[u8], TlsRepr)> {
     let content_type = take(1_usize);
     let version = take(2_usize);
     let length = take(2_usize);
@@ -51,6 +52,9 @@ pub(crate) fn parse_tls_repr(bytes: &[u8]) -> IResult<&[u8], TlsRepr> {
         handshake: None,
     };
     let (rest, bytes) = take(repr.length)(rest)?;
+
+    // Store a copy of the TLS Handshake slice to return
+    let repr_slice_clone = bytes;
     {
         use crate::tls_packet::TlsContentType::*;
         match repr.content_type {
@@ -68,7 +72,7 @@ pub(crate) fn parse_tls_repr(bytes: &[u8]) -> IResult<&[u8], TlsRepr> {
             _ => todo!()
         }
     }
-    Ok((rest, repr))
+    Ok((rest, (repr_slice_clone, repr)))
 }
 
 // Convert TlsInnerPlainText in RFC 8446 into Handshake
@@ -76,9 +80,9 @@ pub(crate) fn parse_tls_repr(bytes: &[u8]) -> IResult<&[u8], TlsRepr> {
 // 1. Handshake can coalesced into a larger TLS record
 // 2. Content type and zero paddings at the end
 // Return handshake slice for hashing
-pub(crate) fn parse_inner_plaintext_for_handshake(bytes: &[u8]) -> IResult<&[u8], (&[u8], Vec<HandshakeRepr>)> {
+pub(crate) fn parse_inner_plaintext_for_handshake(bytes: &[u8]) -> IResult<&[u8], Vec<(&[u8], HandshakeRepr)>> {
     let mut remaining_bytes = bytes;
-    let mut handshake_vec: Vec<HandshakeRepr> = Vec::new();
+    let mut handshake_vec: Vec<(&[u8], HandshakeRepr)> = Vec::new();
     
     loop {
         // Perform check on the number of remaining bytes
@@ -97,20 +101,19 @@ pub(crate) fn parse_inner_plaintext_for_handshake(bytes: &[u8]) -> IResult<&[u8]
             )(remaining_bytes)?;
             return Ok((
                 &[],
-                (
-                    // A concatenation of all handshakes received
-                    // The remaining content_type byte and zero paddings are stripped
-                    &bytes[
-                        ..(bytes.len()-remaining_bytes.len())
-                    ], 
-                    handshake_vec
-                )
+                // // A concatenation of all handshakes received
+                // // The remaining content_type byte and zero paddings are stripped
+                // &bytes[
+                //     ..(bytes.len()-remaining_bytes.len())
+                // ], 
+                handshake_vec
             ));
         }
 
         let (rem, handshake_repr) = parse_handshake(remaining_bytes)?;
+        let handshake_slice = &remaining_bytes[..(remaining_bytes.len()-rem.len())];
         remaining_bytes = rem;
-        handshake_vec.push(handshake_repr);
+        handshake_vec.push((handshake_slice, handshake_repr));
     }
 }
 

src/session.rs

@@ -57,12 +57,13 @@ pub(crate) struct Session {
     // Sequence number: Start from 0, 64 bits
     // Increment by one per record processed (read OR write)
     // Reset to 0 on rekey AND key exchange
-    // TODO: Force rekey if sequence number need to wrap
+    // TODO: Force rekey if sequence number need to wrap (very low priority)
     client_sequence_number: u64,
     server_sequence_number: u64,
     // Certificate public key
     // For Handling CertificateVerify
-    cert_rsa_public_key: Option<RSAPublicKey>,
+    cert_rsa_public_key: Option<RSAPublicKey>,      // TODO: Replace and remove
+    cert_public_key: Option<CertificatePublicKey>,
 }
 
 impl Session {
@@ -93,7 +94,8 @@ impl Session {
             server_application_nonce: None,
             client_sequence_number: 0,
             server_sequence_number: 0,
-            cert_rsa_public_key: None
+            cert_rsa_public_key: None,      // TODO: Remove over-specific public key
+            cert_public_key: None
         }
     }
 
@@ -500,8 +502,7 @@ impl Session {
 
         // Handle Ed25519 and p256 separately
         // These 2 algorithms have a mandated hash function
-        if signature_algorithm == SignatureScheme::ecdsa_secp256r1_sha256
-        {
+        if signature_algorithm == SignatureScheme::ecdsa_secp256r1_sha256 {
             todo!()
         }
 
@@ -1121,6 +1122,7 @@ impl Session {
 
     // TODO: Merge decryption methods
     // Take control of the entire decryption, manually invoke detached decryption
+    // TODO: Bad naming, it should say the KEY of application data
     pub(crate) fn decrypt_application_data_in_place(
         &self,
         associated_data: &[u8],
@@ -1145,6 +1147,7 @@ impl Session {
         let mut processed_nonce: [u8; 12] = [0; 12];
         NetworkEndian::write_uint128(&mut processed_nonce, nonce ^ clipped_seq_num, 12);
 
+        // Duplicate authentication tag
         let buffer_size = buffer.len();
         let tag = GenericArray::clone_from_slice(&buffer[(buffer_size-16)..]);
 
@@ -1156,6 +1159,7 @@ impl Session {
         )
     }
 
+    // Decryption using handshake keys
     pub(crate) fn decrypt_in_place(
         &self,
         associated_data: &[u8],
@@ -1187,6 +1191,44 @@ impl Session {
         )
     }
 
+    // A veriant for handshake decryption in-place and detached
+    // Caller need to manually discard the authentication bytes
+    pub(crate) fn decrypt_in_place_detached(
+        &self,
+        associated_data: &[u8],
+        buffer: &mut [u8]
+    ) -> Result<(), Error> {
+        let (seq_num, nonce, cipher): (u64, &Vec<u8, U12>, &Cipher) = match self.role {
+            TlsRole::Server => {(
+                self.client_sequence_number,
+                self.client_handshake_nonce.as_ref().unwrap(),
+                self.client_handshake_cipher.as_ref().unwrap()
+            )},
+            TlsRole::Client => {(
+                self.server_sequence_number,
+                self.server_handshake_nonce.as_ref().unwrap(),
+                self.server_handshake_cipher.as_ref().unwrap()
+            )},
+        };
+
+        // Calculate XOR'ed nonce
+        let nonce: u128 = NetworkEndian::read_uint128(nonce, 12);
+        let clipped_seq_num: u128 = seq_num.into();
+        let mut processed_nonce: [u8; 12] = [0; 12];
+        NetworkEndian::write_uint128(&mut processed_nonce, nonce ^ clipped_seq_num, 12);
+
+        // Duplicate authentication tag
+        let buffer_size = buffer.len();
+        let tag = GenericArray::clone_from_slice(&buffer[(buffer_size-16)..]);
+
+        cipher.decrypt_in_place_detached(
+            &GenericArray::from_slice(&processed_nonce),
+            associated_data,
+            &mut buffer[..(buffer_size-16)],
+            &tag
+        )
+    }
+
     pub(crate) fn increment_client_sequence_number(&mut self) {
         self.client_sequence_number += 1;
     }
@@ -1401,3 +1443,15 @@ impl Cipher {
         }
     }
 }
+
+pub(crate) enum CertificatePublicKey {
+    RSA {
+        cert_rsa_public_key: RSAPublicKey
+    },
+    ECDSA_SECP256R1_SHA256 {
+        cert_verify_key: p256::ecdsa::VerifyKey
+    },
+    ED25519 {
+        cert_eddsa_key: ed25519_dalek::PublicKey
+    }
+}

src/tls.rs

@@ -23,6 +23,7 @@ use aes_gcm::AeadInPlace;
 
 use nom::bytes::complete::take;
 use nom::error::ErrorKind;
+use nom::combinator::complete;
 
 use alloc::vec::Vec;
 use heapless::Vec as HeaplessVec;
@@ -202,6 +203,7 @@ impl<R: 'static + RngCore + CryptoRng> TlsSocket<R> {
 
         // 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)?;
 
@@ -209,15 +211,84 @@ impl<R: 'static + RngCore + CryptoRng> TlsSocket<R> {
         // Process as a queue
         let tls_repr_vec_size = tls_repr_vec.len();
         for _index in 0..tls_repr_vec_size {
-            let repr = tls_repr_vec.remove(0);
-            self.process(repr)?;
+            let (repr_slice, mut repr) = tls_repr_vec.remove(0);
+            // TODO:
+            // Check TLS content type
+            // If application data, decrypt before process.
+            // If there are multiple handshakes within a record,
+            //     give each of them a unique TLS record wrapper and process
+            //     If a pure application data is found, sliently ignore
+            // Otherwise process directly
+            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();
+                    // log::info!("Decypted data: {:?}", app_data);
+                    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);
+                    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)
+                        }
+                    )?;
+                }
+            }
+
+            
+            else {
+                self.process(repr_slice, repr)?;
+                log::info!("Processed record");
+            }
         }
 
         Ok(self.session.borrow().has_completed_handshake())
     }
 
     // Process TLS ingress during handshake
-    fn process(&self, mut repr: TlsRepr) -> Result<()> {
+    // The slice should ONLY include handshake overhead
+    // i.e. Exclude 5 bytes of TLS Record
+    //      Include 4 bytes of HandshakeRepr, everything within the same handshake
+    fn process(&self, handshake_slice: &[u8], mut repr: TlsRepr) -> Result<()> {
         // Change_cipher_spec check:
         // Must receive CCS before recv peer's FINISH message
         // i.e. Must happen after START and before CONNECTED
@@ -226,8 +297,10 @@ impl<R: 'static + RngCore + CryptoRng> TlsSocket<R> {
         // Drop the message and update `received_change_cipher_spec`
         // Note: CSS doesn't count as a proper record, no need to increment sequence number
         if repr.is_change_cipher_spec() {
-            let mut session = self.session.borrow_mut();
+            log::info!("Change Cipher Spec");
+            let mut session = self.session.try_borrow_mut().expect("Cannot borrow mut");
             session.receive_change_cipher_spec();
+            log::info!("Changed Cipher Spec");
             return Ok(())
         }
 
@@ -237,12 +310,6 @@ impl<R: 'static + RngCore + CryptoRng> TlsSocket<R> {
         match tls_state {
             // During WAIT_SH for a TLS client, client should wait for ServerHello
             TlsState::WAIT_SH => {
-                // Legacy_protocol must be TLS 1.2
-                if repr.version != TlsVersion::Tls12 {
-                    // Abort communication
-                    todo!()
-                }
-
                 // TODO: Validate SH
                 if repr.is_server_hello() {
                     // Check SH content:
@@ -342,18 +409,12 @@ impl<R: 'static + RngCore + CryptoRng> TlsSocket<R> {
                         todo!()
                     }
 
-                    // This is indeed a desirable ServerHello TLS repr
-                    // Reprocess ServerHello into a slice
-                    // Update session with required parameter
-                    let mut array = [0; 512];
-                    let mut buffer = TlsBuffer::new(&mut array);
-                    buffer.enqueue_tls_repr(repr)?;
-                    let slice: &[u8] = buffer.into();
+                    // Get slice without reserialization
                     let mut session = self.session.borrow_mut();
                     session.client_update_for_sh(
                         selected_cipher.unwrap(),
                         server_public.unwrap(),
-                        &slice[5..]
+                        handshake_slice
                     );
                     // Key exchange occurred, seq_num is set to 0
                     // Do NOT update seq_num again. Early return.
@@ -363,34 +424,9 @@ impl<R: 'static + RngCore + CryptoRng> TlsSocket<R> {
 
             // Expect encrypted extensions after receiving SH
             TlsState::WAIT_EE => {
-                // Check that the packet is classified as application data
-                if !repr.is_application_data() {
-                    // Abort communication, this affect IV calculation
-                    todo!()
-                }
-
-                // ExcepytedExtensions are disguised as ApplicationData
-                // Pull out the `payload` from TlsRepr, decrypt as EE
-                let mut payload = repr.payload.take().unwrap();
-                let mut array: [u8; 5] = [0; 5];
-                let mut buffer = TlsBuffer::new(&mut array);
-                buffer.write_u8(repr.content_type.into())?;
-                buffer.write_u16(repr.version.into())?;
-                buffer.write_u16(repr.length)?;
-                let associated_data: &[u8] = buffer.into();
-                {
-                    self.session.borrow_mut().decrypt_in_place(
-                        associated_data,
-                        &mut payload
-                    );
-                }
-
-                let parse_result = parse_inner_plaintext_for_handshake(&payload);
-                let (_, (handshake_slice, mut handshake_vec)) = 
-                    parse_result.map_err(|_| Error::Unrecognized)?;
-
                 // Verify that it is indeed an EE
-                let might_be_ee = handshake_vec.remove(0);
+                // let might_be_ee = handshake_vec.remove(0);
+                let might_be_ee = repr.handshake.take().unwrap();
                 if might_be_ee.get_msg_type() != HandshakeType::EncryptedExtensions {
                     // Process the other handshakes in "handshake_vec"
                     todo!()
@@ -413,47 +449,15 @@ impl<R: 'static + RngCore + CryptoRng> TlsSocket<R> {
                         &ee_slice
                     );
                 
-                // TODO: Handle in WAIT_CERT_CR if there are still unprocessed handshakes
-                // Ideas: 1. Split off WAIT_CERT_CR handling into a separate function
-                //              so WAIT_EE branch can jsut call WAIT_CERT_CR branch
-                //              if there are extra handshake unprocessed
-                // 2. Merge state dependent listeners into 1 branch, execute conditionally
-            },
+                log::info!("Received EE");
+           },
 
             // In this stage, wait for a certificate from server
             // Parse the certificate and check its content
             TlsState::WAIT_CERT_CR => {
-                // Check that the packet is classified as application data
-                // Certificates transfer is disguised as application data
-                if !repr.is_application_data() {
-                    // Abort communication, this affect IV calculation
-                    todo!()
-                }
-
-                // Pull out the `payload` from TlsRepr, decrypt as CERT
-                let mut payload = repr.payload.take().unwrap();
-
-                // Instantiate associated data and decrypt
-                let mut array: [u8; 5] = [0; 5];
-                let mut buffer = TlsBuffer::new(&mut array);
-                buffer.write_u8(repr.content_type.into())?;
-                buffer.write_u16(repr.version.into())?;
-                buffer.write_u16(repr.length)?;
-                let associated_data: &[u8] = buffer.into();
-                {
-                    self.session.borrow_mut().decrypt_in_place(
-                        associated_data,
-                        &mut payload
-                    );
-                }
-
-                // Parse the certificate from TLS payload
-                let parse_result = parse_inner_plaintext_for_handshake(&payload);
-                let (_, (handshake_slice, mut handshake_vec)) = parse_result
-                    .map_err(|_| Error::Unrecognized)?;
-
                 // Verify that it is indeed an Certificate
-                let might_be_cert = handshake_vec.remove(0);
+                // let might_be_cert = handshake_vec.remove(0);
+                let might_be_cert = repr.handshake.take().unwrap();
                 if might_be_cert.get_msg_type() != HandshakeType::Certificate {
                     // Process the other handshakes in "handshake_vec"
                     todo!()
@@ -482,41 +486,15 @@ impl<R: 'static + RngCore + CryptoRng> TlsSocket<R> {
                         &cert_slice,
                         cert.return_rsa_public_key().unwrap()
                     );
+                log::info!("Received WAIT_CERT_CR");
             },
 
             // In this stage, server will eventually send a CertificateVerify
             // Verify that the signature is indeed correct
             TlsState::WAIT_CV => {
-                // CertificateVerify is disguised as Application Data
-                if !repr.is_application_data() {
-                    // Abort communication, this affect IV calculation
-                    todo!()
-                }
-
-                // Pull out the `payload` from TlsRepr, decrypt as CV
-                let mut payload = repr.payload.take().unwrap();
-
-                // Instantiate associated data and decrypt
-                let mut array: [u8; 5] = [0; 5];
-                let mut buffer = TlsBuffer::new(&mut array);
-                buffer.write_u8(repr.content_type.into())?;
-                buffer.write_u16(repr.version.into())?;
-                buffer.write_u16(repr.length)?;
-                let associated_data: &[u8] = buffer.into();
-                {
-                    self.session.borrow_mut().decrypt_in_place(
-                        associated_data,
-                        &mut payload
-                    );
-                }
-
-                // Parse the certificate from TLS payload
-                let parse_result = parse_inner_plaintext_for_handshake(&payload);
-                let (_, (handshake_slice, mut handshake_vec)) = parse_result
-                    .map_err(|_| Error::Unrecognized)?;
-
                 // Ensure that it is CertificateVerify
-                let might_be_cert_verify = handshake_vec.remove(0);
+                // let might_be_cert_verify = handshake_vec.remove(0);
+                let might_be_cert_verify = repr.handshake.take().unwrap();
                 if might_be_cert_verify.get_msg_type() != HandshakeType::CertificateVerify {
                     // Process the other handshakes in "handshake_vec"
                     todo!()
@@ -538,40 +516,14 @@ impl<R: 'static + RngCore + CryptoRng> TlsSocket<R> {
                         sig_alg,
                         signature
                     );
+                log::info!("Received CV");
             },
 
             // Client will receive a Finished handshake from server
             TlsState::WAIT_FINISHED => {
-                // Finished is disguised as Application Data
-                if !repr.is_application_data() {
-                    // Abort communication, this affect IV calculation
-                    todo!()
-                }
-
-                // Pull out the `payload` from TlsRepr, decrypt as Finished
-                let mut payload = repr.payload.take().unwrap();
-
-                // Instantiate associated data and decrypt
-                let mut array: [u8; 5] = [0; 5];
-                let mut buffer = TlsBuffer::new(&mut array);
-                buffer.write_u8(repr.content_type.into())?;
-                buffer.write_u16(repr.version.into())?;
-                buffer.write_u16(repr.length)?;
-                let associated_data: &[u8] = buffer.into();
-                {
-                    self.session.borrow_mut().decrypt_in_place(
-                        associated_data,
-                        &mut payload
-                    );
-                }
-
-                // Parse the TLS inner ciphertext as a Finished handshake
-                let parse_result = parse_inner_plaintext_for_handshake(&payload);
-                let (_, (handshake_slice, mut handshake_vec)) = parse_result
-                    .map_err(|_| Error::Unrecognized)?;
-
                 // Ensure that it is Finished
-                let might_be_server_finished = handshake_vec.remove(0);
+                // let might_be_server_finished = handshake_vec.remove(0);
+                let might_be_server_finished = repr.handshake.take().unwrap();
                 if might_be_server_finished.get_msg_type() != HandshakeType::Finished {
                     // Process the other handshakes in "handshake_vec"
                     todo!()
@@ -592,15 +544,11 @@ impl<R: 'static + RngCore + CryptoRng> TlsSocket<R> {
                         server_finished_slice,
                         might_be_server_finished.get_verify_data().unwrap()
                     );
+                log::info!("Received server FIN");
             }
 
             _ => {},
         }
-
-        // A TLS Record was received and processed and verified
-        // Increment sequence number
-        self.session.borrow_mut().increment_server_sequence_number();
-
         Ok(())
     }
 
@@ -699,18 +647,20 @@ impl<R: 'static + RngCore + CryptoRng> TlsSocket<R> {
     // 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::<TlsRepr>> {
+    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<TlsRepr> = Vec::new();
+        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)) => {
-                    vec.push(repr);
+                Ok((rest, (repr_slice, repr))) => {
+                    vec.push(
+                        (repr_slice, repr)
+                    );
                     if rest.len() == 0 {
                         return Ok(vec);
                     } else {