解答

1. use std::convert::TryFrom;
2. use thiserror::Error;

4. #[derive(Debug, Error)]
5. enum Error {
6. #[error("Invalid varint")]
7. InvalidVarint,
8. #[error("Invalid wire-type")]
9. InvalidWireType,
10. #[error("Unexpected EOF")]
11. UnexpectedEOF,
12. #[error("Invalid length")]
13. InvalidSize(#[from] std::num::TryFromIntError),
14. #[error("Unexpected wire-type)")]
15. UnexpectedWireType,
16. #[error("Invalid string (not UTF-8)")]
17. InvalidString,
18. }

20. /// A wire type as seen on the wire.
21. enum WireType {
22. /// Varint WireType 表明该值为单个 VARINT。
23. Varint,
24. //I64,  -- not needed for this exercise
25. /// The Len WireType indicates that the value is a length represented as a
26. /// VARINT followed by exactly that number of bytes.
27. Len,
28. /// The I32 WireType indicates that the value is precisely 4 bytes in
29. /// little-endian order containing a 32-bit signed integer.
30. I32,
31. }

33. #[derive(Debug)]
34. /// A field's value, typed based on the wire type.
35. enum FieldValue<'a> {
36. Varint(u64),
37. //I64(i64),  -- not needed for this exercise
38. Len(&'a [u8]),
39. I32(i32),
40. }

42. #[derive(Debug)]
43. /// A field, containing the field number and its value.
44. struct Field<'a> {
45. field_num: u64,
46. value: FieldValue<'a>,
47. }

49. trait ProtoMessage<'a>: Default + 'a {
50. fn add_field(&mut self, field: Field<'a>) -> Result<(), Error>;
51. }

53. impl TryFrom<u64> for WireType {
54. type Error = Error;

56. fn try_from(value: u64) -> Result<WireType, Error> {
57. Ok(match value {
58. 0 => WireType::Varint,
59. //1 => WireType::I64,  -- not needed for this exercise
60. 2 => WireType::Len,
61. 5 => WireType::I32,
62. _ => return Err(Error::InvalidWireType),
63. })
64. }
65. }

67. impl<'a> FieldValue<'a> {
68. fn as_string(&self) -> Result<&'a str, Error> {
69. let FieldValue::Len(data) = self else {
70. return Err(Error::UnexpectedWireType);
71. };
72. std::str::from_utf8(data).map_err(|_| Error::InvalidString)
73. }

75. fn as_bytes(&self) -> Result<&'a [u8], Error> {
76. let FieldValue::Len(data) = self else {
77. return Err(Error::UnexpectedWireType);
78. };
79. Ok(data)
80. }

82. fn as_u64(&self) -> Result<u64, Error> {
83. let FieldValue::Varint(value) = self else {
84. return Err(Error::UnexpectedWireType);
85. };
86. Ok(*value)
87. }
88. }

90. /// Parse a VARINT, returning the parsed value and the remaining bytes.
91. fn parse_varint(data: &[u8]) -> Result<(u64, &[u8]), Error> {
92. for i in 0..7 {
93. let Some(b) = data.get(i) else {
94. return Err(Error::InvalidVarint);
95. };
96. if b & 0x80 == 0 {
97. // This is the last byte of the VARINT, so convert it to
98. // a u64 and return it.
99. let mut value = 0u64;
100. for b in data[..=i].iter().rev() {
101. value = (value << 7) | (b & 0x7f) as u64;
102. }
103. return Ok((value, &data[i + 1..]));
104. }
105. }

107. // More than 7 bytes is invalid.
108. Err(Error::InvalidVarint)
109. }

111. /// Convert a tag into a field number and a WireType.
112. fn unpack_tag(tag: u64) -> Result<(u64, WireType), Error> {
113. let field_num = tag >> 3;
114. let wire_type = WireType::try_from(tag & 0x7)?;
115. Ok((field_num, wire_type))
116. }

