From 33f6a1629092dd8120eb591d2644efec0e2e4ba9 Mon Sep 17 00:00:00 2001
From: Victor Julien <vjulien@oisf.net>
Date: Fri, 20 Jan 2023 15:27:23 +0100
Subject: [PATCH] smb: file handling cleanups

---
 rust/src/smb/files.rs | 31 +++++++++++++++++++++++--------
 rust/src/smb/smb.rs   |  2 +-
 rust/src/smb/smb1.rs  |  2 +-
 rust/src/smb/smb2.rs  |  6 +++---
 4 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/rust/src/smb/files.rs b/rust/src/smb/files.rs
index 37f60aec21..938ffaffab 100644
--- a/rust/src/smb/files.rs
+++ b/rust/src/smb/files.rs
@@ -55,14 +55,30 @@ pub fn filetracker_newchunk(ft: &mut FileTransferTracker, files: &mut FileContai
         flags: u16, name: &[u8], data: &[u8],
         chunk_offset: u64, chunk_size: u32, is_last: bool, xid: &u32)
 {
-    match unsafe {SURICATA_SMB_FILE_CONFIG} {
-        Some(sfcm) => {
-            ft.new_chunk(sfcm, files, flags, name, data, chunk_offset,
-                    chunk_size, 0, is_last, xid); }
-        None => panic!("no SURICATA_SMB_FILE_CONFIG"),
+    if let Some(sfcm) = unsafe { SURICATA_SMB_FILE_CONFIG } {
+        ft.new_chunk(sfcm, files, flags, name, data, chunk_offset,
+                chunk_size, 0, is_last, xid);
     }
 }
 
+pub fn filetracker_trunc(ft: &mut FileTransferTracker, files: &mut FileContainer,
+        flags: u16)
+{
+    ft.trunc(files, flags);
+}
+
+pub fn filetracker_close(ft: &mut FileTransferTracker, files: &mut FileContainer,
+        flags: u16)
+{
+    ft.close(files, flags);
+}
+
+fn filetracker_update(ft: &mut FileTransferTracker, files: &mut FileContainer,
+        flags: u16, data: &[u8], gap_size: u32) -> u32
+{
+    ft.update(files, flags, data, gap_size)
+}
+
 impl SMBState {
     pub fn new_file_tx(&mut self, fuid: &[u8], file_name: &[u8], direction: Direction)
         -> &mut SMBTransaction
@@ -185,7 +201,7 @@ impl SMBState {
                         let queued_data = tdf.file_tracker.get_queued_size();
                         if queued_data > 2000000 { // TODO should probably be configurable
                             SCLogDebug!("QUEUED size {} while we've seen GAPs. Truncating file.", queued_data);
-                            tdf.file_tracker.trunc(files, flags);
+                            filetracker_trunc(&mut tdf.file_tracker, files, flags);
                         }
                     }
 
@@ -195,8 +211,7 @@ impl SMBState {
                     }
 
                     let file_data = &data[0..data_to_handle_len];
-                    let cs = tdf.file_tracker.update(files, flags, file_data, gap_size);
-                    cs
+                    filetracker_update(&mut tdf.file_tracker, files, flags, file_data, gap_size)
                 } else {
                     0
                 }
diff --git a/rust/src/smb/smb.rs b/rust/src/smb/smb.rs
index 370d48bd0f..e811d13921 100644
--- a/rust/src/smb/smb.rs
+++ b/rust/src/smb/smb.rs
@@ -1098,7 +1098,7 @@ impl SMBState {
                         tx.request_done = true;
                         tx.response_done = true;
                         let (files, flags) = f.files.get(f.direction);
-                        f.file_tracker.trunc(files, flags);
+                        filetracker_trunc(&mut f.file_tracker, files, flags);
                     } else {
                         post_gap_txs = true;
                     }
diff --git a/rust/src/smb/smb1.rs b/rust/src/smb/smb1.rs
index a85b025edd..3ddcf6ddf7 100644
--- a/rust/src/smb/smb1.rs
+++ b/rust/src/smb/smb1.rs
@@ -156,7 +156,7 @@ fn smb1_close_file(state: &mut SMBState, fid: &[u8], direction: Direction)
             if !tx.request_done {
                 SCLogDebug!("closing file tx {} FID {:?}", tx.id, fid);
                 let (files, flags) = tdf.files.get(direction);
-                tdf.file_tracker.close(files, flags);
+                filetracker_close(&mut tdf.file_tracker, files, flags);
                 tx.request_done = true;
                 tx.response_done = true;
                 SCLogDebug!("tx {} is done", tx.id);
diff --git a/rust/src/smb/smb2.rs b/rust/src/smb/smb2.rs
index 0c1daba6d7..b9fb924ffc 100644
--- a/rust/src/smb/smb2.rs
+++ b/rust/src/smb/smb2.rs
@@ -595,7 +595,7 @@ pub fn smb2_request_record<'b>(state: &mut SMBState, r: &Smb2Record<'b>)
                             if !tx.request_done {
                                 if let Some(SMBTransactionTypeData::FILE(ref mut tdf)) = tx.type_data {
                                     let (files, flags) = tdf.files.get(Direction::ToServer);
-                                    tdf.file_tracker.close(files, flags);
+                                    filetracker_close(&mut tdf.file_tracker, files, flags);
                                 }
                             }
                             tx.request_done = true;
@@ -610,7 +610,7 @@ pub fn smb2_request_record<'b>(state: &mut SMBState, r: &Smb2Record<'b>)
                             if !tx.request_done {
                                 if let Some(SMBTransactionTypeData::FILE(ref mut tdf)) = tx.type_data {
                                     let (files, flags) = tdf.files.get(Direction::ToClient);
-                                    tdf.file_tracker.close(files, flags);
+                                    filetracker_close(&mut tdf.file_tracker, files, flags);
                                 }
                             }
                             tx.request_done = true;
@@ -709,7 +709,7 @@ pub fn smb2_response_record<'b>(state: &mut SMBState, r: &Smb2Record<'b>)
                         if !tx.request_done {
                             if let Some(SMBTransactionTypeData::FILE(ref mut tdf)) = tx.type_data {
                                 let (files, flags) = tdf.files.get(Direction::ToClient);
-                                tdf.file_tracker.close(files, flags);
+                                filetracker_close(&mut tdf.file_tracker, files, flags);
                             }
                         }
                         tx.set_status(r.nt_status, false);