From 4d58aaae90573dc7d860e640e6ec650a6e0248ef Mon Sep 17 00:00:00 2001
From: Victor Julien <victor@inliniac.net>
Date: Sat, 24 Mar 2018 16:17:20 +0100
Subject: [PATCH] smb: clean up partial read/write record handling

---
 rust/src/smb/smb.rs  | 12 ++++++++++++
 rust/src/smb/smb1.rs |  9 +++------
 rust/src/smb/smb2.rs |  9 ++-------
 3 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/rust/src/smb/smb.rs b/rust/src/smb/smb.rs
index 60f1507d29..29db1d3fc4 100644
--- a/rust/src/smb/smb.rs
+++ b/rust/src/smb/smb.rs
@@ -1107,6 +1107,18 @@ impl SMBState {
         }
     }
 
+    pub fn set_file_left(&mut self, direction: u8, rec_size: u32, data_size: u32, fuid: Vec<u8>)
+    {
+        let left = if data_size >= rec_size { 0 } else { rec_size - data_size };
+        if direction == STREAM_TOSERVER {
+            self.file_ts_left = left;
+            self.file_ts_guid = fuid;
+        } else {
+            self.file_tc_left = left;
+            self.file_tc_guid = fuid;
+        }
+    }
+
     pub fn set_skip(&mut self, direction: u8, rec_size: u32, data_size: u32)
     {
         let skip = if data_size >= rec_size { 0 } else { rec_size - data_size };
diff --git a/rust/src/smb/smb1.rs b/rust/src/smb/smb1.rs
index ebaa0314bc..a45508b8c8 100644
--- a/rust/src/smb/smb1.rs
+++ b/rust/src/smb/smb1.rs
@@ -856,9 +856,8 @@ pub fn smb1_write_request_record<'b>(state: &mut SMBState, r: &SmbRecord<'b>)
                     tx.vercmd.set_smb1_cmd(SMB1_COMMAND_WRITE_ANDX);
                 }
             }
-            state.file_ts_left = rd.len - rd.data.len() as u32;
-            state.file_ts_guid = file_fid.to_vec();
-            SCLogDebug!("SMBv1 WRITE RESPONSE: {} bytes left", state.file_tc_left);
+
+            state.set_file_left(STREAM_TOSERVER, rd.len, rd.data.len() as u32, file_fid.to_vec());
 
             if r.command == SMB1_COMMAND_WRITE_AND_CLOSE {
                 SCLogDebug!("closing FID {:?}", file_fid);
@@ -939,9 +938,7 @@ pub fn smb1_read_response_record<'b>(state: &mut SMBState, r: &SmbRecord<'b>)
                     smb_read_dcerpc_record(state, vercmd, hdr, &pure_fid, &rd.data);
                 }
 
-                state.file_tc_left = rd.len - rd.data.len() as u32;
-                state.file_tc_guid = file_fid.to_vec();
-                SCLogDebug!("SMBv1 READ RESPONSE: {} bytes left", state.file_tc_left);
+                state.set_file_left(STREAM_TOCLIENT, rd.len, rd.data.len() as u32, file_fid.to_vec());
             }
             _ => {
                 events.push(SMBEvent::MalformedData);
diff --git a/rust/src/smb/smb2.rs b/rust/src/smb/smb2.rs
index 0da71ab617..fd843c88cf 100644
--- a/rust/src/smb/smb2.rs
+++ b/rust/src/smb/smb2.rs
@@ -186,9 +186,7 @@ pub fn smb2_read_response_record<'b>(state: &mut SMBState, r: &Smb2Record<'b>)
                 }
             }
 
-            state.file_tc_left = rd.len - rd.data.len() as u32;
-            state.file_tc_guid = file_guid.to_vec();
-            SCLogDebug!("SMBv2 READ RESPONSE: {} bytes left", state.file_tc_left);
+            state.set_file_left(STREAM_TOCLIENT, rd.len, rd.data.len() as u32, file_guid.to_vec());
         }
         _ => {
             state.set_event(SMBEvent::MalformedData);
@@ -257,10 +255,7 @@ pub fn smb2_write_request_record<'b>(state: &mut SMBState, r: &Smb2Record<'b>)
                             r.session_id, r.tree_id, 0); // TODO move into new_file_tx
                 }
             }
-            state.file_ts_left = wr.wr_len - wr.data.len() as u32;
-            state.file_ts_guid = file_guid.to_vec();
-            SCLogDebug!("SMBv2 WRITE REQUEST: {} bytes left", state.file_ts_left);
-
+            state.set_file_left(STREAM_TOSERVER, wr.wr_len, wr.data.len() as u32, file_guid.to_vec());
         },
         _ => {
             state.set_event(SMBEvent::MalformedData);