118. /// Parse a field, returning the remaining bytes
119. fn parse_field(data: &[u8]) -> Result<(Field, &[u8]), Error> {
120. let (tag, remainder) = parse_varint(data)?;
121. let (field_num, wire_type) = unpack_tag(tag)?;
122. let (fieldvalue, remainder) = match wire_type {
123. WireType::Varint => {
124. let (value, remainder) = parse_varint(remainder)?;
125. (FieldValue::Varint(value), remainder)
126. }
127. WireType::Len => {
128. let (len, remainder) = parse_varint(remainder)?;
129. let len: usize = len.try_into()?;
130. if remainder.len() < len {
131. return Err(Error::UnexpectedEOF);
132. }
133. let (value, remainder) = remainder.split_at(len);
134. (FieldValue::Len(value), remainder)
135. }
136. WireType::I32 => {
137. if remainder.len() < 4 {
138. return Err(Error::UnexpectedEOF);
139. }
140. let (value, remainder) = remainder.split_at(4);
141. // Unwrap error because `value` is definitely 4 bytes long.
142. let value = i32::from_le_bytes(value.try_into().unwrap());
143. (FieldValue::I32(value), remainder)
144. }
145. };
146. Ok((Field { field_num, value: fieldvalue }, remainder))
147. }

149. /// Parse a message in the given data, calling `T::add_field` for each field in
150. /// the message.
151. ///
152. /// The entire input is consumed.
153. fn parse_message<'a, T: ProtoMessage<'a>>(mut data: &'a [u8]) -> Result<T, Error> {
154. let mut result = T::default();
155. while !data.is_empty() {
156. let parsed = parse_field(data)?;
157. result.add_field(parsed.0)?;
158. data = parsed.1;
159. }
160. Ok(result)
161. }

163. #[derive(Debug, Default)]
164. struct PhoneNumber<'a> {
165. number: &'a str,
166. type_: &'a str,
167. }

169. #[derive(Debug, Default)]
170. struct Person<'a> {
171. name: &'a str,
172. id: u64,
173. phone: Vec<PhoneNumber<'a>>,
174. }

176. impl<'a> ProtoMessage<'a> for Person<'a> {
177. fn add_field(&mut self, field: Field<'a>) -> Result<(), Error> {
178. match field.field_num {
179. 1 => self.name = field.value.as_string()?,
180. 2 => self.id = field.value.as_u64()?,
181. 3 => self.phone.push(parse_message(field.value.as_bytes()?)?),
182. _ => {} // skip everything else
183. }
184. Ok(())
185. }
186. }

188. impl<'a> ProtoMessage<'a> for PhoneNumber<'a> {
189. fn add_field(&mut self, field: Field<'a>) -> Result<(), Error> {
190. match field.field_num {
191. 1 => self.number = field.value.as_string()?,
192. 2 => self.type_ = field.value.as_string()?,
193. _ => {} // skip everything else
194. }
195. Ok(())
196. }
197. }

199. fn main() {
200. let person: Person = parse_message(&[
201. 0x0a, 0x07, 0x6d, 0x61, 0x78, 0x77, 0x65, 0x6c, 0x6c, 0x10, 0x2a, 0x1a,
202. 0x16, 0x0a, 0x0e, 0x2b, 0x31, 0x32, 0x30, 0x32, 0x2d, 0x35, 0x35, 0x35,
203. 0x2d, 0x31, 0x32, 0x31, 0x32, 0x12, 0x04, 0x68, 0x6f, 0x6d, 0x65, 0x1a,
204. 0x18, 0x0a, 0x0e, 0x2b, 0x31, 0x38, 0x30, 0x30, 0x2d, 0x38, 0x36, 0x37,
205. 0x2d, 0x35, 0x33, 0x30, 0x38, 0x12, 0x06, 0x6d, 0x6f, 0x62, 0x69, 0x6c,
206. 0x65,
207. ])
208. .unwrap();
209. println!("{:#?}", person);
210. }

212. #[cfg(test)]
213. mod test {
214. use super::*;

216. #[test]
217. fn as_string() {
218. assert!(FieldValue::Varint(10).as_string().is_err());
219. assert!(FieldValue::I32(10).as_string().is_err());
220. assert_eq!(FieldValue::Len(b"hello").as_string().unwrap(), "hello");
221. }

223. #[test]
224. fn as_bytes() {
225. assert!(FieldValue::Varint(10).as_bytes().is_err());
226. assert!(FieldValue::I32(10).as_bytes().is_err());
227. assert_eq!(FieldValue::Len(b"hello").as_bytes().unwrap(), b"hello");
228. }

230. #[test]
231. fn as_u64() {
232. assert_eq!(FieldValue::Varint(10).as_u64().unwrap(), 10u64);
233. assert!(FieldValue::I32(10).as_u64().is_err());
234. assert!(FieldValue::Len(b"hello").as_u64().is_err());
235. }
236. }