rs_jsonrpc_server_utils/
stream_codec.rs

1use std::{io, str};
2use tokio_io::codec::{Decoder, Encoder};
3use bytes::BytesMut;
4
5/// Separator for enveloping messages in streaming codecs
6#[derive(Debug, Clone)]
7pub enum Separator {
8	/// No envelope is expected between messages. Decoder will try to figure out
9	/// message boundaries by accumulating incoming bytes until valid JSON is formed.
10	/// Encoder will send messages without any boundaries between requests.
11	Empty,
12	/// Byte is used as an sentitel between messages
13	Byte(u8),
14}
15
16impl Default for Separator {
17	fn default() -> Self {
18		Separator::Byte(b'\n')
19	}
20}
21
22/// Stream codec for streaming protocols (ipc, tcp)
23#[derive(Debug, Default)]
24pub struct StreamCodec {
25	incoming_separator: Separator,
26	outgoing_separator: Separator,
27}
28
29impl StreamCodec {
30	/// Default codec with streaming input data. Input can be both enveloped and not.
31	pub fn stream_incoming() -> Self {
32		StreamCodec::new(Separator::Empty, Default::default())
33	}
34
35	/// New custom stream codec
36	pub fn new(incoming_separator: Separator, outgoing_separator: Separator) -> Self {
37		StreamCodec {
38			incoming_separator: incoming_separator,
39			outgoing_separator: outgoing_separator,
40		}
41	}
42}
43
44fn is_whitespace(byte: u8) -> bool {
45	match byte {
46		0x0D | 0x0A | 0x20 | 0x09 => true,
47		_ => false,
48	}
49}
50
51impl Decoder for StreamCodec {
52	type Item = String;
53	type Error = io::Error;
54
55	fn decode(&mut self, buf: &mut BytesMut) -> io::Result<Option<Self::Item>> {
56		if let Separator::Byte(separator) = self.incoming_separator {
57			if let Some(i) = buf.as_ref().iter().position(|&b| b == separator) {
58				let line = buf.split_to(i);
59				buf.split_to(1);
60
61				match str::from_utf8(&line.as_ref()) {
62					Ok(s) => Ok(Some(s.to_string())),
63					Err(_) => Err(io::Error::new(io::ErrorKind::Other, "invalid UTF-8")),
64				}
65			} else {
66				Ok(None)
67			}			
68		} else {
69			let mut depth = 0;
70			let mut in_str = false;
71			let mut is_escaped = false;
72			let mut start_idx = 0;
73			let mut whitespaces = 0;
74
75			for idx in 0..buf.as_ref().len() {
76				let byte = buf.as_ref()[idx];
77
78				if (byte == b'{' || byte == b'[') && !in_str {
79					if depth == 0 {
80						start_idx = idx;
81					}
82					depth += 1;
83				}
84				else if (byte == b'}' || byte == b']') && !in_str {
85					depth -= 1;
86				}
87				else if byte == b'"' && !is_escaped {
88					in_str = !in_str;
89				}
90				else if is_whitespace(byte) {
91					whitespaces += 1;
92				}
93				if byte == b'\\' && !is_escaped && in_str {
94					is_escaped = true;
95				} else {
96					is_escaped = false;
97				}
98
99				if depth == 0 && idx != start_idx && idx - start_idx + 1 > whitespaces {
100					let bts = buf.split_to(idx + 1);
101					match String::from_utf8(bts.as_ref().to_vec()) {
102						Ok(val) => { return Ok(Some(val)) },
103						Err(_) => { return Ok(None); } // skip non-utf requests (TODO: log error?)
104					};
105				}
106			}
107			Ok(None)
108		}
109	}
110}
111
112impl Encoder for StreamCodec {
113	type Item = String;
114	type Error = io::Error;
115	
116	fn encode(&mut self, msg: String, buf: &mut BytesMut) -> io::Result<()> {
117		let mut payload = msg.into_bytes();
118		if let Separator::Byte(separator) = self.outgoing_separator {
119			payload.push(separator);
120		}
121		buf.extend_from_slice(&payload);
122		Ok(())
123	}
124}
125
126#[cfg(test)]
127mod tests {
128
129	use super::StreamCodec;
130	use tokio_io::codec::Decoder;
131	use bytes::{BytesMut, BufMut};
132
133	#[test]
134	fn simple_encode() {
135		let mut buf = BytesMut::with_capacity(2048);
136		buf.put_slice(b"{ test: 1 }{ test: 2 }{ test: 3 }");
137
138		let mut codec = StreamCodec::stream_incoming();
139
140		let request = codec.decode(&mut buf)
141			.expect("There should be no error in simple test")
142			.expect("There should be at least one request in simple test");
143
144		assert_eq!(request, "{ test: 1 }");
145	}
146
147	#[test]
148	fn escape() {
149		let mut buf = BytesMut::with_capacity(2048);
150		buf.put_slice(br#"{ test: "\"\\" }{ test: "\ " }{ test: "\}" }[ test: "\]" ]"#);
151
152		let mut codec = StreamCodec::stream_incoming();
153
154		let request = codec.decode(&mut buf)
155			.expect("There should be no error in first escape test")
156			.expect("There should be a request in first escape test");
157
158		assert_eq!(request, r#"{ test: "\"\\" }"#);
159
160		let request2 = codec.decode(&mut buf)
161			.expect("There should be no error in 2nd escape test")
162			.expect("There should be a request in 2nd escape test");
163		assert_eq!(request2, r#"{ test: "\ " }"#);
164
165		let request3 = codec.decode(&mut buf)
166			.expect("There should be no error in 3rd escape test")
167			.expect("There should be a request in 3rd escape test");
168		assert_eq!(request3, r#"{ test: "\}" }"#);
169
170		let request4 = codec.decode(&mut buf)
171			.expect("There should be no error in 4th escape test")
172			.expect("There should be a request in 4th escape test");
173		assert_eq!(request4, r#"[ test: "\]" ]"#);
174	}
175
176	#[test]
177	fn whitespace() {
178		let mut buf = BytesMut::with_capacity(2048);
179		buf.put_slice(b"{ test: 1 }\n\n\n\n{ test: 2 }\n\r{\n test: 3 }  ");
180
181		let mut codec = StreamCodec::stream_incoming();
182
183		let request = codec.decode(&mut buf)
184			.expect("There should be no error in first whitespace test")
185			.expect("There should be a request in first whitespace test");
186
187		assert_eq!(request, "{ test: 1 }");
188
189		let request2 = codec.decode(&mut buf)
190			.expect("There should be no error in first 2nd test")
191			.expect("There should be aa request in 2nd whitespace test");
192		// TODO: maybe actually trim it out
193		assert_eq!(request2, "\n\n\n\n{ test: 2 }");
194
195		let request3 = codec.decode(&mut buf)
196			.expect("There should be no error in first 3rd test")
197			.expect("There should be a request in 3rd whitespace test");
198		assert_eq!(request3, "\n\r{\n test: 3 }");
199
200		let request4 = codec.decode(&mut buf)
201			.expect("There should be no error in first 4th test");
202		assert!(request4.is_none(), "There should be no 4th request because it contains only whitespaces");
203	}
204
205	#[test]
206	fn fragmented_encode() {
207		let mut buf = BytesMut::with_capacity(2048);
208		buf.put_slice(b"{ test: 1 }{ test: 2 }{ tes");
209
210		let mut codec = StreamCodec::stream_incoming();
211
212		let request = codec.decode(&mut buf)
213			.expect("There should be no error in first fragmented test")
214			.expect("There should be at least one request in first fragmented test");
215		assert_eq!(request, "{ test: 1 }");
216		codec.decode(&mut buf)
217			.expect("There should be no error in second fragmented test")
218			.expect("There should be at least one request in second fragmented test");
219		assert_eq!(String::from_utf8(buf.as_ref().to_vec()).unwrap(), "{ tes");
220
221		buf.put_slice(b"t: 3 }");
222		let request = codec.decode(&mut buf)
223			.expect("There should be no error in third fragmented test")
224			.expect("There should be at least one request in third fragmented test");
225		assert_eq!(request, "{ test: 3 }");
226	}
227
228	#[test]
229	fn huge() {
230		let request = r#"
231		{
232			"jsonrpc":"2.0",
233			"method":"say_hello",
234			"params": [
235				42,
236				0,
237				{
238					"from":"0xb60e8dd61c5d32be8058bb8eb970870f07233155",
239					"gas":"0x2dc6c0",
240					"data":"0x606060405260003411156010576002565b6001805433600160a060020a0319918216811790925560028054909116909117905561291f806100406000396000f3606060405236156100e55760e060020a600035046304029f2381146100ed5780630a1273621461015f57806317c1dd87146102335780631f9ea25d14610271578063266fa0e91461029357806349593f5314610429578063569aa0d8146104fc57806359a4669f14610673578063647a4d5f14610759578063656104f5146108095780636e9febfe1461082b57806370de8c6e1461090d57806371bde852146109ed5780638f30435d14610ab4578063916dbc1714610da35780639f5a7cd414610eef578063c91540f614610fe6578063eae99e1c146110b5578063fedc2a281461115a575b61122d610002565b61122d6004808035906020019082018035906020019191908080601f01602080910402602001604051908101604052809392919081815260200183838082843750949650509335935050604435915050606435600154600090600160a060020a03908116339091161461233357610002565b61122f6004808035906020019082018035906020019191908080601f016020809104026020016040519081016040528093929190818152602001838380828437509496505093359350506044359150506064355b60006000600060005086604051808280519060200190808383829060006004602084601f0104600f02600301f1509050019150509081526020016040518091039020600050905042816005016000508560ff1660028110156100025760040201835060010154604060020a90046001604060020a0316116115df576115d6565b6112416004355b604080516001604060020a038316408152606060020a33600160a060020a031602602082015290519081900360340190205b919050565b61122d600435600254600160a060020a0390811633909116146128e357610002565b61125e6004808035906020019082018035906020019191908080601f01602080910402602001604051908101604052809392919081815260200183838082843750949650509335935050505060006000600060006000600060005087604051808280519060200190808383829060006004602084601f0104600f02600301f1509050019150509081526020016040518091039020600050905080600001600050600087600160a060020a0316815260200190815260200160002060005060000160059054906101000a90046001604060020a03169450845080600001600050600087600160a060020a03168152602001908152602001600020600050600001600d9054906101000a90046001604060020a03169350835080600001600050600087600160a060020a0316815260200190815260200160002060005060000160009054906101000a900460ff169250825080600001600050600087600160a060020a0316815260200190815260200160002060005060000160019054906101000a900463ffffffff16915081505092959194509250565b61122d6004808035906020019082018035906020019191908080601f01602080910402602001604051908101604052809392919081815260200183838082843750949650509335935050604435915050606435608435600060006000600060005088604051808280519060200190808383829060006004602084601f0104600f02600301f15090500191505090815260200160405180910390206000509250346000141515611c0e5760405133600160a060020a0316908290349082818181858883f193505050501515611c1a57610002565b6112996004808035906020019082018035906020019191908080601f01602080910402602001604051908101604052809392919081815260200183838082843750949650509335935050604435915050600060006000600060006000600060006000508a604051808280519060200190808383829060006004602084601f0104600f02600301f15090500191505090815260200160405180910390206000509050806001016000508960ff16600281101561000257600160a060020a038a168452828101600101602052604084205463ffffffff1698506002811015610002576040842054606060020a90046001604060020a031697506002811015610002576040842054640100000000900463ffffffff169650600281101561000257604084206001015495506002811015610002576040842054604060020a900463ffffffff169450600281101561000257505060409091205495999498509296509094509260a060020a90046001604060020a0316919050565b61122d6004808035906020019082018035906020019191908080601f016020809104026020016040519081016040528093929190818152602001838380828437509496505050505050506000600060005082604051808280519060200190808383829060006004602084601f0104600f02600301f15090500191505090815260200160405180910390206000509050348160050160005082600d0160009054906101000a900460ff1660ff16600281101561000257600402830160070180546001608060020a0381169093016001608060020a03199390931692909217909155505b5050565b6112e26004808035906020019082018035906020019191908080601f01602080910003423423094734987103498712093847102938740192387401349857109487501938475"
241				}
242			]
243		}"#;
244
245		let mut buf = BytesMut::with_capacity(65536);
246		buf.put_slice(request.as_bytes());
247
248		let mut codec = StreamCodec::stream_incoming();
249
250		let parsed_request = codec.decode(&mut buf)
251			.expect("There should be no error in huge test")
252			.expect("There should be at least one request huge test");
253		assert_eq!(request, parsed_request);
254	}
255
256	#[test]
257	fn simple_line_codec() {
258		let mut buf = BytesMut::with_capacity(2048);
259		buf.put_slice(b"{ test: 1 }\n{ test: 2 }\n{ test: 3 }");
260
261		let mut codec = StreamCodec::default();
262
263		let request = codec.decode(&mut buf)
264			.expect("There should be no error in simple test")
265			.expect("There should be at least one request in simple test");
266		let request2 = codec.decode(&mut buf)
267			.expect("There should be no error in simple test")
268			.expect("There should be at least one request in simple test");
269
270		assert_eq!(request, "{ test: 1 }");
271		assert_eq!(request2, "{ test: 2 }");
272	}	
273}
rs_jsonrpc_server_utils/stream_codec.rs

rs_jsonrpc_server_utils/
stream_codec.rs
