Resolve TODOs, add tests for signing + query dir, fix tree disconnect bug

afiffon · afiffon · commit 145606a24427 · 2025-03-20T19:37:19.000+02:00
diff --git a/smb/src/connection.rs b/smb/src/connection.rs
@@ -409,7 +409,6 @@ impl ConnectionMessageHandler {
 impl MessageHandler for ConnectionMessageHandler {
     #[maybe_async]
     async fn sendo(&self, mut msg: OutgoingMessage) -> crate::Result<SendMessageResult> {
-        // TODO: Add assertion in the struct regarding the selected dialect!
         let priority_value = match self.conn_info.get() {
             Some(neg_info) => match neg_info.negotiation.dialect_rev {
                 Dialect::Smb0311 => 1,
diff --git a/smb/src/packets/fscc/query_file_info.rs b/smb/src/packets/fscc/query_file_info.rs
@@ -113,11 +113,9 @@ pub struct FileCompressionInformation {
     pub compression_unit: u8,
     pub chunk_shift: u8,
     pub cluster_shift: u8,
-    #[br(parse_with = binrw::helpers::read_u24)]
-    #[br(assert(reserved == 0))]
-    #[bw(align_before = 4)]
-    #[bw(ignore)]
-    reserved: u32, // 3-bytes. TODO: Define a normal u24 type for such cases?
+
+    #[bw(calc = [0; 3])]
+    _reserved: [u8; 3],
 }
 
 #[binrw::binrw]
diff --git a/smb/src/packets/smb2/create.rs b/smb/src/packets/smb2/create.rs
@@ -451,15 +451,13 @@ pub struct DurableHandleRequest {
 #[derive(Debug, PartialEq, Eq, Default)]
 pub struct DurableHandleResponse {
     #[bw(calc = 0)]
-    #[br(assert(response == 0))]
-    response: u128,
+    _reserved: u64,
 }
 
-/// TODO: SMB2_FILEID
 #[binrw::binrw]
 #[derive(Debug, PartialEq, Eq)]
 pub struct DurableHandleReconnect {
-    pub durable_request: u128,
+    pub durable_request: FileId,
 }
 #[binrw::binrw]
 #[derive(Debug, PartialEq, Eq, Default)]
diff --git a/smb/src/packets/smb2/ioctl/msg.rs b/smb/src/packets/smb2/ioctl/msg.rs
@@ -36,7 +36,7 @@ pub struct IoctlRequest {
     #[bw(calc = 0)]
     #[br(assert(reserved2 == 0))]
     reserved2: u32,
-    // TODO: treat FSCTLs differently when having a concrerte data structure.
+
     #[bw(write_with = PosMarker::write_aoff_size, args(&_input_offset, &_input_count))]
     #[br(map_stream = |s| s.take_seek(_input_count.value as u64), args(ctl_code, flags))]
     pub buffer: IoctlReqData,
diff --git a/smb/src/packets/smb2/negotiate.rs b/smb/src/packets/smb2/negotiate.rs
@@ -22,7 +22,7 @@ pub struct NegotiateRequest {
     _reserved: u16,
     pub capabilities: GlobalCapabilities,
     pub client_guid: Guid,
-    // TODO: The 3 fields below are possibly a union in older versions of SMB.
+
     #[bw(calc = PosMarker::default())]
     negotiate_context_offset: PosMarker<u32>,
     #[bw(try_calc(u16::try_from(negotiate_context_list.as_ref().map(|v| v.len()).unwrap_or(0))))]
@@ -32,7 +32,7 @@ pub struct NegotiateRequest {
     reserved2: u16,
     #[br(count = dialect_count)]
     pub dialects: Vec<Dialect>,
-    // Only on SMB 3.1.1 we have negotiate contexts.
+    // Only on SMB 3.1.1 supporting clients we have negotiation contexts.
     // Align to 8 bytes.
     #[brw(if(dialects.contains(&Dialect::Smb0311)), align_before = 8)]
     #[br(count = negotiate_context_count, seek_before = SeekFrom::Start(negotiate_context_offset.value as u64))]
diff --git a/smb/src/packets/smb2/query_dir.rs b/smb/src/packets/smb2/query_dir.rs
@@ -80,6 +80,8 @@ impl QueryDirectoryResponse {
 
 #[cfg(test)]
 mod tests {
+    use time::macros::datetime;
+
     use super::*;
     #[test]
     pub fn test_both_directory_information_attribute_parse() {
@@ -135,6 +137,100 @@ mod tests {
         .read_output()
         .unwrap();
 
-        // TODO: Test the contents of the result, not just that it parses.
+        assert_eq!(
+            vec![
+                ChainedItem::new(FileIdBothDirectoryInformationInner {
+                    file_index: 0,
+                    creation_time: FileTime::from(datetime!(2024-12-11 12:32:31.7084985)),
+                    last_access_time: FileTime::from(datetime!(2025-01-03 10:18:15.6499175)),
+                    last_write_time: FileTime::from(datetime!(2024-12-27 14:22:59.9648231)),
+                    change_time: FileTime::from(datetime!(2024-12-27 14:22:59.9648231)),
+                    end_of_file: 0,
+                    allocation_size: 0,
+                    file_attributes: FileAttributes::new().with_directory(true),
+                    ea_size: 0,
+                    short_name_length: 0,
+                    short_name: [0; 12],
+                    fild_id: 562949953454203,
+                    file_name: ".".into(),
+                }),
+                ChainedItem::new(FileIdBothDirectoryInformationInner {
+                    file_index: 0,
+                    creation_time: FileTime::from(datetime!(2024-12-11 9:25:15.4208828)),
+                    last_access_time: FileTime::from(datetime!(2025-01-02 19:05:31.8723088)),
+                    last_write_time: FileTime::from(datetime!(2024-12-11 12:32:35.4544738)),
+                    change_time: FileTime::from(datetime!(2024-12-11 12:32:35.4544738)),
+                    end_of_file: 0,
+                    allocation_size: 0,
+                    file_attributes: FileAttributes::new().with_directory(true),
+                    ea_size: 0,
+                    short_name_length: 0,
+                    short_name: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,],
+                    fild_id: 1125899906967338,
+                    file_name: "..".into(),
+                }),
+                ChainedItem::new(FileIdBothDirectoryInformationInner {
+                    file_index: 0,
+                    creation_time: FileTime::from(datetime!(2024-12-27 14:22:48.7929947)),
+                    last_access_time: FileTime::from(datetime!(2024-12-27 14:22:48.7929947)),
+                    last_write_time: FileTime::from(datetime!(2024-12-27 14:22:48.7929947)),
+                    change_time: FileTime::from(datetime!(2024-12-27 14:22:49.7460575)),
+                    end_of_file: 0,
+                    allocation_size: 0,
+                    file_attributes: FileAttributes::new().with_archive(true),
+                    ea_size: 0,
+                    short_name_length: 0,
+                    short_name: [0; 12],
+                    fild_id: 2814749767148784,
+                    file_name: "a.txt".into(),
+                }),
+                ChainedItem::new(FileIdBothDirectoryInformationInner {
+                    file_index: 0,
+                    creation_time: FileTime::from(datetime!(2024-12-27 14:22:51.5742424)),
+                    last_access_time: FileTime::from(datetime!(2024-12-27 14:23:06.9505662)),
+                    last_write_time: FileTime::from(datetime!(2024-12-27 14:23:06.9505662)),
+                    change_time: FileTime::from(datetime!(2024-12-27 14:23:06.9505662)),
+                    end_of_file: 6,
+                    allocation_size: 8,
+                    file_attributes: FileAttributes::new().with_archive(true),
+                    ea_size: 0,
+                    short_name_length: 0,
+                    short_name: [0; 12],
+                    fild_id: 1125899906906297,
+                    file_name: "b.txt".into(),
+                }),
+                ChainedItem::new(FileIdBothDirectoryInformationInner {
+                    file_index: 0,
+                    creation_time: FileTime::from(datetime!(2024-12-27 14:22:52.0116823)),
+                    last_access_time: FileTime::from(datetime!(2024-12-27 14:23:14.7795682)),
+                    last_write_time: FileTime::from(datetime!(2024-12-27 14:23:14.7795682)),
+                    change_time: FileTime::from(datetime!(2024-12-27 14:23:14.7795682)),
+                    end_of_file: 486,
+                    allocation_size: 488,
+                    file_attributes: FileAttributes::new().with_archive(true),
+                    ea_size: 0,
+                    short_name_length: 0,
+                    short_name: [0; 12],
+                    fild_id: 1125899906906299,
+                    file_name: "c.txt".into(),
+                },),
+                ChainedItem::new(FileIdBothDirectoryInformationInner {
+                    file_index: 0,
+                    creation_time: FileTime::from(datetime!(2024-12-27 14:22:52.167941),),
+                    last_access_time: FileTime::from(datetime!(2025-01-02 19:05:44.7804931),),
+                    last_write_time: FileTime::from(datetime!(2025-01-02 19:05:44.7804931),),
+                    change_time: FileTime::from(datetime!(2025-01-02 19:05:44.7804931),),
+                    end_of_file: 15910,
+                    allocation_size: 16384,
+                    file_attributes: FileAttributes::new().with_archive(true),
+                    ea_size: 0,
+                    short_name_length: 0,
+                    short_name: [0; 12],
+                    fild_id: 1125899906906300,
+                    file_name: "d.txt".into(),
+                })
+            ],
+            _res
+        );
     }
 }
diff --git a/smb/src/session/signer.rs b/smb/src/session/signer.rs
@@ -90,5 +90,34 @@ impl Clone for MessageSigner {
 
 #[cfg(test)]
 mod tests {
-    // TODO: Add tests.
+    use crate::{crypto::make_signing_algo, packets::smb2::SigningAlgorithmId};
+
+    use super::*;
+
+    const TEST_SIGNING_KEY: [u8; 16] = [
+        0xAC, 0x36, 0xE9, 0x54, 0x3C, 0xD8, 0x88, 0xF0, 0xA8, 0x41, 0x23, 0xE4, 0x6B, 0xB2, 0xA0,
+        0xD7,
+    ];
+
+    #[test]
+    fn test_calc_signature() {
+        // Some random session logoff request for testing.
+        let raw_data = vec![
+            0xfeu8, 0x53, 0x4d, 0x42, 0x40, 0x0, 0x1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x2, 0x0, 0x1, 0x0,
+            0x18, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x9, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
+            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x53, 0x20, 0xc, 0x21, 0x0, 0x0, 0x0, 0x0, 0x76,
+            0x23, 0x4b, 0x3c, 0x81, 0x2f, 0x51, 0xab, 0x8a, 0x5c, 0xf9, 0xfa, 0x43, 0xd4, 0xeb,
+            0x28, 0x4, 0x0, 0x0, 0x0,
+        ];
+        let mut header = Header::read_le(&mut Cursor::new(
+            &raw_data.as_slice()[..=Header::STRUCT_SIZE],
+        ))
+        .unwrap();
+
+        let mut signer = MessageSigner::new(
+            make_signing_algo(SigningAlgorithmId::AesGmac, &TEST_SIGNING_KEY).unwrap(),
+        );
+        let signature = signer.calculate_signature(&mut header, &raw_data).unwrap();
+        assert_eq!(signature, 0x28ebd443faf95c8aab512f813c4b2376);
+    }
 }
diff --git a/smb/src/tree.rs b/smb/src/tree.rs
@@ -138,7 +138,7 @@ impl TreeMessageHandler {
 
     #[maybe_async]
     async fn disconnect(&mut self) -> crate::Result<()> {
-        let info = self.connect_info.take();
+        let info = self.connect_info.get();
 
         if !info.is_some() {
             return Err(Error::InvalidState(
@@ -155,6 +155,7 @@ impl TreeMessageHandler {
             ))
             .await?;
 
+        self.connect_info.take();
         log::info!("Disconnected from tree {}", self.tree_name);
 
         Ok(())
