Last active
January 30, 2023 04:18
-
-
Save lxdlam/e3102327d3be4689660be48c3f68cea5 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
// You need at least `nom = "7"` to work. | |
mod protocol { | |
use std::{fmt, str::FromStr}; | |
use nom::Finish; | |
use self::parser::parse; | |
const DELIMITER: &str = "\r\n"; | |
#[derive(Debug, PartialEq, Eq, Clone)] | |
pub enum Object { | |
SimpleString(String), | |
Error(String), | |
Integer(i64), | |
BulkString(Option<String>), | |
Array(Option<Vec<Object>>), | |
} | |
impl fmt::Display for Object { | |
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { | |
match self { | |
Object::SimpleString(s) => write!(f, "+{}{}", s, DELIMITER), | |
Object::Error(s) => write!(f, "-{}{}", s, DELIMITER), | |
Object::Integer(i) => write!(f, ":{}{}", i, DELIMITER), | |
Object::BulkString(s) => { | |
if let Some(s) = s { | |
write!(f, "${}{}{}{}", s.len(), DELIMITER, s, DELIMITER) | |
} else { | |
write!(f, "$-1\r\n") | |
} | |
} | |
Object::Array(v) => { | |
if let Some(v) = v { | |
write!( | |
f, | |
"*{}{}{}", | |
v.len(), | |
DELIMITER, | |
v.iter() | |
.map(ToString::to_string) | |
.collect::<Vec<String>>() | |
.join("") | |
) | |
} else { | |
write!(f, "*-1\r\n") | |
} | |
} | |
} | |
} | |
} | |
#[cfg(test)] | |
mod encode_tests { | |
use super::*; | |
macro_rules! encode_test { | |
{ | |
$(($name:ident,$expected:literal,$actual:expr)),* | |
} => { | |
$(#[test] | |
fn $name() { | |
let actual = $actual.to_string(); | |
assert_eq!($expected, actual); | |
})* | |
}; | |
} | |
encode_test! { | |
(test_simple_string, "+OK\r\n", Object::SimpleString("OK".into())), | |
(test_error, "-Error Message\r\n", Object::Error("Error Message".into())), | |
(test_integer, ":1000\r\n", Object::Integer(1000)), | |
(test_zero, ":0\r\n", Object::Integer(0)), | |
(test_negative, ":-10000\r\n", Object::Integer(-10000)), | |
(test_null_bulk_string, "$-1\r\n", Object::BulkString(None)), | |
(test_empty_bulk_string, "$0\r\n\r\n", Object::BulkString(Some("".into()))), | |
(test_bulk_string, "$14\r\nHello \rWorld!\n\r\n", Object::BulkString(Some("Hello \rWorld!\n".into()))), | |
( | |
test_complex_bulk_string, | |
"$206\r\n$198\r\nMcpJxQSaoHhgLbAsdPTRjGFCZqrwlEivnIuKkfOVYzWDNXmyUetB\r\n\r\n\r\x0c \n\t\x0b\r\n\r\n^|,!%#`:_<*-=?;$[\\&.>\"(~+]'/}{)@\r\n\r\nETbgCoRcXOWezIHKxmQqVvhADyGZJplufNFjakdYPrwMLsBitUSn\r\n\r\n_)~&!\\.[<=*]^>+;$%(/'@,}\"{:-|#`?\r\n\r\n\r\n \t\r\n\r\n", | |
Object::BulkString(Some("$198\r\nMcpJxQSaoHhgLbAsdPTRjGFCZqrwlEivnIuKkfOVYzWDNXmyUetB\r\n\r\n\r\x0c \n\t\x0b\r\n\r\n^|,!%#`:_<*-=?;$[\\&.>\"(~+]'/}{)@\r\n\r\nETbgCoRcXOWezIHKxmQqVvhADyGZJplufNFjakdYPrwMLsBitUSn\r\n\r\n_)~&!\\.[<=*]^>+;$%(/'@,}\"{:-|#`?\r\n\r\n\r\n \t\r\n".into())) | |
), | |
(test_utf8_bulk_string, "$19\r\n你好,\n世界!\r\n", Object::BulkString(Some("你好,\n世界!".into()))), | |
(test_null_array, "*-1\r\n", Object::Array(None)), | |
(test_empty_array, "*0\r\n", Object::Array(Some(vec![]))), | |
( | |
test_array, | |
"*4\r\n+OK\r\n-Error Message\r\n:1000\r\n$14\r\nHello \rWorld!\n\r\n", | |
Object::Array(Some(vec![ | |
Object::SimpleString("OK".into()), | |
Object::Error("Error Message".into()), | |
Object::Integer(1000), | |
Object::BulkString(Some("Hello \rWorld!\n".into())), | |
])) | |
), | |
( | |
test_null_in_array, | |
"*3\r\n$3\r\nfoo\r\n$-1\r\n$3\r\nbar\r\n", | |
Object::Array(Some(vec![ | |
Object::BulkString(Some("foo".into())), | |
Object::BulkString(None), | |
Object::BulkString(Some("bar".into())), | |
])) | |
), | |
( | |
test_nested_array, | |
"*1\r\n*1\r\n*1\r\n*2\r\n:-1\r\n+OK\r\n", | |
Object::Array(Some(vec![Object::Array(Some(vec![Object::Array(Some( | |
vec![Object::Array(Some(vec![ | |
Object::Integer(-1), | |
Object::SimpleString("OK".into()), | |
]))], | |
))]))])) | |
), | |
( | |
test_complex_array, | |
"*13\r\n+OK\r\n-Error Message\r\n:1000\r\n:0\r\n:-10000\r\n$0\r\n\r\n$14\r\nHello \rWorld!\n\r\n$206\r\n$198\r\nMcpJxQSaoHhgLbAsdPTRjGFCZqrwlEivnIuKkfOVYzWDNXmyUetB\r\n\r\n\r\x0c \n\t\x0b\r\n\r\n^|,!%#`:_<*-=?;$[\\&.>\"(~+]'/}{)@\r\n\r\nETbgCoRcXOWezIHKxmQqVvhADyGZJplufNFjakdYPrwMLsBitUSn\r\n\r\n_)~&!\\.[<=*]^>+;$%(/'@,}\"{:-|#`?\r\n\r\n\r\n \t\r\n\r\n*0\r\n*4\r\n+OK\r\n-Error Message\r\n:1000\r\n$14\r\nHello \rWorld!\n\r\n*3\r\n$3\r\nfoo\r\n$-1\r\n$3\r\nbar\r\n*1\r\n*1\r\n*1\r\n*2\r\n:-1\r\n+OK\r\n*-1\r\n", | |
Object::Array(Some(vec![ | |
Object::SimpleString("OK".into()), | |
Object::Error("Error Message".into()), | |
Object::Integer(1000), | |
Object::Integer(0), | |
Object::Integer(-10000), | |
Object::BulkString(Some("".into())), | |
Object::BulkString(Some("Hello \rWorld!\n".into())), | |
Object::BulkString(Some("$198\r\nMcpJxQSaoHhgLbAsdPTRjGFCZqrwlEivnIuKkfOVYzWDNXmyUetB\r\n\r\n\r\x0c \n\t\x0b\r\n\r\n^|,!%#`:_<*-=?;$[\\&.>\"(~+]'/}{)@\r\n\r\nETbgCoRcXOWezIHKxmQqVvhADyGZJplufNFjakdYPrwMLsBitUSn\r\n\r\n_)~&!\\.[<=*]^>+;$%(/'@,}\"{:-|#`?\r\n\r\n\r\n \t\r\n".into())), | |
Object::Array(Some(vec![])), | |
Object::Array(Some(vec![ | |
Object::SimpleString("OK".into()), | |
Object::Error("Error Message".into()), | |
Object::Integer(1000), | |
Object::BulkString(Some("Hello \rWorld!\n".into()))])), | |
Object::Array(Some(vec![ | |
Object::BulkString(Some("foo".into())), | |
Object::BulkString(None), | |
Object::BulkString(Some("bar".into()))])), | |
Object::Array(Some(vec![ | |
Object::Array(Some(vec![ | |
Object::Array(Some(vec![ | |
Object::Array(Some(vec![ | |
Object::Integer(-1), | |
Object::SimpleString("OK".into()) | |
]))]))]))])), | |
Object::Array(None) | |
])) | |
) | |
} | |
} | |
mod parser { | |
use std::str::from_utf8; | |
use nom::{ | |
branch::alt, | |
bytes::complete::{tag, take, take_until1}, | |
character::complete::i64, | |
combinator::{all_consuming, cond, flat_map, map}, | |
multi::count, | |
sequence::{delimited, pair}, | |
IResult, | |
}; | |
use super::{Object, DELIMITER}; | |
pub(super) fn parse(input: &[u8]) -> IResult<&[u8], Object> { | |
all_consuming(parse_item)(input) | |
} | |
fn parse_item(input: &[u8]) -> IResult<&[u8], Object> { | |
alt(( | |
map( | |
delimited(tag("+"), take_until1(DELIMITER), tag(DELIMITER)), | |
|s: &[u8]| Object::SimpleString(from_utf8(s).unwrap().into()), | |
), | |
map( | |
delimited(tag("-"), take_until1(DELIMITER), tag(DELIMITER)), | |
|s: &[u8]| Object::Error(from_utf8(s).unwrap().into()), | |
), | |
map(delimited(tag(":"), i64, tag(DELIMITER)), Object::Integer), | |
flat_map(delimited(tag("$"), i64, tag(DELIMITER)), |i| { | |
map( | |
cond(i != -1, pair(take(i as usize), tag(DELIMITER))), | |
|s: Option<(&[u8], _)>| { | |
if let Some((s, _)) = s { | |
Object::BulkString(Some(from_utf8(s).unwrap().into())) | |
} else { | |
Object::BulkString(None) | |
} | |
}, | |
) | |
}), | |
flat_map(delimited(tag("*"), i64, tag(DELIMITER)), |i| { | |
map(cond(i != -1, count(parse_item, i as usize)), Object::Array) | |
}), | |
))(input) | |
} | |
} | |
impl FromStr for Object { | |
type Err = nom::error::Error<String>; | |
fn from_str(s: &str) -> Result<Self, Self::Err> { | |
match parse(s.as_bytes()).finish() { | |
Ok((_, obj)) => Ok(obj), | |
Err(nom::error::Error { input, code }) => Err(nom::error::Error { | |
input: std::str::from_utf8(input).unwrap().into(), | |
code, | |
}), | |
} | |
} | |
} | |
#[cfg(test)] | |
mod parse_tests { | |
use nom::error::ErrorKind; | |
use super::*; | |
macro_rules! parse_test { | |
{ | |
$(($name:ident, $case:literal)),* | |
} => { | |
$(#[test] | |
fn $name() { | |
match $case.parse::<Object>() { | |
Ok(obj) => assert_eq!( | |
$case, | |
obj.to_string(), | |
"encode failed, case={}, obj={:?}", | |
$case, | |
obj | |
), | |
Err(e) => assert!(false, "parse error. err={}", e), | |
} | |
})* | |
}; | |
{ | |
$(($name:ident, $case:literal, $input:literal, $code:expr)),* | |
} => { | |
$(#[test] | |
fn $name() { | |
match $case.parse::<Object>() { | |
Err(nom::error::Error { input, code }) => { | |
assert_eq!($input, input); | |
assert_eq!($code, code); | |
} | |
Ok(obj) => { | |
assert!(false, "should parse fail. case={}, parsed={}", $case, obj) | |
} | |
} | |
})* | |
}; | |
} | |
parse_test! { | |
(test_simple_string, "+OK\r\n"), | |
(test_error, "-Error Message\r\n"), | |
(test_integer, ":1000\r\n"), | |
(test_zero, ":0\r\n"), | |
(test_negative, ":-10000\r\n"), | |
(test_null_bulk_string, "$-1\r\n"), | |
(test_empty_bulk_string, "$0\r\n\r\n"), | |
(test_bulk_string, "$14\r\nHello \rWorld!\n\r\n"), | |
( | |
test_complex_bulk_string, | |
"$206\r\n$198\r\nMcpJxQSaoHhgLbAsdPTRjGFCZqrwlEivnIuKkfOVYzWDNXmyUetB\r\n\r\n\r\x0c \n\t\x0b\r\n\r\n^|,!%#`:_<*-=?;$[\\&.>\"(~+]'/}{)@\r\n\r\nETbgCoRcXOWezIHKxmQqVvhADyGZJplufNFjakdYPrwMLsBitUSn\r\n\r\n_)~&!\\.[<=*]^>+;$%(/'@,}\"{:-|#`?\r\n\r\n\r\n \t\r\n\r\n" | |
), | |
(test_utf8_bulk_string, "$19\r\n你好,\n世界!\r\n"), | |
(test_null_array, "*-1\r\n"), | |
(test_empty_array, "*0\r\n"), | |
( | |
test_array, | |
"*4\r\n+OK\r\n-Error Message\r\n:1000\r\n$14\r\nHello \rWorld!\n\r\n" | |
), | |
( | |
test_null_in_array, | |
"*3\r\n$3\r\nfoo\r\n$-1\r\n$3\r\nbar\r\n" | |
), | |
( | |
test_nested_array, | |
"*1\r\n*1\r\n*1\r\n*2\r\n:-1\r\n+OK\r\n" | |
), | |
( | |
test_complex_array, | |
"*13\r\n+OK\r\n-Error Message\r\n:1000\r\n:0\r\n:-10000\r\n$0\r\n\r\n$14\r\nHello \rWorld!\n\r\n$206\r\n$198\r\nMcpJxQSaoHhgLbAsdPTRjGFCZqrwlEivnIuKkfOVYzWDNXmyUetB\r\n\r\n\r\x0c \n\t\x0b\r\n\r\n^|,!%#`:_<*-=?;$[\\&.>\"(~+]'/}{)@\r\n\r\nETbgCoRcXOWezIHKxmQqVvhADyGZJplufNFjakdYPrwMLsBitUSn\r\n\r\n_)~&!\\.[<=*]^>+;$%(/'@,}\"{:-|#`?\r\n\r\n\r\n \t\r\n\r\n*0\r\n*4\r\n+OK\r\n-Error Message\r\n:1000\r\n$14\r\nHello \rWorld!\n\r\n*3\r\n$3\r\nfoo\r\n$-1\r\n$3\r\nbar\r\n*1\r\n*1\r\n*1\r\n*2\r\n:-1\r\n+OK\r\n*-1\r\n" | |
) | |
} | |
parse_test! { | |
(test_invalid_head, "?123\r\n", "?123\r\n", ErrorKind::Tag), | |
(test_missing_delimeter, ":123", ":123", ErrorKind::Tag), | |
(test_wrong_data_type, ":hello\r\n", ":hello\r\n", ErrorKind::Tag) | |
} | |
} | |
} | |
fn main() {} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment