1 files changed, 280 insertions, 0 deletions
diff --git a/rust/src/tftp/tftp.rs b/rust/src/tftp/tftp.rs
new file mode 100644
index 0000000..1b093fe
--- /dev/null
+++ b/rust/src/tftp/tftp.rs
@@ -0,0 +1,280 @@
+/* Copyright (C) 2017-2021 Open Information Security Foundation
+ *
+ * You can copy, redistribute or modify this Program under the terms of
+ * the GNU General Public License version 2 as published by the Free
+ * Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * version 2 along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
+ * 02110-1301, USA.
+ */
+
+// written by Clément Galland <clement.galland@epita.fr>
+
+use std::str;
+use std;
+use nom7::IResult;
+use nom7::combinator::map_res;
+use nom7::bytes::streaming::{tag, take_while};
+use nom7::number::streaming::be_u8;
+
+use crate::applayer::{AppLayerTxData,AppLayerStateData};
+
+const READREQUEST:  u8 = 1;
+const WRITEREQUEST: u8 = 2;
+const DATA:         u8 = 3;
+const ACK:          u8 = 4;
+const ERROR:        u8 = 5;
+
+#[derive(Debug, PartialEq, Eq)]
+pub struct TFTPTransaction {
+    pub opcode : u8,
+    pub filename : String,
+    pub mode : String,
+    id: u64,
+    tx_data: AppLayerTxData,
+}
+
+pub struct TFTPState {
+    state_data: AppLayerStateData,
+    pub transactions : Vec<TFTPTransaction>,
+    /// tx counter for assigning incrementing id's to tx's
+    tx_id: u64,
+}
+
+impl TFTPState {
+    fn get_tx_by_id(&mut self, tx_id: u64) -> Option<&TFTPTransaction> {
+        self.transactions.iter().find(|&tx| tx.id == tx_id + 1)
+    }
+
+    fn free_tx(&mut self, tx_id: u64) {
+        let tx = self.transactions.iter().position(|tx| tx.id == tx_id + 1);
+        debug_assert!(tx.is_some());
+        if let Some(idx) = tx {
+            let _ = self.transactions.remove(idx);
+        }
+    }
+}
+
+impl TFTPTransaction {
+    pub fn new(opcode : u8, filename : String, mode : String) -> TFTPTransaction {
+        TFTPTransaction {
+            opcode,
+            filename,
+            mode : mode.to_lowercase(),
+            id : 0,
+            tx_data: AppLayerTxData::new(),
+        }
+    }
+    pub fn is_mode_ok(&self) -> bool {
+        match self.mode.as_str() {
+            "netascii" | "mail" | "octet" => true,
+            _ => false
+        }
+    }
+    pub fn is_opcode_ok(&self) -> bool {
+        match self.opcode {
+            READREQUEST | WRITEREQUEST | ACK | DATA | ERROR => true,
+            _ => false
+        }
+    }
+}
+
+#[no_mangle]
+pub extern "C" fn rs_tftp_state_alloc() -> *mut std::os::raw::c_void {
+    let state = TFTPState { state_data: AppLayerStateData::new(), transactions : Vec::new(), tx_id: 0, };
+    let boxed = Box::new(state);
+    return Box::into_raw(boxed) as *mut _;
+}
+
+#[no_mangle]
+pub extern "C" fn rs_tftp_state_free(state: *mut std::os::raw::c_void) {
+    std::mem::drop(unsafe { Box::from_raw(state as *mut TFTPState) });
+}
+
+#[no_mangle]
+pub extern "C" fn rs_tftp_state_tx_free(state: &mut TFTPState,
+                                        tx_id: u64) {
+    state.free_tx(tx_id);
+}
+
+#[no_mangle]
+pub extern "C" fn rs_tftp_get_tx(state: &mut TFTPState,
+                                    tx_id: u64) -> *mut std::os::raw::c_void {
+    match state.get_tx_by_id(tx_id) {
+        Some(tx) => tx as *const _ as *mut _,
+        None     => std::ptr::null_mut(),
+    }
+}
+
+#[no_mangle]
+pub extern "C" fn rs_tftp_get_tx_cnt(state: &mut TFTPState) -> u64 {
+    return state.tx_id;
+}
+
+fn getstr(i: &[u8]) -> IResult<&[u8], &str> {
+    map_res(
+        take_while(|c| c != 0),
+        str::from_utf8
+    )(i)
+}
+
+fn tftp_request(slice: &[u8]) -> IResult<&[u8], TFTPTransaction> {
+    let (i, _) = tag([0])(slice)?;
+    let (i, opcode) = be_u8(i)?;
+    let (i, filename) = getstr(i)?;
+    let (i, _) = tag([0])(i)?;
+    let (i, mode) = getstr(i)?;
+    Ok((i,
+        TFTPTransaction::new(opcode, String::from(filename), String::from(mode))
+       )
+      )
+}
+
+fn parse_tftp_request(input: &[u8]) -> Option<TFTPTransaction> {
+    match tftp_request(input) {
+        Ok((_, tx)) => {
+            if !tx.is_mode_ok() {
+                return None;
+            }
+            if !tx.is_opcode_ok() {
+                return None;
+            }
+            return Some(tx);
+        }
+        Err(_) => {
+            return None;
+        }
+    }
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn rs_tftp_request(state: &mut TFTPState,
+                                  input: *const u8,
+                                  len: u32) -> i64 {
+    let buf = std::slice::from_raw_parts(input, len as usize);
+    match parse_tftp_request(buf) {
+        Some(mut tx) => {
+            state.tx_id += 1;
+            tx.id = state.tx_id;
+            state.transactions.push(tx);
+            0
+        },
+        None => {
+           -1
+        }
+    }
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn rs_tftp_get_tx_data(
+    tx: *mut std::os::raw::c_void)
+    -> *mut AppLayerTxData
+{
+    let tx = cast_pointer!(tx, TFTPTransaction);
+    return &mut tx.tx_data;
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn rs_tftp_get_state_data(
+    state: *mut std::os::raw::c_void)
+    -> *mut AppLayerStateData
+{
+    let state = cast_pointer!(state, TFTPState);
+    return &mut state.state_data;
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    static READ_REQUEST: [u8; 20] = [
+            0x00, 0x01, 0x72, 0x66, 0x63, 0x31, 0x33, 0x35, 0x30, 0x2e, 0x74, 0x78, 0x74, 0x00, 0x6f, 0x63, 0x74, 0x65, 0x74, 0x00,
+    ];
+    /* filename not terminated */
+    static READ_REQUEST_INVALID_1: [u8; 20] = [
+            0x00, 0x01, 0x72, 0x66, 0x63, 0x31, 0x33, 0x35, 0x30, 0x2e, 0x74, 0x78, 0x74, 0x6e, 0x6f, 0x63, 0x74, 0x65, 0x74, 0x00,
+    ];
+    /* garbage */
+    static READ_REQUEST_INVALID_2: [u8; 3] = [
+            0xff, 0xff, 0xff,
+    ];
+    static WRITE_REQUEST: [u8; 20] = [
+            0x00, 0x02, 0x72, 0x66, 0x63, 0x31, 0x33, 0x35, 0x30, 0x2e, 0x74, 0x78, 0x74, 0x00, 0x6f, 0x63, 0x74, 0x65, 0x74, 0x00,
+    ];
+    /* filename not terminated */
+    static INVALID_OPCODE: [u8; 20] = [
+            0x00, 0x06, 0x72, 0x66, 0x63, 0x31, 0x33, 0x35, 0x30, 0x2e, 0x74, 0x78, 0x74, 0x6e, 0x6f, 0x63, 0x74, 0x65, 0x74, 0x00,
+    ];
+    static INVALID_MODE: [u8; 20] = [
+            0x00, 0x01, 0x72, 0x66, 0x63, 0x31, 0x33, 0x35, 0x30, 0x2e, 0x74, 0x78, 0x74, 0x00, 0x63, 0x63, 0x63, 0x63, 0x63, 0x00,
+    ];
+
+    #[test]
+    pub fn test_parse_tftp_read_request_1() {
+        let tx = TFTPTransaction {
+            opcode: READREQUEST,
+            filename: String::from("rfc1350.txt"),
+            mode: String::from("octet"),
+            id: 0,
+            tx_data: AppLayerTxData::new(),
+        };
+
+        match parse_tftp_request(&READ_REQUEST[..]) {
+            Some(txp) => {
+                assert_eq!(tx, txp);
+            }
+            None => {
+                assert!(true);
+            }
+        }
+    }
+
+    #[test]
+    pub fn test_parse_tftp_write_request_1() {
+        let tx = TFTPTransaction {
+            opcode: WRITEREQUEST,
+            filename: String::from("rfc1350.txt"),
+            mode: String::from("octet"),
+            id: 0,
+            tx_data: AppLayerTxData::new(),
+        };
+
+        match parse_tftp_request(&WRITE_REQUEST[..]) {
+            Some(txp) => {
+                assert_eq!(tx, txp);
+            }
+            None => {
+                assert!(true, "fadfasd");
+            }
+        }
+    }
+
+    // Invalid request: filename not terminated
+    #[test]
+    pub fn test_parse_tftp_read_request_2() {
+        assert_eq!(None, parse_tftp_request(&READ_REQUEST_INVALID_1[..]));
+    }
+
+    // Invalid request: garbage input
+    #[test]
+    pub fn test_parse_tftp_read_request_3() {
+        assert_eq!(None, parse_tftp_request(&READ_REQUEST_INVALID_2[..]));
+    }
+
+    #[test]
+    pub fn test_parse_tftp_invalid_opcode_1() {
+        assert_eq!(None, parse_tftp_request(&INVALID_OPCODE[..]));
+    }
+
+    #[test]
+    pub fn test_parse_tftp_invalid_mode() {
+
+        assert_eq!(None, parse_tftp_request(&INVALID_MODE[..]));
+    }
+}