Merge pull request #2 from moriyoshi/correctly-parse-quoted-string

fix: parse quoted strings correctly

src/lib.rs

@@ -50,7 +50,10 @@ pub fn real_ip(
 pub fn get_forwarded_for() -> impl Filter<Extract = (Vec<IpAddr>,), Error = Infallible> + Clone {
     warp::header("x-forwarded-for")
         .map(|list: CommaSeparated<IpAddr>| list.into_inner())
-        .or(warp::header("x-real-ip").map(|ip| vec![ip]))
+        .or(warp::header("x-real-ip").map(
+            |ip: String| IpAddr::from_str(maybe_bracketed(&maybe_quoted(ip)))
+                .map_or_else(|_| Vec::<IpAddr>::new(), |x| vec![x])
+        ))
         .unify()
         .or(warp::header("forwarded").map(|header: String| {
             parse(&header)
@@ -72,6 +75,143 @@ pub fn get_forwarded_for() -> impl Filter<Extract = (Vec<IpAddr>,), Error = Infa
         .unify()
 }
 
+enum CommaSeparatedIteratorState {
+    Default,
+    Quoted,
+    QuotedPair,
+    Token,
+    PostambleForQuoted,
+}
+
+struct CommaSeparatedIterator<'a> {
+    /// target
+    target: &'a str,
+    /// iterator
+    char_indices: std::str::CharIndices<'a>,
+    /// current scanner state
+    state: CommaSeparatedIteratorState,
+    /// start position of the last token found
+    s: usize,
+}
+
+impl<'a> CommaSeparatedIterator<'a> {
+    pub fn new(target: &'a str) -> Self {
+        Self {
+            target: target,
+            char_indices: target.char_indices(),
+            state: CommaSeparatedIteratorState::Default,
+            s: 0,
+        }
+    }
+}
+
+impl<'a> Iterator for CommaSeparatedIterator<'a> {
+    type Item = &'a str;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            match self.char_indices.next() {
+                Some((i, c)) => match match self.state {
+                    CommaSeparatedIteratorState::Default => match c {
+                        '"' => {
+                            self.s = i;
+                            (None, CommaSeparatedIteratorState::Quoted)
+                        }
+                        ' ' | '\t' => (None, CommaSeparatedIteratorState::Default),
+                        ',' => (
+                            Some(Some(&self.target[i..i])),
+                            CommaSeparatedIteratorState::Default,
+                        ),
+                        _ => {
+                            self.s = i;
+                            (None, CommaSeparatedIteratorState::Token)
+                        }
+                    },
+                    CommaSeparatedIteratorState::Quoted => match c {
+                        '"' => (
+                            Some(Some(&self.target[self.s..i + 1])),
+                            CommaSeparatedIteratorState::PostambleForQuoted,
+                        ),
+                        '\\' => (None, CommaSeparatedIteratorState::QuotedPair),
+                        _ => (None, CommaSeparatedIteratorState::Quoted),
+                    },
+                    CommaSeparatedIteratorState::QuotedPair => {
+                        (None, CommaSeparatedIteratorState::Quoted)
+                    }
+                    CommaSeparatedIteratorState::Token => match c {
+                        ',' => (
+                            Some(Some(&self.target[self.s..i])),
+                            CommaSeparatedIteratorState::Default,
+                        ),
+                        _ => (None, CommaSeparatedIteratorState::Token),
+                    },
+                    CommaSeparatedIteratorState::PostambleForQuoted => match c {
+                        ',' => (None, CommaSeparatedIteratorState::Default),
+                        _ => (None, CommaSeparatedIteratorState::PostambleForQuoted),
+                    },
+                } {
+                    (Some(next), next_state) => {
+                        self.state = next_state;
+                        return next;
+                    }
+                    (None, next_state) => {
+                        self.state = next_state;
+                    }
+                },
+                None => break,
+            }
+        }
+        return match self.state {
+            CommaSeparatedIteratorState::Default
+            | CommaSeparatedIteratorState::PostambleForQuoted => None,
+            CommaSeparatedIteratorState::Quoted | CommaSeparatedIteratorState::QuotedPair => {
+                self.state = CommaSeparatedIteratorState::Default;
+                Some(&self.target[self.s..])
+            }
+            CommaSeparatedIteratorState::Token => {
+                self.state = CommaSeparatedIteratorState::Default;
+                Some(&self.target[self.s..])
+            }
+        };
+    }
+}
+
+pub fn maybe_quoted<T: AsRef<str>>(x: T) -> String {
+    let x = x.as_ref();
+    let mut i = x.chars();
+    if i.next() == Some('"') {
+        let mut s = String::with_capacity(x.len());
+        let mut state = 0;
+        for c in i {
+            state = match state {
+                0 => match c {
+                    '"' => break,
+                    '\\' => 1,
+                    _ => {
+                        s.push(c);
+                        0
+                    }
+                },
+                _ => {
+                    s.push(c);
+                    0
+                }
+            };
+        }
+        s
+    } else {
+        x.to_string()
+    }
+}
+
+pub fn maybe_bracketed<'a>(x: &'a str) -> &'a str {
+    if x.as_bytes()[0] == ('[' as u8) && x.as_bytes()[x.len() - 1] == (']' as u8) {
+        &x[1..x.len() - 1]
+    } else {
+        x
+    }
+}
+
 /// Newtype so we can implement FromStr
 struct CommaSeparated<T>(Vec<T>);
 
@@ -85,11 +225,36 @@ impl<T: FromStr> FromStr for CommaSeparated<T> {
     type Err = T::Err;
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
-        let vec = s
+        let vec = CommaSeparatedIterator::new(s)
-            .split(',')
+            .map(|x| T::from_str(maybe_bracketed(&maybe_quoted(x.trim()))))
-            .map(str::trim)
-            .map(T::from_str)
             .collect::<Result<Vec<_>, _>>()?;
         Ok(CommaSeparated(vec))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::{CommaSeparatedIterator, maybe_quoted, maybe_bracketed};
+
+    #[test]
+    fn test_comma_separated_iterator() {
+        assert_eq!(vec!["abc", "def", "ghi", "jkl ", "mno", "pqr"], CommaSeparatedIterator::new("abc,def, ghi,\tjkl , mno,\tpqr").collect::<Vec<&str>>());
+        assert_eq!(vec!["abc", "\"def\"", "\"ghi\"", "\"jkl\"", "\"mno\"", "pqr"], CommaSeparatedIterator::new("abc,\"def\", \"ghi\",\t\"jkl\" , \"mno\",\tpqr").collect::<Vec<&str>>());
+    }
+
+    #[test]
+    fn test_maybe_quoted() {
+        assert_eq!("abc", maybe_quoted("abc"));
+        assert_eq!("abc", maybe_quoted("\"abc\""));
+        assert_eq!("a\"bc", maybe_quoted("\"a\\\"bc\""));
+    }
+
+    #[test]
+    fn test_maybe_bracketed() {
+        assert_eq!("abc", maybe_bracketed("abc"));
+        assert_eq!("abc", maybe_bracketed("[abc]"));
+        assert_eq!("[abc", maybe_bracketed("[abc"));
+        assert_eq!("abc]", maybe_bracketed("abc]"));
+    }
+
+}