implement proper json deserialization

1 files changed, 77 insertions, 22 deletions

diff --git a/src/i3bar_protocol.rs b/src/i3bar_protocol.rs
index d838fb1..7da061c 100644
--- a/src/i3bar_protocol.rs
+++ b/src/i3bar_protocol.rs
@@ -3,6 +3,7 @@ use crate::pointer_btn::PointerBtn;
 use crate::text::Align;
 use serde::Deserialize;
 use serde::Serialize;
+use serde_json::{Deserializer, Result as SerdeResult};
 use std::io::{Error, ErrorKind, Result};
 
 #[derive(Deserialize, Default)]
@@ -81,11 +82,17 @@ pub enum Protocol {
     Unknown,
     PlainText(Option<String>),
     JsonNotStarted(JsonHeader),
-    Json(JsonHeader, Option<Vec<Block>>),
+    Json(JsonHeader, Option<Vec<Block>>, Option<String>),
 }
 
 impl Protocol {
     pub fn process_line(&mut self, line: String) -> Result<()> {
+        macro_rules! invalid {
+            ($fmt:expr) => {
+                return Err(Error::new(ErrorKind::InvalidData, format!($fmt)))
+            };
+        }
+
         match self {
             Self::Unknown => {
                 if let Ok(header) = serde_json::from_str::<JsonHeader>(&line) {
@@ -98,32 +105,38 @@ impl Protocol {
             }
             Self::PlainText(s) => *s = Some(line),
             Self::JsonNotStarted(header) => {
-                let line = line.trim();
+                let line = line.trim_start();
                 if !line.is_empty() {
-                    if !line.starts_with('[') && line.is_empty() {
-                        return Err(Error::new(ErrorKind::InvalidData, "Expected '['"));
+                    if !line.starts_with('[') {
+                        invalid!("Expected '['");
                     }
-                    *self = Self::Json(*header, None);
+                    *self = Self::Json(*header, None, None);
                     self.process_line(line[1..].to_owned())?;
                 }
             }
-            Self::Json(_, blocks) => {
-                let line = line.trim();
-                if !line.is_empty() {
-                    if !line.ends_with(',') {
-                        return Err(Error::new(ErrorKind::InvalidData, "Expected ','"));
-                    }
-                    *blocks = match serde_json::from_str::<Vec<Block>>(&line[..(line.len() - 1)]) {
-                        Ok(b) => Some(b),
-                        Err(e) => {
-                            return Err(Error::new(
-                                ErrorKind::InvalidData,
-                                format!("Invalid JSON: {e}"),
-                            ))
+            Self::Json(_, blocks, start_orig) => match start_orig {
+                Some(start) => {
+                    start.push_str(&line);
+                    match de_last::<Vec<Block>>(start) {
+                        Err(e) => invalid!("Invalid JSON: {e}"),
+                        Ok((rem, new_blocks)) => {
+                            *start_orig = (!rem.is_empty()).then(|| rem.to_owned());
+                            if let Some(new_blocks) = new_blocks {
+                                *blocks = Some(new_blocks);
+                            }
                         }
-                    };
+                    }
                 }
-            }
+                None => match de_last::<Vec<Block>>(&line) {
+                    Err(e) => invalid!("Invalid JSON: {e}"),
+                    Ok((rem, new_blocks)) => {
+                        *start_orig = (!rem.is_empty()).then(|| rem.to_owned());
+                        if let Some(new_blocks) = new_blocks {
+                            *blocks = Some(new_blocks);
+                        }
+                    }
+                },
+            },
         }
         Ok(())
     }
@@ -136,15 +149,57 @@ impl Protocol {
                 ..Default::default()
             }]),
             Self::JsonNotStarted(_) => None,
-            Self::Json(_, blocks) => blocks.take(),
+            Self::Json(_, blocks, _) => blocks.take(),
         }
     }
 
     pub fn supports_clicks(&self) -> bool {
         match self {
             Self::JsonNotStarted(h) => h.click_events,
-            Self::Json(h, _) => h.click_events,
+            Self::Json(h, _, _) => h.click_events,
             _ => false,
         }
     }
 }
+
+/// Deserialize the last complete object. Returns (`remaining`, `object`). See tests for examples.
+fn de_last<'a, T: Deserialize<'a>>(mut s: &'a str) -> SerdeResult<(&'a str, Option<T>)> {
+    let mut last = None;
+    s = s.trim_start_matches(|x: char| x.is_ascii_whitespace() || x == ',');
+    loop {
+        let (rem, obj) = de_once(s)?;
+        last = match obj {
+            Some(obj) => Some(obj),
+            None => return Ok((rem, last)),
+        };
+        s = rem.trim_start_matches(|x: char| x.is_ascii_whitespace() || x == ',');
+    }
+}
+
+/// Deserialize the first complete object. Returns (`remaining`, `object`).
+fn de_once<'a, T: Deserialize<'a>>(s: &'a str) -> SerdeResult<(&'a str, Option<T>)> {
+    let mut de = Deserializer::from_str(s).into_iter();
+    match de.next() {
+        Some(Ok(obj)) => Ok((&s[de.byte_offset()..], Some(obj))),
+        Some(Err(e)) if e.is_eof() => Ok((&s[de.byte_offset()..], None)),
+        Some(Err(e)) => Err(e),
+        None => Ok((&s[de.byte_offset()..], None)),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn streaming_json() {
+        let s = ",[2], [3], [4, 3],[32][3] ";
+        assert_eq!(de_last::<Vec<u8>>(s).unwrap(), ("", Some(vec![3])));
+
+        let s = ",[2], [3], [4, 3],[32][3] [2, 4";
+        assert_eq!(de_last::<Vec<u8>>(s).unwrap(), ("[2, 4", Some(vec![3])));
+
+        let s = ",[2], [3], [4, 3],[32] invalid";
+        assert!(de_last::<Vec<u8>>(s).is_err());
+    }
+}