From: Jason Ish <jason.ish@oisf.net>
Date: Thu, 21 May 2020 19:06:50 +0000 (-0600)
Subject: jsonbuilder: add reset marks
X-Git-Tag: suricata-6.0.0-beta1~372
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=60bfbd43fdead02ff156f5335b1961f6d5bc9689;p=thirdparty%2Fsuricata.git

jsonbuilder: add reset marks

Add methods to get the state of a JsonBuilder (called a mark),
then allow restoring to the mark.
---

diff --git a/rust/cbindgen.toml b/rust/cbindgen.toml
index dfcf8d03d5..8c41789e58 100644
--- a/rust/cbindgen.toml
+++ b/rust/cbindgen.toml
@@ -74,6 +74,7 @@ documentation_style = "doxy"
 include = [
     "AppLayerGetTxIterTuple",
     "SIPState",
+    "CMark",
 ]
 
 # A list of items to not include in the generated bindings
diff --git a/rust/src/jsonbuilder.rs b/rust/src/jsonbuilder.rs
index 5e72ca2ba7..003c6a00dd 100644
--- a/rust/src/jsonbuilder.rs
+++ b/rust/src/jsonbuilder.rs
@@ -54,14 +54,49 @@ enum Type {
 }
 
 #[derive(Debug, Clone, Copy, PartialEq)]
+#[repr(C)]
 enum State {
-    None,
+    None = 0,
     ObjectFirst,
     ObjectNth,
     ArrayFirst,
     ArrayNth,
 }
 
+impl State {
+    fn from_u64(v: u64) -> Result<State, JsonError> {
+        let s = match v {
+            0 => State::None,
+            1 => State::ObjectFirst,
+            2 => State::ObjectNth,
+            3 => State::ArrayFirst,
+            4 => State::ArrayNth,
+            _ => {
+                return Err(JsonError::InvalidState);
+            }
+        };
+        Ok(s)
+    }
+}
+
+#[derive(Debug, Clone)]
+struct Mark {
+    position: usize,
+    state_index: usize,
+    state: State,
+}
+
+/// A "mark" or saved state for a JsonBuilder object.
+///
+/// The name is full, and the types are u64 as this object is used
+/// directly in C as well.
+#[repr(C)]
+pub struct JsonBuilderMark {
+    position: u64,
+    state_index: u64,
+    state: u64,
+}
+
 #[derive(Debug, Clone)]
 pub struct JsonBuilder {
     buf: String,
@@ -158,6 +193,24 @@ impl JsonBuilder {
         self.state[n] = state;
     }
 
+    pub fn get_mark(&self) -> JsonBuilderMark {
+        JsonBuilderMark {
+            position: self.buf.len() as u64,
+            state: self.current_state() as u64,
+            state_index: self.state.len() as u64,
+        }
+    }
+
+    pub fn restore_mark(&mut self, mark: &JsonBuilderMark) -> Result<(), JsonError> {
+        let state = State::from_u64(mark.state)?;
+        if mark.position < (self.buf.len() as u64) && mark.state_index < (self.state.len() as u64) {
+            self.buf.truncate(mark.position as usize);
+            self.state.truncate(mark.state_index as usize);
+            self.state[(mark.state_index as usize) - 1] = state;
+        }
+        Ok(())
+    }
+
     /// Open an object under the given key.
     ///
     /// For example:
@@ -307,9 +360,7 @@ impl JsonBuilder {
     }
 
     pub fn set_jsont(
-        &mut self,
-        key: &str,
-        jsont: &mut json::JsonT,
+        &mut self, key: &str, jsont: &mut json::JsonT,
     ) -> Result<&mut Self, JsonError> {
         match self.current_state() {
             State::ObjectNth => self.buf.push(','),
@@ -541,9 +592,7 @@ pub unsafe extern "C" fn jb_open_array(js: &mut JsonBuilder, key: *const c_char)
 
 #[no_mangle]
 pub unsafe extern "C" fn jb_set_string(
-    js: &mut JsonBuilder,
-    key: *const c_char,
-    val: *const c_char,
+    js: &mut JsonBuilder, key: *const c_char, val: *const c_char,
 ) -> bool {
     if val == std::ptr::null() {
         return false;
@@ -558,9 +607,7 @@ pub unsafe extern "C" fn jb_set_string(
 
 #[no_mangle]
 pub unsafe extern "C" fn jb_set_jsont(
-    jb: &mut JsonBuilder,
-    key: *const c_char,
-    jsont: &mut json::JsonT,
+    jb: &mut JsonBuilder, key: *const c_char, jsont: &mut json::JsonT,
 ) -> bool {
     if let Ok(key) = CStr::from_ptr(key).to_str() {
         return jb.set_jsont(key, jsont).is_ok();
@@ -575,9 +622,7 @@ pub unsafe extern "C" fn jb_append_object(jb: &mut JsonBuilder, obj: &JsonBuilde
 
 #[no_mangle]
 pub unsafe extern "C" fn jb_set_object(
-    js: &mut JsonBuilder,
-    key: *const c_char,
-    val: &mut JsonBuilder,
+    js: &mut JsonBuilder, key: *const c_char, val: &mut JsonBuilder,
 ) -> bool {
     if let Ok(key) = CStr::from_ptr(key).to_str() {
         return js.set_object(key, val).is_ok();
@@ -632,6 +677,19 @@ pub unsafe extern "C" fn jb_ptr(js: &mut JsonBuilder) -> *const u8 {
     js.buf.as_ptr()
 }
 
+#[no_mangle]
+pub unsafe extern "C" fn jb_get_mark(js: &mut JsonBuilder, mark: &mut JsonBuilderMark) {
+    let m = js.get_mark();
+    mark.position = m.position;
+    mark.state_index = m.state_index;
+    mark.state = m.state;
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn jb_restore_mark(js: &mut JsonBuilder, mark: &mut JsonBuilderMark) -> bool {
+    js.restore_mark(mark).is_ok()
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
@@ -870,6 +928,29 @@ mod test {
         jb.append_string_from_bytes(&[0xf0, 0xf1, 0xf2]).unwrap();
         assert_eq!(jb.buf, r#"["\\xf0\\xf1\\xf2""#);
     }
+
+    #[test]
+    fn test_marks() {
+        let mut jb = JsonBuilder::new_object();
+        jb.set_string("foo", "bar").unwrap();
+        assert_eq!(jb.buf, r#"{"foo":"bar""#);
+        assert_eq!(jb.current_state(), State::ObjectNth);
+        assert_eq!(jb.state.len(), 2);
+        let mark = jb.get_mark();
+
+        // Mutate such that states are transitioned.
+        jb.open_array("bar").unwrap();
+        jb.start_object().unwrap();
+        assert_eq!(jb.buf, r#"{"foo":"bar","bar":[{"#);
+        assert_eq!(jb.current_state(), State::ObjectFirst);
+        assert_eq!(jb.state.len(), 4);
+
+        // Restore to mark.
+        jb.restore_mark(&mark).unwrap();
+        assert_eq!(jb.buf, r#"{"foo":"bar""#);
+        assert_eq!(jb.current_state(), State::ObjectNth);
+        assert_eq!(jb.state.len(), 2);
+    }
 }
 
 // Escape table as seen in serde-json (MIT/Apache license)