From a8de682c1543bb121cd35bfdfcfd1249cabdcb0c Mon Sep 17 00:00:00 2001
From: Narthana Epa <narthana.epa@gmail.com>
Date: Sun, 1 Sep 2024 09:43:40 +0530
Subject: [PATCH] Escape raw data

---
 src/lib.rs      |  4 +--
 src/response.rs | 82 +++++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 82 insertions(+), 4 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index c0ed444..ba6008d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -298,7 +298,7 @@ mod test {
         let mut output = std::io::Cursor::new(vec![]);
         let mut listener = Listener::new(Config {
             timeout: None,
-            command: vec!["echo", "-n", "1234"]
+            command: vec!["echo", "1234"]
                 .into_iter()
                 .map(std::string::ToString::to_string)
                 .collect(),
@@ -343,7 +343,7 @@ mod test {
                     OK
                     OK
                     OK
-                    D 1234
+                    D 1234%0A
                     OK
                     OK closing connection
                 "},
diff --git a/src/response.rs b/src/response.rs
index e845233..5c96307 100644
--- a/src/response.rs
+++ b/src/response.rs
@@ -1,4 +1,7 @@
-use std::fmt::{self, Display, Formatter};
+use std::{
+    borrow::Cow,
+    fmt::{self, Display, Formatter},
+};
 
 #[derive(Debug, PartialEq, Eq)]
 pub enum Response {
@@ -20,10 +23,85 @@ impl Display for Response {
                 s.as_ref().map(|s| format!(" {s}")).unwrap_or_default(),
             ),
             Err(code, msg) => write!(f, "ERR {code} {msg}"),
-            D(s) => write!(f, "D {s}"),
+            D(s) => write!(f, "D {}", escape(s)),
             Comment(s) => write!(f, "# {s}"),
             S(k, v) => write!(f, "S {k} {v}"),
             Inquire(k, v) => write!(f, "INQUIRE {k} {v}"),
         }
     }
 }
+
+/// Encode a string to be used in a response. It will percent escape `%`, `\n`, and `\r`.
+fn escape(s: &str) -> Cow<'_, str> {
+    // TODO: Split into lines of length at most 1000 bytes.
+    let mut s = s;
+    let mut escaped = String::with_capacity(s.len());
+
+    loop {
+        let unescaped_len = s
+            .chars()
+            .take_while(|c| !matches!(c, '%' | '\n' | '\r'))
+            .count();
+
+        let (unescaped, rest) = if unescaped_len >= s.len() {
+            if escaped.is_empty() {
+                return Cow::from(s);
+            }
+            (s, "")
+        } else {
+            s.split_at(unescaped_len)
+        };
+
+        if !unescaped.is_empty() {
+            escaped.push_str(unescaped);
+        }
+        if rest.is_empty() {
+            break;
+        }
+        let (first, rest) = rest.split_at(1);
+        match first {
+            "%" => escaped.push_str("%25"),
+            "\n" => escaped.push_str("%0A"),
+            "\r" => escaped.push_str("%0D"),
+            _ => unreachable!(),
+        }
+        s = rest;
+    }
+
+    Cow::from(escaped)
+}
+
+#[cfg(test)]
+mod test {
+    use std::borrow::Cow;
+
+    #[test]
+    fn escape() {
+        [
+            ("", ""),
+            ("a", "a"),
+            ("a\n", "a%0A"),
+            ("a\r", "a%0D"),
+            ("a%", "a%25"),
+            ("a%b", "a%25b"),
+            ("a%b\n", "a%25b%0A"),
+            ("a%b\r", "a%25b%0D"),
+            ("a\nb", "a%0Ab"),
+            ("a\rb", "a%0Db"),
+            ("a\nb\n", "a%0Ab%0A"),
+            ("a\rb\r", "a%0Db%0D"),
+            ("a\nb\r", "a%0Ab%0D"),
+            ("a\rb\n", "a%0Db%0A"),
+            ("a\nb\r\n", "a%0Ab%0D%0A"),
+            ("a\nb\r\nc", "a%0Ab%0D%0Ac"),
+            ("a\nb\r\nc\n", "a%0Ab%0D%0Ac%0A"),
+            ("a\nb\r\nc\nd", "a%0Ab%0D%0Ac%0Ad"),
+            ("a\nb\r\nc\nd\n", "a%0Ab%0D%0Ac%0Ad%0A"),
+        ]
+        .into_iter()
+        .map(|(input, expected)| (input, Cow::from(expected)))
+        .for_each(|(input, expected)| {
+            assert_eq!(super::escape(input), *expected);
+        });
+    }
+}