Add TCB argument to repl functions in protos

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

@@ -25,6 +25,7 @@ use crate::smack::{
     Smack, SmackFlags, BASE_STATE, NO_MATCH, SMACK_CASE_INSENSITIVE, UNANCHORED_STATE,
 };
 use crate::Masscanned;
+use crate::proto::{TCPControlBlock};
 
 pub const HTTP_VERBS: [&str; 9] = [
     "GET", "PUT", "POST", "HEAD", "DELETE", "CONNECT", "OPTIONS", "TRACE", "PATCH",
@@ -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,6 +224,7 @@ 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();

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,7 +37,7 @@ 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};
@@ -51,8 +51,16 @@ const PROTO_RPC_UDP: usize = 6;
 const PROTO_SMB1: usize = 7;
 const PROTO_SMB2: usize = 8;
 
-struct TCPControlBlock {
+enum ProtocolState {
-    proto_state: usize,
+    HTTPProtocolState,
+    RPCProtocolState,
+}
+
+pub struct TCPControlBlock {
+    /* state used to detect protocols (not specific) */
+    smack_state: usize,
+    /* internal state of protocol parser (e.g., HTTP parsing) */
+    proto_state: Option<ProtocolState>,
 }
 
 lazy_static! {
@@ -126,26 +134,29 @@ 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_state: None,
                 },
             );
         }
         let mut i = 0;
-        let mut tcb = ct.get_mut(&cookie).unwrap();
+        let mut t = ct.get_mut(&cookie).unwrap();
-        let mut state = tcb.proto_state;
+        let mut state = t.smack_state;
         id = PROTO_SMACK.search_next(&mut state, data, &mut i);
-        tcb.proto_state = state;
+        t.smack_state = state;
+        tcb = Some(t);
     } else {
         /* proto over else (e.g., UDP) */
         let mut i = 0;
@@ -158,14 +169,14 @@ pub fn repl<'a>(
     }
     /* proto over else (e.g., UDP) */
     match id {
-        PROTO_HTTP => http::repl(data, masscanned, client_info),
+        PROTO_HTTP => http::repl(data, masscanned, client_info, tcb),
-        PROTO_STUN => stun::repl(data, masscanned, &mut client_info),
+        PROTO_STUN => stun::repl(data, masscanned, &mut client_info, tcb),
-        PROTO_SSH => ssh::repl(data, masscanned, &mut client_info),
+        PROTO_SSH => ssh::repl(data, masscanned, &mut client_info, tcb),
-        PROTO_GHOST => ghost::repl(data, masscanned, &mut client_info),
+        PROTO_GHOST => ghost::repl(data, masscanned, &mut client_info, tcb),
-        PROTO_RPC_TCP => rpc::repl_tcp(data, masscanned, &mut client_info),
+        PROTO_RPC_TCP => rpc::repl_tcp(data, masscanned, &mut client_info, tcb),
-        PROTO_RPC_UDP => rpc::repl_udp(data, masscanned, &mut client_info),
+        PROTO_RPC_UDP => rpc::repl_udp(data, masscanned, &mut client_info, tcb),
-        PROTO_SMB1 => smb::repl_smb1(data, masscanned, &mut client_info),
+        PROTO_SMB1 => smb::repl_smb1(data, masscanned, &mut client_info, tcb),
-        PROTO_SMB2 => smb::repl_smb2(data, masscanned, &mut client_info),
+        PROTO_SMB2 => smb::repl_smb2(data, masscanned, &mut client_info, tcb),
         _ => {
             debug!("id: {}", id);
             None

src/proto/rpc.rs

@@ -19,6 +19,7 @@ use std::convert::TryInto;
 use std::net::IpAddr;
 
 use crate::client::ClientInfo;
+use crate::proto::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,
@@ -373,6 +374,7 @@ 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();
     rpc_parse(&mut pstate, data);
@@ -402,6 +404,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;

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");