TCB: move to layer 4

src/client/client_info.rs

@@ -21,7 +21,7 @@ use std::net::IpAddr;
 use pnet::packet::ip::IpNextHeaderProtocol;
 use pnet::util::MacAddr;
 
-#[derive(PartialEq, Hash, Copy, Clone)]
+#[derive(PartialEq, Hash, Copy, Clone, Debug)]
 pub struct ClientInfoSrcDst<A: Hash + PartialEq + Clone> {
     pub src: Option<A>,
     pub dst: Option<A>,
@@ -35,7 +35,7 @@ pub struct ClientInfoSrcDst<A: Hash + PartialEq + Clone> {
  * - source and dest. transport port
  * - syn cookie
  **/
-#[derive(Copy, Clone, PartialEq)]
+#[derive(Copy, Clone, PartialEq, Debug)]
 pub struct ClientInfo {
     pub mac: ClientInfoSrcDst<MacAddr>,
     pub ip: ClientInfoSrcDst<IpAddr>,

src/layer_4/tcp.rs

@@ -49,16 +49,24 @@ pub fn repl<'a, 'b>(
             };
             /* Compute syncookie */
             if let Ok(cookie) = synackcookie::generate(&client_info, &masscanned.synack_key) {
+                client_info.cookie = Some(cookie);
+                if !proto::is_tcb_set(cookie) {
+                    /* First Ack: check syncookie, create tcb */
                     if cookie != ackno {
                         masscanned.log.tcp_drop(tcp_req, client_info);
                         return None;
                     }
-                client_info.cookie = Some(cookie);
+                    proto::add_tcb(cookie);
+                }
             }
             warn!("ACK to PSH-ACK on port {}", tcp_req.get_destination());
             let payload = tcp_req.payload();
             /* Any answer to upper-layer protocol? */
-            if let Some(repl) = proto::repl(&payload, masscanned, &mut client_info) {
+            let mut payload_repl = None;
+            proto::get_tcb(client_info.cookie.unwrap(), |tcb| {
+                payload_repl = proto::repl(&payload, masscanned, &mut client_info, tcb);
+            });
+            if let Some(repl) = payload_repl {
                 tcp_repl = MutableTcpPacket::owned(
                     [vec![0; MutableTcpPacket::minimum_packet_size()], repl].concat(),
                 )

src/layer_4/udp.rs

@@ -34,7 +34,7 @@ pub fn repl<'a, 'b>(
     client_info.port.dst = Some(udp_req.get_destination());
     let payload = udp_req.payload();
     let mut udp_repl;
-    if let Some(repl) = proto::repl(&payload, masscanned, &mut client_info) {
+    if let Some(repl) = proto::repl(&payload, masscanned, &mut client_info, None) {
         udp_repl = MutableUdpPacket::owned(
             [vec![0; MutableUdpPacket::minimum_packet_size()], repl].concat(),
         )

src/proto/mod.rs

@@ -17,15 +17,13 @@
 use lazy_static::lazy_static;
 use log::*;
 use pnet::packet::ip::IpNextHeaderProtocols;
-use std::collections::HashMap;
-use std::sync::Mutex;
 
 use crate::client::ClientInfo;
 use crate::smack::{Smack, SmackFlags, BASE_STATE, NO_MATCH, SMACK_CASE_SENSITIVE};
 use crate::Masscanned;
 
 mod http;
-use http::{ProtocolState as HTTPProtocolState, HTTP_VERBS};
+use http::HTTP_VERBS;
 
 mod stun;
 use stun::{STUN_PATTERN_CHANGE_REQUEST, STUN_PATTERN_EMPTY, STUN_PATTERN_MAGIC};
@@ -37,11 +35,14 @@ mod ghost;
 use ghost::GHOST_PATTERN_SIGNATURE;
 
 mod rpc;
-use rpc::{ProtocolState as RPCProtocolState, RPC_CALL_TCP, RPC_CALL_UDP};
+use rpc::{RPC_CALL_TCP, RPC_CALL_UDP};
 
 mod smb;
 use smb::{SMB1_PATTERN_MAGIC, SMB2_PATTERN_MAGIC};
 
+mod tcb;
+pub use tcb::{add_tcb, get_tcb, is_tcb_set, ProtocolState, TCPControlBlock};
+
 const PROTO_NONE: usize = 0;
 const PROTO_HTTP: usize = 1;
 const PROTO_STUN: usize = 2;
@@ -52,23 +53,8 @@ const PROTO_RPC_UDP: usize = 6;
 const PROTO_SMB1: usize = 7;
 const PROTO_SMB2: usize = 8;
 
-enum ProtocolState {
-    HTTP(HTTPProtocolState),
-    RPC(RPCProtocolState),
-}
-
-pub struct TCPControlBlock {
-    /* state used to detect protocols (not specific) */
-    smack_state: usize,
-    /* detected protocol */
-    proto_id: usize,
-    /* internal state of protocol parser (e.g., HTTP parsing) */
-    proto_state: Option<ProtocolState>,
-}
-
 lazy_static! {
     static ref PROTO_SMACK: Smack = proto_init();
-    static ref CONTABLE: Mutex<HashMap<u32, TCPControlBlock>> = Mutex::new(HashMap::new());
 }
 
 fn proto_init() -> Smack {
@@ -134,36 +120,22 @@ pub fn repl<'a>(
     data: &'a [u8],
     masscanned: &Masscanned,
     mut client_info: &mut ClientInfo,
+    mut tcb: Option<&mut TCPControlBlock>,
 ) -> Option<Vec<u8>> {
     debug!("packet payload: {:?}", data);
     let mut id;
-    let mut ct = CONTABLE.lock().unwrap();
-    let mut tcb = None;
     if client_info.transport == Some(IpNextHeaderProtocols::Tcp) && client_info.cookie == None {
         error!("Unexpected empty cookie");
         return None;
-    } else if client_info.cookie != None {
+    } else if let Some(t) = &mut tcb {
         /* proto over TCP */
-        let cookie = client_info.cookie.unwrap();
-        if !ct.contains_key(&cookie) {
-            ct.insert(
-                cookie,
-                TCPControlBlock {
-                    smack_state: BASE_STATE,
-                    proto_id: PROTO_NONE,
-                    proto_state: None,
-                },
-            );
-        }
         let mut i = 0;
-        let mut t = ct.get_mut(&cookie).unwrap();
         if t.proto_id == PROTO_NONE {
             let mut state = t.smack_state;
             t.proto_id = PROTO_SMACK.search_next(&mut state, data, &mut i);
             t.smack_state = state;
         }
         id = t.proto_id;
-        tcb = Some(t);
     } else {
         /* proto over else (e.g., UDP) */
         let mut i = 0;
@@ -204,176 +176,6 @@ mod tests {
     use pnet::util::MacAddr;
 
     use crate::logger::MetaLogger;
-    use crate::synackcookie;
-
-    #[test]
-    fn test_proto_tcb_proto_id() {
-        let mut client_info = ClientInfo::new();
-        let test_ip_addr = Ipv4Addr::new(3, 2, 1, 0);
-        client_info.ip.src = Some(IpAddr::V4(test_ip_addr));
-        client_info.port.src = Some(65000);
-        client_info.port.dst = Some(80);
-        client_info.transport = Some(IpNextHeaderProtocols::Tcp);
-        let masscanned_ip_addr = Ipv4Addr::new(0, 1, 2, 3);
-        client_info.ip.dst = Some(IpAddr::V4(masscanned_ip_addr));
-        let mut ips = HashSet::new();
-        ips.insert(IpAddr::V4(masscanned_ip_addr));
-        /* Construct masscanned context object */
-        let masscanned = Masscanned {
-            synack_key: [0, 0],
-            mac: MacAddr::from_str("00:11:22:33:44:55").expect("error parsing MAC address"),
-            iface: None,
-            ip_addresses: Some(&ips),
-            log: MetaLogger::new(),
-        };
-        let cookie = synackcookie::generate(&client_info, &masscanned.synack_key).unwrap();
-        client_info.cookie = Some(cookie);
-        {
-            let ct = CONTABLE.lock().unwrap();
-            if ct.contains_key(&cookie) {
-                panic!("expected no TCB entry, found one");
-            }
-        }
-        /***** TEST PROTOCOL ID IN TCB *****/
-        let payload = b"GET / HTTP/1.1\r\n";
-        repl(&payload.to_vec(), &masscanned, &mut client_info);
-        {
-            let mut ct = CONTABLE.lock().unwrap();
-            if !ct.contains_key(&cookie) {
-                panic!("expected a TCB entry, not found");
-            }
-            let t = ct.get_mut(&cookie).unwrap();
-            assert!(t.proto_id == PROTO_HTTP);
-        }
-        /***** SENDING MORE DATA *****/
-        let payload = b"garbage data with no specific format (no protocol)\r\n\r\n";
-        repl(&payload.to_vec(), &masscanned, &mut client_info);
-        {
-            let mut ct = CONTABLE.lock().unwrap();
-            if !ct.contains_key(&cookie) {
-                panic!("expected a TCB entry, not found");
-            }
-            let t = ct.get_mut(&cookie).unwrap();
-            assert!(t.proto_id == PROTO_HTTP);
-        }
-    }
-
-    #[test]
-    fn test_proto_tcb_proto_state_http() {
-        let mut client_info = ClientInfo::new();
-        let test_ip_addr = Ipv4Addr::new(3, 2, 1, 0);
-        client_info.ip.src = Some(IpAddr::V4(test_ip_addr));
-        client_info.port.src = Some(65001);
-        client_info.port.dst = Some(80);
-        client_info.transport = Some(IpNextHeaderProtocols::Tcp);
-        let masscanned_ip_addr = Ipv4Addr::new(0, 1, 2, 3);
-        client_info.ip.dst = Some(IpAddr::V4(masscanned_ip_addr));
-        let mut ips = HashSet::new();
-        ips.insert(IpAddr::V4(masscanned_ip_addr));
-        /* Construct masscanned context object */
-        let masscanned = Masscanned {
-            synack_key: [0, 0],
-            mac: MacAddr::from_str("00:11:22:33:44:55").expect("error parsing MAC address"),
-            iface: None,
-            ip_addresses: Some(&ips),
-            log: MetaLogger::new(),
-        };
-        let cookie = synackcookie::generate(&client_info, &masscanned.synack_key).unwrap();
-        client_info.cookie = Some(cookie);
-        {
-            let ct = CONTABLE.lock().unwrap();
-            if ct.contains_key(&cookie) {
-                panic!("expected no TCB entry, found one");
-            }
-        }
-        /***** TEST PROTOCOL ID IN TCB *****/
-        let payload = b"GET / HTTP/1.1\r\n";
-        repl(&payload.to_vec(), &masscanned, &mut client_info);
-        {
-            let mut ct = CONTABLE.lock().unwrap();
-            if !ct.contains_key(&cookie) {
-                panic!("expected a TCB entry, not found");
-            }
-            let t = ct.get_mut(&cookie).unwrap();
-            assert!(t.proto_id == PROTO_HTTP);
-            if let Some(ProtocolState::HTTP(_)) = t.proto_state {
-            } else {
-                panic!("expected a HTTP protocole state, found None");
-            }
-        }
-        /***** SENDING MORE DATA *****/
-        let payload = b"Field: empty\r\n\r\n";
-        /* Should have an answer here */
-        if let None = repl(&payload.to_vec(), &masscanned, &mut client_info) {
-            panic!("expected an HTTP response, got nothing");
-        }
-        {
-            let mut ct = CONTABLE.lock().unwrap();
-            if !ct.contains_key(&cookie) {
-                panic!("expected a TCB entry, not found");
-            }
-            let t = ct.get_mut(&cookie).unwrap();
-            assert!(t.proto_id == PROTO_HTTP);
-        }
-    }
-
-    #[test]
-    fn test_proto_tcb_proto_state_rpc() {
-        let mut client_info = ClientInfo::new();
-        let test_ip_addr = Ipv4Addr::new(3, 2, 1, 0);
-        client_info.ip.src = Some(IpAddr::V4(test_ip_addr));
-        client_info.port.src = Some(65002);
-        client_info.port.dst = Some(80);
-        client_info.transport = Some(IpNextHeaderProtocols::Tcp);
-        let masscanned_ip_addr = Ipv4Addr::new(0, 1, 2, 3);
-        client_info.ip.dst = Some(IpAddr::V4(masscanned_ip_addr));
-        let mut ips = HashSet::new();
-        ips.insert(IpAddr::V4(masscanned_ip_addr));
-        /* Construct masscanned context object */
-        let masscanned = Masscanned {
-            synack_key: [0, 0],
-            mac: MacAddr::from_str("00:11:22:33:44:55").expect("error parsing MAC address"),
-            iface: None,
-            ip_addresses: Some(&ips),
-            log: MetaLogger::new(),
-        };
-        let cookie = synackcookie::generate(&client_info, &masscanned.synack_key).unwrap();
-        client_info.cookie = Some(cookie);
-        {
-            let ct = CONTABLE.lock().unwrap();
-            if ct.contains_key(&cookie) {
-                panic!("expected no TCB entry, found one");
-            }
-        }
-        /***** TEST PROTOCOL ID IN TCB *****/
-        let payload = b"\x80\x00\x00\x28\x72\xfe\x1d\x13\x00\x00\x00\x00\x00\x00\x00\x02\x00\x01\x86\xa0\x00\x01\x97\x7c\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00";
-        repl(&payload[0..28].to_vec(), &masscanned, &mut client_info);
-        {
-            let mut ct = CONTABLE.lock().unwrap();
-            if !ct.contains_key(&cookie) {
-                panic!("expected a TCB entry, not found");
-            }
-            let t = ct.get_mut(&cookie).unwrap();
-            assert!(t.proto_id == PROTO_RPC_TCP);
-            if let Some(ProtocolState::RPC(_)) = t.proto_state {
-            } else {
-                panic!("expected a RPC protocole state, found None");
-            }
-        }
-        /***** SENDING MORE DATA *****/
-        /* Should have an answer here */
-        if let None = repl(&payload[28..].to_vec(), &masscanned, &mut client_info) {
-            panic!("expected a RPC response, got nothing");
-        }
-        {
-            let mut ct = CONTABLE.lock().unwrap();
-            if !ct.contains_key(&cookie) {
-                panic!("expected a TCB entry, not found");
-            }
-            let t = ct.get_mut(&cookie).unwrap();
-            assert!(t.proto_id == PROTO_RPC_TCP);
-        }
-    }
 
     #[test]
     fn test_proto_dispatch_stun() {
@@ -401,7 +203,7 @@ mod tests {
          */
         let payload =
             b"\x00\x01\x00\x00\x21\x12\xa4\x42\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00";
-        let _stun_resp = if let Some(r) = repl(payload, &masscanned, &mut client_info) {
+        let _stun_resp = if let Some(r) = repl(payload, &masscanned, &mut client_info, None) {
             r
         } else {
             panic!("expected an answer, got nothing");
@@ -415,7 +217,7 @@ mod tests {
          */
         let payload =
             b"\x00\x01\x00\x00\xaa\xbb\xcc\xdd\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00";
-        let _stun_resp = if let Some(r) = repl(payload, &masscanned, &mut client_info) {
+        let _stun_resp = if let Some(r) = repl(payload, &masscanned, &mut client_info, None) {
             r
         } else {
             panic!("expected an answer, got nothing");
@@ -428,7 +230,7 @@ mod tests {
          */
         let payload =
             b"\x00\x01\x00\x08\x01\xdb\xd4]4\x9f\xe2RQ\x19\x05,\x93\x14f4\x00\x03\x00\x04\x00\x00\x00\x00";
-        let _stun_resp = if let Some(r) = repl(payload, &masscanned, &mut client_info) {
+        let _stun_resp = if let Some(r) = repl(payload, &masscanned, &mut client_info, None) {
             r
         } else {
             panic!("expected an answer, got nothing");
@@ -466,8 +268,8 @@ mod tests {
             "SSH-2.0-OpenSSH_6.7p1 Raspbian-5+deb8u3",
         ];
         for payload in payloads.iter() {
-            let _ssh_resp = if let Some(r) = repl(payload.as_bytes(), &masscanned, &mut client_info)
+            let _ssh_resp =
-            {
+                if let Some(r) = repl(payload.as_bytes(), &masscanned, &mut client_info, None) {
                     r
                 } else {
                     panic!("expected an answer, got nothing");
@@ -498,7 +300,7 @@ mod tests {
         ];
         for payload in payloads.iter() {
             let _ghost_resp =
-                if let Some(r) = repl(&payload.to_vec(), &masscanned, &mut client_info) {
+                if let Some(r) = repl(&payload.to_vec(), &masscanned, &mut client_info, None) {
                     r
                 } else {
                     panic!("expected an answer, got nothing");
@@ -527,12 +329,12 @@ mod tests {
         };
         /***** TEST COMPLETE REQUEST *****/
         let payload = b"GET / HTTP/1.1\r\n\r\n";
-        if let None = repl(&payload.to_vec(), &masscanned, &mut client_info) {
+        if let None = repl(&payload.to_vec(), &masscanned, &mut client_info, None) {
             panic!("expected an answer, got nothing");
         }
         /***** TEST INCOMPLETE REQUEST *****/
         let payload = b"GET / HTTP/1.1\r\n";
-        if let Some(_) = repl(&payload.to_vec(), &masscanned, &mut client_info) {
+        if let Some(_) = repl(&payload.to_vec(), &masscanned, &mut client_info, None) {
             panic!("expected no answer, got one");
         }
     }

src/proto/tcb.rs

@@ -0,0 +1,267 @@
+// This file is part of masscanned.
+// Copyright 2021 - The IVRE project
+//
+// Masscanned is free software: you can redistribute it and/or modify it
+// under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// Masscanned is distributed in the hope that it will be useful, but WITHOUT
+// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+// or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
+// License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with Masscanned. If not, see <http://www.gnu.org/licenses/>.
+
+use lazy_static::lazy_static;
+
+use std::collections::HashMap;
+use std::sync::Mutex;
+
+use super::http::ProtocolState as HTTPProtocolState;
+use super::rpc::ProtocolState as RPCProtocolState;
+use crate::proto::{BASE_STATE, PROTO_NONE};
+
+pub enum ProtocolState {
+    HTTP(HTTPProtocolState),
+    RPC(RPCProtocolState),
+}
+
+pub struct TCPControlBlock {
+    /* state used to detect protocols (not specific) */
+    pub smack_state: usize,
+    /* detected protocol */
+    pub proto_id: usize,
+    /* internal state of protocol parser (e.g., HTTP parsing) */
+    pub proto_state: Option<ProtocolState>,
+}
+
+lazy_static! {
+    static ref CONTABLE: Mutex<HashMap<u32, TCPControlBlock>> = Mutex::new(HashMap::new());
+}
+
+pub fn is_tcb_set(cookie: u32) -> bool {
+    CONTABLE.lock().unwrap().contains_key(&cookie)
+}
+
+pub fn get_tcb<F>(cookie: u32, mut f: F)
+where
+    F: FnMut(Option<&mut TCPControlBlock>),
+{
+    f(CONTABLE.lock().unwrap().get_mut(&cookie));
+}
+
+pub fn add_tcb(cookie: u32) {
+    let mut ct = CONTABLE.lock().unwrap();
+    let tcb = TCPControlBlock {
+        smack_state: BASE_STATE,
+        proto_id: PROTO_NONE,
+        proto_state: None,
+    };
+    if !ct.contains_key(&cookie) {
+        ct.insert(cookie, tcb);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::collections::HashSet;
+    use std::net::{IpAddr, Ipv4Addr};
+    use std::str::FromStr;
+
+    use pnet::{
+        packet::{ip::IpNextHeaderProtocols, tcp::TcpPacket},
+        util::MacAddr,
+    };
+
+    use crate::client::ClientInfo;
+    use crate::layer_4::tcp;
+    use crate::logger::MetaLogger;
+    use crate::proto::{PROTO_HTTP, PROTO_RPC_TCP};
+    use crate::synackcookie;
+    use crate::Masscanned;
+
+    fn get_dummy_tcp(&client_info: &ClientInfo) -> Vec<u8> {
+        /* Craft a TCP ACK+PUSH packet with correct ports and ack */
+        let mut pkt = Vec::new();
+        pkt.extend_from_slice(&client_info.port.src.unwrap().to_be_bytes());
+        pkt.extend_from_slice(&client_info.port.dst.unwrap().to_be_bytes());
+        pkt.extend_from_slice(b"\x00\x00\x00\x00");
+        pkt.extend_from_slice(&(client_info.cookie.unwrap() + 1).to_be_bytes());
+        pkt.extend_from_slice(b"P\x18 \x00\x00\x00\x00\x00");
+        pkt
+    }
+
+    #[test]
+    fn test_proto_tcb_proto_id() {
+        let mut client_info = ClientInfo::new();
+        let test_ip_addr = Ipv4Addr::new(3, 2, 1, 0);
+        client_info.ip.src = Some(IpAddr::V4(test_ip_addr));
+        client_info.port.src = Some(65000);
+        client_info.port.dst = Some(80);
+        client_info.transport = Some(IpNextHeaderProtocols::Tcp);
+        let masscanned_ip_addr = Ipv4Addr::new(0, 1, 2, 3);
+        client_info.ip.dst = Some(IpAddr::V4(masscanned_ip_addr));
+        let mut ips = HashSet::new();
+        ips.insert(IpAddr::V4(masscanned_ip_addr));
+        /* Construct masscanned context object */
+        let masscanned = Masscanned {
+            synack_key: [0, 0],
+            mac: MacAddr::from_str("00:11:22:33:44:55").expect("error parsing MAC address"),
+            iface: None,
+            ip_addresses: Some(&ips),
+            log: MetaLogger::new(),
+        };
+        let cookie = synackcookie::generate(&client_info, &masscanned.synack_key).unwrap();
+        client_info.cookie = Some(cookie);
+        assert!(!is_tcb_set(cookie), "expected no TCB entry, found one");
+        /***** TEST PROTOCOL ID IN TCB *****/
+        let payload = [get_dummy_tcp(&client_info), b"GET / HTTP/1.1\r\n".to_vec()].concat();
+        tcp::repl(
+            &TcpPacket::new(&payload).unwrap(),
+            &masscanned,
+            &mut client_info,
+        );
+        assert!(is_tcb_set(cookie), "expected a TCB entry, not found");
+        get_tcb(cookie, |t| {
+            let t = t.unwrap();
+            assert!(t.proto_id == PROTO_HTTP);
+        });
+
+        /***** SENDING MORE DATA *****/
+        let payload = [
+            get_dummy_tcp(&client_info),
+            b"garbage data with no specific format (no protocol)\r\n\r\n".to_vec(),
+        ]
+        .concat();
+        tcp::repl(
+            &TcpPacket::new(&payload).unwrap(),
+            &masscanned,
+            &mut client_info,
+        );
+        assert!(is_tcb_set(cookie), "expected a TCB entry, not found");
+        get_tcb(cookie, |t| {
+            let t = t.unwrap();
+            assert!(t.proto_id == PROTO_HTTP);
+        });
+    }
+
+    #[test]
+    fn test_proto_tcb_proto_state_http() {
+        let mut client_info = ClientInfo::new();
+        let test_ip_addr = Ipv4Addr::new(3, 2, 1, 0);
+        client_info.ip.src = Some(IpAddr::V4(test_ip_addr));
+        client_info.port.src = Some(65001);
+        client_info.port.dst = Some(80);
+        client_info.transport = Some(IpNextHeaderProtocols::Tcp);
+        let masscanned_ip_addr = Ipv4Addr::new(0, 1, 2, 3);
+        client_info.ip.dst = Some(IpAddr::V4(masscanned_ip_addr));
+        let mut ips = HashSet::new();
+        ips.insert(IpAddr::V4(masscanned_ip_addr));
+        /* Construct masscanned context object */
+        let masscanned = Masscanned {
+            synack_key: [0, 0],
+            mac: MacAddr::from_str("00:11:22:33:44:55").expect("error parsing MAC address"),
+            iface: None,
+            ip_addresses: Some(&ips),
+            log: MetaLogger::new(),
+        };
+        let cookie = synackcookie::generate(&client_info, &masscanned.synack_key).unwrap();
+        client_info.cookie = Some(cookie);
+        assert!(!is_tcb_set(cookie), "expected no TCB entry, found one");
+        /***** TEST PROTOCOL ID IN TCB *****/
+        let payload = [get_dummy_tcp(&client_info), b"GET / HTTP/1.1\r\n".to_vec()].concat();
+        tcp::repl(
+            &TcpPacket::new(&payload).unwrap(),
+            &masscanned,
+            &mut client_info,
+        );
+        assert!(is_tcb_set(cookie), "expected a TCB entry, not found");
+        get_tcb(cookie, |t| {
+            let t = t.unwrap();
+            assert!(t.proto_id == PROTO_HTTP);
+            if let Some(ProtocolState::HTTP(_)) = t.proto_state {
+            } else {
+                panic!("expected a HTTP protocole state, found None");
+            }
+        });
+        /***** SENDING MORE DATA *****/
+        let payload = [
+            get_dummy_tcp(&client_info),
+            b"Field: empty\r\n\r\n".to_vec(),
+        ]
+        .concat();
+        /* Should have an answer here */
+        if let None = tcp::repl(
+            &TcpPacket::new(&payload).unwrap(),
+            &masscanned,
+            &mut client_info,
+        ) {
+            panic!("expected an HTTP response, got nothing");
+        }
+        assert!(is_tcb_set(cookie), "expected a TCB entry, not found");
+        get_tcb(cookie, |t| {
+            let t = t.unwrap();
+            assert!(t.proto_id == PROTO_HTTP);
+        })
+    }
+
+    #[test]
+    fn test_proto_tcb_proto_state_rpc() {
+        let mut client_info = ClientInfo::new();
+        let test_ip_addr = Ipv4Addr::new(3, 2, 1, 0);
+        client_info.ip.src = Some(IpAddr::V4(test_ip_addr));
+        client_info.port.src = Some(65002);
+        client_info.port.dst = Some(80);
+        client_info.transport = Some(IpNextHeaderProtocols::Tcp);
+        let masscanned_ip_addr = Ipv4Addr::new(0, 1, 2, 3);
+        client_info.ip.dst = Some(IpAddr::V4(masscanned_ip_addr));
+        let mut ips = HashSet::new();
+        ips.insert(IpAddr::V4(masscanned_ip_addr));
+        /* Construct masscanned context object */
+        let masscanned = Masscanned {
+            synack_key: [0, 0],
+            mac: MacAddr::from_str("00:11:22:33:44:55").expect("error parsing MAC address"),
+            iface: None,
+            ip_addresses: Some(&ips),
+            log: MetaLogger::new(),
+        };
+        let cookie = synackcookie::generate(&client_info, &masscanned.synack_key).unwrap();
+        client_info.cookie = Some(cookie);
+        assert!(!is_tcb_set(cookie), "expected no TCB entry, found one");
+        /***** TEST PROTOCOL ID IN TCB *****/
+        let full_payload = b"\x80\x00\x00\x28\x72\xfe\x1d\x13\x00\x00\x00\x00\x00\x00\x00\x02\x00\x01\x86\xa0\x00\x01\x97\x7c\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00";
+        let payload = [get_dummy_tcp(&client_info), full_payload[0..28].to_vec()].concat();
+        tcp::repl(
+            &TcpPacket::new(&payload).unwrap(),
+            &masscanned,
+            &mut client_info,
+        );
+        assert!(is_tcb_set(cookie), "expected a TCB entry, not found");
+        get_tcb(cookie, |t| {
+            let t = t.unwrap();
+            assert!(t.proto_id == PROTO_RPC_TCP);
+            if let Some(ProtocolState::RPC(_)) = t.proto_state {
+            } else {
+                panic!("expected a RPC protocole state, found None");
+            }
+        });
+        /***** SENDING MORE DATA *****/
+        /* Should have an answer here */
+        let payload = [get_dummy_tcp(&client_info), full_payload[28..].to_vec()].concat();
+        if let None = tcp::repl(
+            &TcpPacket::new(&payload).unwrap(),
+            &masscanned,
+            &mut client_info,
+        ) {
+            panic!("expected a RPC response, got nothing");
+        }
+        assert!(is_tcb_set(cookie), "expected a TCB entry, not found");
+        get_tcb(cookie, |t| {
+            let t = t.unwrap();
+            assert!(t.proto_id == PROTO_RPC_TCP);
+        });
+    }
+}