Merge pull request #21 from ivre/fix-tcb

Keep TCP and proto state in TCB

src/proto/ghost.rs

@@ -21,6 +21,7 @@ use flate2::write::ZlibEncoder;
 use flate2::Compression;
 
 use crate::client::ClientInfo;
+use crate::proto::TCPControlBlock;
 use crate::Masscanned;
 
 pub const GHOST_PATTERN_SIGNATURE: &[u8; 5] = b"Gh0st";
@@ -29,6 +30,7 @@ pub fn repl<'a>(
     _data: &'a [u8],
     _masscanned: &Masscanned,
     _client_info: &mut ClientInfo,
+    _tcb: Option<&mut TCPControlBlock>,
 ) -> Option<Vec<u8>> {
     debug!("receiving Gh0st data, sending one null byte payload");
     // Packet structure:

src/proto/http.rs

@@ -21,6 +21,7 @@ use lazy_static::lazy_static;
 use std::str;
 
 use crate::client::ClientInfo;
+use crate::proto::{ProtocolState as GenericProtocolState, TCPControlBlock};
 use crate::smack::{
     Smack, SmackFlags, BASE_STATE, NO_MATCH, SMACK_CASE_INSENSITIVE, UNANCHORED_STATE,
 };
@@ -62,7 +63,7 @@ const HTTP_STATE_CONTENT: usize = 64;
 
 const HTTP_STATE_FAIL: usize = 0xFFFF;
 
-struct ProtocolState {
+pub struct ProtocolState {
     state: usize,
     state_bis: usize,
     smack_state: usize,
@@ -223,9 +224,28 @@ pub fn repl<'a>(
     data: &'a [u8],
     _masscanned: &Masscanned,
     _client_info: &ClientInfo,
+    tcb: Option<&mut TCPControlBlock>,
 ) -> Option<Vec<u8>> {
     debug!("receiving HTTP data");
-    let mut pstate = ProtocolState::new();
+    let mut state = ProtocolState::new();
+    let mut pstate = {
+        if let Some(t) = tcb {
+            match t.proto_state {
+                None => t.proto_state = Some(GenericProtocolState::HTTP(ProtocolState::new())),
+                Some(GenericProtocolState::HTTP(_)) => {}
+                _ => {
+                    panic!()
+                }
+            };
+            if let Some(GenericProtocolState::HTTP(p)) = &mut t.proto_state {
+                p
+            } else {
+                panic!();
+            }
+        } else {
+            &mut state
+        }
+    };
     http_parse(&mut pstate, data);
     if pstate.state == HTTP_STATE_FAIL {
         debug!("data in not correctly formatted - not responding");

src/proto/mod.rs

@@ -25,7 +25,7 @@ use crate::smack::{Smack, SmackFlags, BASE_STATE, NO_MATCH, SMACK_CASE_SENSITIVE
 use crate::Masscanned;
 
 mod http;
-use http::HTTP_VERBS;
+use http::{ProtocolState as HTTPProtocolState, HTTP_VERBS};
 
 mod stun;
 use stun::{STUN_PATTERN_CHANGE_REQUEST, STUN_PATTERN_EMPTY, STUN_PATTERN_MAGIC};
@@ -37,11 +37,12 @@ mod ghost;
 use ghost::GHOST_PATTERN_SIGNATURE;
 
 mod rpc;
-use rpc::{RPC_CALL_TCP, RPC_CALL_UDP};
+use rpc::{ProtocolState as RPCProtocolState, RPC_CALL_TCP, RPC_CALL_UDP};
 
 mod smb;
 use smb::{SMB1_PATTERN_MAGIC, SMB2_PATTERN_MAGIC};
 
+const PROTO_NONE: usize = 0;
 const PROTO_HTTP: usize = 1;
 const PROTO_STUN: usize = 2;
 const PROTO_SSH: usize = 3;
@@ -51,8 +52,18 @@ const PROTO_RPC_UDP: usize = 6;
 const PROTO_SMB1: usize = 7;
 const PROTO_SMB2: usize = 8;
 
-struct TCPControlBlock {
-    proto_state: usize,
+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! {
@@ -126,26 +137,33 @@ pub fn repl<'a>(
 ) -> 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 {
         /* proto over TCP */
         let cookie = client_info.cookie.unwrap();
-        let mut ct = CONTABLE.lock().unwrap();
         if !ct.contains_key(&cookie) {
             ct.insert(
                 cookie,
                 TCPControlBlock {
-                    proto_state: BASE_STATE,
+                    smack_state: BASE_STATE,
+                    proto_id: PROTO_NONE,
+                    proto_state: None,
                 },
             );
         }
         let mut i = 0;
-        let mut tcb = ct.get_mut(&cookie).unwrap();
-        let mut state = tcb.proto_state;
-        id = PROTO_SMACK.search_next(&mut state, data, &mut i);
-        tcb.proto_state = state;
+        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;
@@ -158,15 +176,18 @@ pub fn repl<'a>(
     }
     /* proto over else (e.g., UDP) */
     match id {
-        PROTO_HTTP => http::repl(data, masscanned, client_info),
-        PROTO_STUN => stun::repl(data, masscanned, &mut client_info),
-        PROTO_SSH => ssh::repl(data, masscanned, &mut client_info),
-        PROTO_GHOST => ghost::repl(data, masscanned, &mut client_info),
-        PROTO_RPC_TCP => rpc::repl_tcp(data, masscanned, &mut client_info),
-        PROTO_RPC_UDP => rpc::repl_udp(data, masscanned, &mut client_info),
-        PROTO_SMB1 => smb::repl_smb1(data, masscanned, &mut client_info),
-        PROTO_SMB2 => smb::repl_smb2(data, masscanned, &mut client_info),
+        PROTO_HTTP => http::repl(data, masscanned, client_info, tcb),
+        PROTO_STUN => stun::repl(data, masscanned, &mut client_info, tcb),
+        PROTO_SSH => ssh::repl(data, masscanned, &mut client_info, tcb),
+        PROTO_GHOST => ghost::repl(data, masscanned, &mut client_info, tcb),
+        PROTO_RPC_TCP => rpc::repl_tcp(data, masscanned, &mut client_info, tcb),
+        PROTO_RPC_UDP => rpc::repl_udp(data, masscanned, &mut client_info, tcb),
+        PROTO_SMB1 => smb::repl_smb1(data, masscanned, &mut client_info, tcb),
+        PROTO_SMB2 => smb::repl_smb2(data, masscanned, &mut client_info, tcb),
         _ => {
+            if let Some(t) = &mut tcb {
+                t.proto_id = PROTO_NONE;
+            }
             debug!("id: {}", id);
             None
         }
@@ -183,6 +204,176 @@ 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() {

src/proto/rpc.rs

@@ -19,6 +19,7 @@ use std::convert::TryInto;
 use std::net::IpAddr;
 
 use crate::client::ClientInfo;
+use crate::proto::{ProtocolState as GenericProtocolState, TCPControlBlock};
 use crate::Masscanned;
 
 // last fragment (1 bit) + fragment len (31 bits) / length XID (random) / message type: call (0) / RPC version (0-255) / Program: Portmap (99840 - 100095) / Program version (*, random versions used, see below) / / Procedure: ??? (0-255)
@@ -47,7 +48,7 @@ enum RpcState {
 }
 
 #[derive(Debug)]
-struct ProtocolState {
+pub struct ProtocolState {
     state: RpcState,
     last_frag: bool,
     frag_len: u32,
@@ -210,7 +211,7 @@ fn push_string_pad(buffer: &mut Vec<u8>, data: String) {
     }
 }
 
-fn build_repl_portmap(pstate: ProtocolState, client_info: &ClientInfo) -> Vec<u8> {
+fn build_repl_portmap(pstate: &mut ProtocolState, client_info: &ClientInfo) -> Vec<u8> {
     let mut resp = Vec::<u8>::new();
     match pstate.procedure {
         // 0 => {}
@@ -313,7 +314,7 @@ fn build_repl_portmap(pstate: ProtocolState, client_info: &ClientInfo) -> Vec<u8
     resp
 }
 
-fn build_repl_unknownprog(pstate: ProtocolState, _client_info: &ClientInfo) -> Vec<u8> {
+fn build_repl_unknownprog(pstate: &mut ProtocolState, _client_info: &ClientInfo) -> Vec<u8> {
     warn!(
         "Unknown program {}, procedure {}: accepted state 1",
         pstate.program, pstate.procedure
@@ -322,7 +323,7 @@ fn build_repl_unknownprog(pstate: ProtocolState, _client_info: &ClientInfo) -> V
     vec![0, 0, 0, 1]
 }
 
-fn build_repl(pstate: ProtocolState, client_info: &ClientInfo) -> Vec<u8> {
+fn build_repl(pstate: &mut ProtocolState, client_info: &ClientInfo) -> Vec<u8> {
     // TODO: test RPC versions, drop non calls?
     let mut resp = Vec::<u8>::new();
     push_u32(&mut resp, pstate.xid);
@@ -373,8 +374,27 @@ pub fn repl_tcp<'a>(
     data: &'a [u8],
     _masscanned: &Masscanned,
     client_info: &ClientInfo,
+    tcb: Option<&mut TCPControlBlock>,
 ) -> Option<Vec<u8>> {
-    let mut pstate = ProtocolState::new();
+    let mut state = ProtocolState::new();
+    let mut pstate = {
+        if let Some(t) = tcb {
+            match t.proto_state {
+                None => t.proto_state = Some(GenericProtocolState::RPC(ProtocolState::new())),
+                Some(GenericProtocolState::RPC(_)) => {}
+                _ => {
+                    panic!()
+                }
+            };
+            if let Some(GenericProtocolState::RPC(p)) = &mut t.proto_state {
+                p
+            } else {
+                panic!();
+            }
+        } else {
+            &mut state
+        }
+    };
     rpc_parse(&mut pstate, data);
     // warn!("RPC {:#?}", pstate);
     let resp = match pstate.state {
@@ -402,6 +422,7 @@ pub fn repl_udp<'a>(
     data: &'a [u8],
     _masscanned: &Masscanned,
     client_info: &ClientInfo,
+    _tcb: Option<&mut TCPControlBlock>,
 ) -> Option<Vec<u8>> {
     let mut pstate = ProtocolState::new();
     pstate.state = RpcState::Xid;
@@ -410,7 +431,7 @@ pub fn repl_udp<'a>(
     rpc_parse(&mut pstate, data);
     // warn!("RPC {:#?}", pstate);
     match pstate.state {
-        RpcState::End => Some(build_repl(pstate, client_info)),
+        RpcState::End => Some(build_repl(&mut pstate, client_info)),
         _ => None,
     }
 }
@@ -452,7 +473,7 @@ mod tests {
         assert!(pstate.creds_data.len() == 0);
         assert!(pstate.verif_flavor == 0);
         assert!(pstate.verif_data.len() == 0);
-        let resp = build_repl(pstate, &CLIENT_INFO);
+        let resp = build_repl(&mut pstate, &CLIENT_INFO);
         assert!(resp == b"\x72\xfe\x1d\x13\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x02\x00\x00\x00\x04");
     }
 
@@ -471,7 +492,7 @@ mod tests {
         assert!(pstate.creds_data.len() == 0);
         assert!(pstate.verif_flavor == 0);
         assert!(pstate.verif_data.len() == 0);
-        let resp = build_repl(pstate, &CLIENT_INFO);
+        let resp = build_repl(&mut pstate, &CLIENT_INFO);
         assert!(resp == b"\x72\xfe\x1d\x13\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x02\x00\x00\x00\x04");
     }
 
@@ -491,7 +512,7 @@ mod tests {
         assert!(pstate.creds_data.len() == 0);
         assert!(pstate.verif_flavor == 0);
         assert!(pstate.verif_data.len() == 0);
-        let resp = build_repl(pstate, &CLIENT_INFO);
+        let resp = build_repl(&mut pstate, &CLIENT_INFO);
         assert!(resp == b"\x72\xfe\x1d\x13\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x02\x00\x00\x00\x04");
     }
 
@@ -519,7 +540,7 @@ mod tests {
         assert!(pstate.creds_data.len() == 0);
         assert!(pstate.verif_flavor == 0);
         assert!(pstate.verif_data.len() == 0);
-        let resp = build_repl(pstate, &CLIENT_INFO);
+        let resp = build_repl(&mut pstate, &CLIENT_INFO);
         assert!(resp == b"\x72\xfe\x1d\x13\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x02\x00\x00\x00\x04");
     }
 

src/proto/smb.rs

@@ -20,6 +20,7 @@ use std::convert::TryInto;
 use std::time::SystemTime;
 
 use crate::client::ClientInfo;
+use crate::proto::TCPControlBlock;
 use crate::Masscanned;
 
 // NBTSession + SMB Header
@@ -962,6 +963,7 @@ pub fn repl_smb1<'a>(
     data: &'a [u8],
     _masscanned: &Masscanned,
     _client_info: &ClientInfo,
+    _tcb: Option<&mut TCPControlBlock>,
 ) -> Option<Vec<u8>> {
     let mut nbtsession: NBTSession<SMB1Header> = NBTSession::new();
     for byte in data {
@@ -974,6 +976,7 @@ pub fn repl_smb2<'a>(
     data: &'a [u8],
     _masscanned: &Masscanned,
     _client_info: &ClientInfo,
+    _tcb: Option<&mut TCPControlBlock>,
 ) -> Option<Vec<u8>> {
     let mut nbtsession: NBTSession<SMB2Header> = NBTSession::new();
     for byte in data {
@@ -1057,7 +1060,7 @@ mod tests {
         };
         let client_info = ClientInfo::new();
         let answer =
-            repl_smb1(SMB1_REQ_PAYLOAD, &masscanned, &client_info).expect("Error: no answer");
+            repl_smb1(SMB1_REQ_PAYLOAD, &masscanned, &client_info, None).expect("Error: no answer");
         let expected = [
             0, 0, 1, 149, 255, 83, 77, 66, 114, 0, 0, 0, 0, 152, 7, 200, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 254, 255, 0, 0, 0, 0, 17, 1, 0, 3, 50, 0, 50, 0, 0, 0, 1, 0, 0, 0, 1, 0,
@@ -1129,7 +1132,7 @@ mod tests {
         };
         let client_info = ClientInfo::new();
         let answer =
-            repl_smb2(SMB2_REQ_PAYLOAD, &masscanned, &client_info).expect("Error: no answer");
+            repl_smb2(SMB2_REQ_PAYLOAD, &masscanned, &client_info, None).expect("Error: no answer");
         let expected = [
             0, 0, 1, 192, 254, 83, 77, 66, 64, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,

src/proto/ssh.rs

@@ -19,6 +19,7 @@ use log::*;
 use std::str;
 
 use crate::client::ClientInfo;
+use crate::proto::TCPControlBlock;
 use crate::Masscanned;
 
 pub const SSH_PATTERN_CLIENT_PROTOCOL: &[u8; 7] = b"SSH-2.0";
@@ -27,6 +28,7 @@ pub fn repl<'a>(
     data: &'a [u8],
     _masscanned: &Masscanned,
     mut _client_info: &mut ClientInfo,
+    _tcb: Option<&mut TCPControlBlock>,
 ) -> Option<Vec<u8>> {
     debug!("receiving SSH data");
     let repl_data = b"SSH-2.0-1\r\n".to_vec();

src/proto/stun.rs

@@ -24,6 +24,7 @@ use byteorder::{BigEndian, ByteOrder};
 use std::io;
 
 use crate::client::ClientInfo;
+use crate::proto::TCPControlBlock;
 use crate::Masscanned;
 
 /* RFC 5389: The magic cookie field MUST contain the fixed value 0x2112A442 in
@@ -354,6 +355,7 @@ pub fn repl<'a>(
     data: &'a [u8],
     _masscanned: &Masscanned,
     mut client_info: &mut ClientInfo,
+    _tcb: Option<&mut TCPControlBlock>,
 ) -> Option<Vec<u8>> {
     debug!("receiving STUN data");
     let stun_req: StunPacket = if let Ok(s) = StunPacket::new(&data) {
@@ -443,7 +445,7 @@ mod tests {
             ip_addresses: Some(&ips),
             log: MetaLogger::new(),
         };
-        let payload_resp = if let Some(r) = repl(payload, &masscanned, &mut client_info) {
+        let payload_resp = if let Some(r) = repl(payload, &masscanned, &mut client_info, None) {
             r
         } else {
             panic!("expected an answer, got None");
@@ -507,7 +509,7 @@ mod tests {
         client_info.ip.dst = Some(IpAddr::V6(masscanned_ip_addr));
         client_info.port.src = Some(55000);
         client_info.port.dst = Some(65000);
-        let payload_resp = if let Some(r) = repl(payload, &masscanned, &mut client_info) {
+        let payload_resp = if let Some(r) = repl(payload, &masscanned, &mut client_info, None) {
             r
         } else {
             panic!("expected an answer, got None");
@@ -559,7 +561,7 @@ mod tests {
         client_info.ip.dst = Some(IpAddr::V4(masscanned_ip_addr));
         client_info.port.src = Some(55000);
         client_info.port.dst = Some(65000);
-        let payload_resp = if let Some(r) = repl(payload, &masscanned, &mut client_info) {
+        let payload_resp = if let Some(r) = repl(payload, &masscanned, &mut client_info, None) {
             r
         } else {
             panic!("expected an answer, got None");
@@ -609,7 +611,7 @@ mod tests {
         client_info.ip.dst = Some(IpAddr::V4(masscanned_ip_addr));
         client_info.port.src = Some(55000);
         client_info.port.dst = Some(65535);
-        let payload_resp = if let Some(r) = repl(payload, &masscanned, &mut client_info) {
+        let payload_resp = if let Some(r) = repl(payload, &masscanned, &mut client_info, None) {
             r
         } else {
             panic!("expected an answer, got None");

test/src/tests/http.py

@@ -74,6 +74,76 @@ def test_ipv4_tcp_http():
         assert tcp.payload.load.startswith(b"HTTP/1.1 401 Unauthorized\n")
 
 
+@test
+def test_ipv4_tcp_http_segmented():
+    sport = 24593
+    dports = [80, 443, 5000, 53228]
+    for dport in dports:
+        seq_init = int(RandInt())
+        syn = (
+            Ether(dst=MAC_ADDR)
+            / IP(dst=IPV4_ADDR)
+            / TCP(flags="S", sport=sport, dport=dport, seq=seq_init)
+        )
+        syn_ack = srp1(syn, timeout=1)
+        assert syn_ack is not None, "expecting answer, got nothing"
+        check_ip_checksum(syn_ack)
+        assert TCP in syn_ack, "expecting TCP, got %r" % syn_ack.summary()
+        syn_ack = syn_ack[TCP]
+        assert syn_ack.flags == "SA", "expecting TCP SA, got %r" % syn_ack.flags
+        ack = (
+            Ether(dst=MAC_ADDR)
+            / IP(dst=IPV4_ADDR)
+            / TCP(
+                flags="A",
+                sport=sport,
+                dport=dport,
+                seq=seq_init + 1,
+                ack=syn_ack.seq + 1,
+            )
+        )
+        _ = srp1(ack, timeout=1)
+        # request is not complete yet
+        req = (
+            Ether(dst=MAC_ADDR)
+            / IP(dst=IPV4_ADDR)
+            / TCP(
+                flags="PA",
+                sport=sport,
+                dport=dport,
+                seq=seq_init + 1,
+                ack=syn_ack.seq + 1,
+            )
+            / Raw("GET / HTTP/1.1\r\n")
+        )
+        resp = srp1(req, timeout=1)
+        assert resp is not None, "expecting answer, got nothing"
+        check_ip_checksum(resp)
+        assert TCP in resp, "expecting TCP, got %r" % resp.summary()
+        assert resp[TCP].flags == "A", (
+            'expecting TCP flag "A", got %r' % resp[TCP].flags
+        )
+        req = (
+            Ether(dst=MAC_ADDR)
+            / IP(dst=IPV4_ADDR)
+            / TCP(
+                flags="PA",
+                sport=sport,
+                dport=dport,
+                seq=seq_init + len(req) + 1,
+                ack=syn_ack.seq + 1,
+            )
+            / Raw("\r\n")
+        )
+        resp = srp1(req, timeout=1)
+        assert resp is not None, "expecting answer, got nothing"
+        check_ip_checksum(resp)
+        assert TCP in resp, "expecting TCP, got %r" % resp.summary()
+        tcp = resp[TCP]
+        assert tcp.flags == "PA"
+        assert tcp.payload.load.startswith(b"HTTP/1.1 401 Unauthorized\n")
+
+
 @test
 def test_ipv4_tcp_http_incomplete():
     sport = 24595
@@ -126,7 +196,7 @@ def test_ipv4_tcp_http_incomplete():
 
 @test
 def test_ipv6_tcp_http():
-    sport = 24592
+    sport = 24594
     dports = [80, 443, 5000, 53228]
     for dport in dports:
         seq_init = int(RandInt())
@@ -213,7 +283,7 @@ def test_ipv6_udp_http():
 
 @test
 def test_ipv4_tcp_http_ko():
-    sport = 24592
+    sport = 24596
     dports = [80, 443, 5000, 53228]
     for dport in dports:
         seq_init = int(RandInt())
@@ -277,7 +347,7 @@ def test_ipv4_udp_http_ko():
 
 @test
 def test_ipv6_tcp_http_ko():
-    sport = 24592
+    sport = 24597
     dports = [80, 443, 5000, 53228]
     for dport in dports:
         seq_init = int(RandInt())

test/src/tests/rpc.py

@@ -56,17 +56,26 @@ def test_rpc_nmap():
         result = results[0]
         assert len(result["ports"]) == 1, f"Expected 1 port, got {len(result['ports'])}"
         port = result["ports"][0]
-        assert port["port"] == 111 and port["protocol"] == (
+        assert port["port"] == 111, f"Expected port 111, got {port['port']}"
+        assert port["protocol"] == (
             "tcp" if scan == "S" else "udp"
-        )
-        assert port["service_name"] in {"rpcbind", "nfs"}
-        assert port["service_extrainfo"] in {"RPC #100000", "RPC #100003"}
+        ), f"Unexpected proto {port['protocol']} for scan {scan}"
+        assert port["service_name"] in {
+            "rpcbind",
+            "nfs",
+        }, f"Unexpected service_name: {port['service_name']}"
+        assert port["service_extrainfo"] in {
+            "RPC #100000",
+            "RPC #100003",
+        }, f"Unexpected service_extrainfo: {port['service_extrainfo']}"
         assert (
             len(port["scripts"]) == 1
         ), f"Expected 1 script, got {len(port['scripts'])}"
         script = port["scripts"][0]
         assert script["id"] == "rpcinfo", "Expected rpcinfo script, not found"
-        assert len(script["rpcinfo"]) == 1
+        assert (
+            len(script["rpcinfo"]) == 1
+        ), f"Expected 1 rpcinfo, got {len(script['rpcinfo'])}"
 
 
 @test