From 7ab071a58da141d4316c9452337114c0414b49dd Mon Sep 17 00:00:00 2001
From: Victor Julien <victor@inliniac.net>
Date: Tue, 13 Mar 2018 13:36:19 +0100
Subject: [PATCH] rust/smb: implement minimal record parsing in probing

---
 rust/src/smb/smb.rs          | 29 ++++++++++++-----------------
 src/app-layer-smb-tcp-rust.c | 29 ++++++-----------------------
 2 files changed, 18 insertions(+), 40 deletions(-)

diff --git a/rust/src/smb/smb.rs b/rust/src/smb/smb.rs
index 66c15b623a..a48501f996 100644
--- a/rust/src/smb/smb.rs
+++ b/rust/src/smb/smb.rs
@@ -1696,26 +1696,21 @@ pub extern "C" fn rs_smb_parse_response_tcp_gap(
     return -1;
 }
 
-/// TOSERVER probe function
 #[no_mangle]
-pub extern "C" fn rs_smb_probe_tcp_ts(_input: *const libc::uint8_t, _len: libc::uint32_t)
+pub extern "C" fn rs_smb_probe_tcp(input: *const libc::uint8_t, len: libc::uint32_t)
                                -> libc::int8_t
 {
-//    let slice: &[u8] = unsafe {
-//        std::slice::from_raw_parts(input as *mut u8, len as usize)
-//    };
-    //return smb3_probe(slice, STREAM_TOSERVER);
-    return 1
-}
-/// TOCLIENT probe function
-#[no_mangle]
-pub extern "C" fn rs_smb_probe_tcp_tc(_input: *const libc::uint8_t, _len: libc::uint32_t)
-                               -> libc::int8_t
-{
-//    let slice: &[u8] = unsafe {
-//        std::slice::from_raw_parts(input as *mut u8, len as usize)
-//    };
-    //return smb3_probe(slice, STREAM_TOCLIENT);
+    let slice: &[u8] = unsafe {
+        std::slice::from_raw_parts(input as *mut u8, len as usize)
+    };
+    match parse_nbss_record_partial(slice) {
+        IResult::Done(_, ref hdr) => {
+            if hdr.is_smb() {
+                return 1;
+            }
+        },
+        _ => { },
+    }
     return 1
 }
 
diff --git a/src/app-layer-smb-tcp-rust.c b/src/app-layer-smb-tcp-rust.c
index d2cd338978..f37edfc68d 100644
--- a/src/app-layer-smb-tcp-rust.c
+++ b/src/app-layer-smb-tcp-rust.c
@@ -77,7 +77,7 @@ static int RustSMBTCPParseResponse(Flow *f, void *state,
     return res;
 }
 
-static uint16_t RustSMBTCPProbeTS(Flow *f,
+static uint16_t RustSMBTCPProbe(Flow *f,
         uint8_t *input, uint32_t len, uint32_t *offset)
 {
     SCLogDebug("RustSMBTCPProbe");
@@ -87,24 +87,7 @@ static uint16_t RustSMBTCPProbeTS(Flow *f,
     }
 
     // Validate and return ALPROTO_FAILED if needed.
-    if (!rs_smb_probe_tcp_ts(input, len)) {
-        return ALPROTO_FAILED;
-    }
-
-    return ALPROTO_SMB;
-}
-
-static uint16_t RustSMBTCPProbeTC(Flow *f,
-        uint8_t *input, uint32_t len, uint32_t *offset)
-{
-    SCLogDebug("RustSMBTCPProbe");
-
-    if (len < MIN_REC_SIZE) {
-        return ALPROTO_UNKNOWN;
-    }
-
-    // Validate and return ALPROTO_FAILED if needed.
-    if (!rs_smb_probe_tcp_tc(input, len)) {
+    if (!rs_smb_probe_tcp(input, len)) {
         return ALPROTO_FAILED;
     }
 
@@ -232,20 +215,20 @@ void RegisterRustSMBTCPParsers(void)
 
         if (RunmodeIsUnittests()) {
             AppLayerProtoDetectPPRegister(IPPROTO_TCP, "445", ALPROTO_SMB, 0,
-                    MIN_REC_SIZE, STREAM_TOSERVER, RustSMBTCPProbeTS,
+                    MIN_REC_SIZE, STREAM_TOSERVER, RustSMBTCPProbe,
                     NULL);
         } else {
             int have_cfg = AppLayerProtoDetectPPParseConfPorts("tcp",
                     IPPROTO_TCP, proto_name, ALPROTO_SMB, 0,
-                    MIN_REC_SIZE, RustSMBTCPProbeTS, RustSMBTCPProbeTC);
+                    MIN_REC_SIZE, RustSMBTCPProbe, RustSMBTCPProbe);
             /* if we have no config, we enable the default port 445 */
             if (!have_cfg) {
                 SCLogWarning(SC_ERR_SMB_CONFIG, "no SMB TCP config found, "
                                                 "enabling SMB detection on "
                                                 "port 445.");
                 AppLayerProtoDetectPPRegister(IPPROTO_TCP, "445", ALPROTO_SMB, 0,
-                        MIN_REC_SIZE, STREAM_TOSERVER, RustSMBTCPProbeTS,
-                        RustSMBTCPProbeTC);
+                        MIN_REC_SIZE, STREAM_TOSERVER, RustSMBTCPProbe,
+                        RustSMBTCPProbe);
             }
         }
     } else {
-- 
2.47.2