1use ::core::future::Future;
2
3use tokio::io::{AsyncRead, AsyncReadExt as _, AsyncWrite, AsyncWriteExt as _};
4use tokio_util::bytes::{Buf as _, BytesMut};
5use tokio_util::codec::{Decoder, Encoder};
6
7fn invalid_utf8() -> std::io::Error {
8 std::io::Error::new(std::io::ErrorKind::InvalidInput, "value is not valid UTF8")
9}
10
11pub trait AsyncReadUtf8: AsyncRead {
12 #[cfg_attr(
13 feature = "tracing",
14 tracing::instrument(level = "trace", ret, skip_all)
15 )]
16 fn read_char_utf8(&mut self) -> impl Future<Output = std::io::Result<char>>
17 where
18 Self: Unpin,
19 {
20 async move {
21 let b = self.read_u8().await?;
22 let i = if b & 0x80 == 0 {
23 u32::from(b)
24 } else if b & 0b1110_0000 == 0b1100_0000 {
25 if b < 0b1100_0010 {
26 return Err(invalid_utf8());
27 }
28 let b2 = self.read_u8().await?;
29 if b2 & 0b1100_0000 != 0b1000_0000 {
30 return Err(invalid_utf8());
31 }
32 u32::from(b & 0b0001_1111) << 6 | u32::from(b2 & 0b0011_1111)
33 } else if b & 0b1111_0000 == 0b1110_0000 {
34 let b2 = self.read_u8().await?;
35 if b2 & 0b1100_0000 != 0b1000_0000 || (b == 0b1110_0000 && b2 < 0b1010_0000) {
36 return Err(invalid_utf8());
37 }
38 let b3 = self.read_u8().await?;
39 if b3 & 0b1100_0000 != 0b1000_0000 {
40 return Err(invalid_utf8());
41 }
42 u32::from(b & 0b0000_1111) << 12
43 | u32::from(b2 & 0b0011_1111) << 6
44 | u32::from(b3 & 0b0011_1111)
45 } else if b & 0b1111_1000 == 0b1111_0000 {
46 let b2 = self.read_u8().await?;
47 if b2 & 0b1100_0000 != 0b1000_0000 || (b == 0b1111_0000 && b2 < 0b1001_0000) {
48 return Err(invalid_utf8());
49 }
50 let b3 = self.read_u8().await?;
51 if b3 & 0b1100_0000 != 0b1000_0000 {
52 return Err(invalid_utf8());
53 }
54 let b4 = self.read_u8().await?;
55 if b4 & 0b1100_0000 != 0b1000_0000 {
56 return Err(invalid_utf8());
57 }
58 u32::from(b & 0b0000_0111) << 18
59 | u32::from(b2 & 0b0011_1111) << 12
60 | u32::from(b3 & 0b0011_1111) << 6
61 | u32::from(b4 & 0b0011_1111)
62 } else {
63 return Err(invalid_utf8());
64 };
65 i.try_into()
66 .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))
67 }
68 }
69}
70
71impl<T: AsyncRead> AsyncReadUtf8 for T {}
72
73pub trait AsyncWriteUtf8: AsyncWrite {
74 #[cfg_attr(
75 feature = "tracing",
76 tracing::instrument(level = "trace", ret, skip_all)
77 )]
78 fn write_char_utf8(&mut self, x: char) -> impl Future<Output = std::io::Result<()>>
79 where
80 Self: Unpin,
81 {
82 async move { self.write_all(x.encode_utf8(&mut [0; 4]).as_bytes()).await }
83 }
84}
85
86impl<T: AsyncWrite> AsyncWriteUtf8 for T {}
87
88#[derive(Copy, Clone, Debug, Default, Eq, PartialEq)]
89pub struct Utf8Codec;
90
91impl Decoder for Utf8Codec {
92 type Item = char;
93 type Error = std::io::Error;
94
95 fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
96 let Some(b) = src.first().copied() else {
97 src.reserve(1);
98 return Ok(None);
99 };
100 let i = if b & 0x80 == 0 {
101 src.advance(1);
102 u32::from(b)
103 } else if b & 0b1110_0000 == 0b1100_0000 {
104 if b < 0b1100_0010 {
105 return Err(invalid_utf8());
106 }
107 let Some(b2) = src.get(1).copied() else {
108 src.reserve(1);
109 return Ok(None);
110 };
111 if b2 & 0b1100_0000 != 0b1000_0000 {
112 return Err(invalid_utf8());
113 }
114 src.advance(2);
115 u32::from(b & 0b0001_1111) << 6 | u32::from(b2 & 0b0011_1111)
116 } else if b & 0b1111_0000 == 0b1110_0000 {
117 let Some(b2) = src.get(1).copied() else {
118 src.reserve(2);
119 return Ok(None);
120 };
121 if b2 & 0b1100_0000 != 0b1000_0000 || (b == 0b1110_0000 && b2 < 0b1010_0000) {
122 return Err(invalid_utf8());
123 }
124 let Some(b3) = src.get(2).copied() else {
125 src.reserve(1);
126 return Ok(None);
127 };
128 if b3 & 0b1100_0000 != 0b1000_0000 {
129 return Err(invalid_utf8());
130 }
131 src.advance(3);
132 u32::from(b & 0b0000_1111) << 12
133 | u32::from(b2 & 0b0011_1111) << 6
134 | u32::from(b3 & 0b0011_1111)
135 } else if b & 0b1111_1000 == 0b1111_0000 {
136 let Some(b2) = src.get(1).copied() else {
137 src.reserve(3);
138 return Ok(None);
139 };
140 if b2 & 0b1100_0000 != 0b1000_0000 || (b == 0b1111_0000 && b2 < 0b1001_0000) {
141 return Err(invalid_utf8());
142 }
143 let Some(b3) = src.get(2).copied() else {
144 src.reserve(2);
145 return Ok(None);
146 };
147 if b3 & 0b1100_0000 != 0b1000_0000 {
148 return Err(invalid_utf8());
149 }
150 let Some(b4) = src.get(3).copied() else {
151 src.reserve(1);
152 return Ok(None);
153 };
154 if b4 & 0b1100_0000 != 0b1000_0000 {
155 return Err(invalid_utf8());
156 }
157 src.advance(4);
158 u32::from(b & 0b0000_0111) << 18
159 | u32::from(b2 & 0b0011_1111) << 12
160 | u32::from(b3 & 0b0011_1111) << 6
161 | u32::from(b4 & 0b0011_1111)
162 } else {
163 return Err(invalid_utf8());
164 };
165 let c = i
166 .try_into()
167 .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
168 Ok(Some(c))
169 }
170
171 fn decode_eof(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
172 match self.decode(src)? {
173 Some(frame) => Ok(Some(frame)),
174 None => {
175 if src.is_empty() {
176 Ok(None)
177 } else {
178 Err(std::io::Error::from(std::io::ErrorKind::UnexpectedEof))
179 }
180 }
181 }
182 }
183}
184
185impl Encoder<char> for Utf8Codec {
186 type Error = std::io::Error;
187
188 fn encode(&mut self, x: char, dst: &mut BytesMut) -> Result<(), Self::Error> {
189 dst.extend_from_slice(x.encode_utf8(&mut [0; 4]).as_bytes());
190 Ok(())
191 }
192}
193
194impl Encoder<&char> for Utf8Codec {
195 type Error = std::io::Error;
196
197 fn encode(&mut self, x: &char, dst: &mut BytesMut) -> Result<(), Self::Error> {
198 self.encode(*x, dst)
199 }
200}
201
202impl Encoder<&&char> for Utf8Codec {
203 type Error = std::io::Error;
204
205 fn encode(&mut self, x: &&char, dst: &mut BytesMut) -> Result<(), Self::Error> {
206 self.encode(**x, dst)
207 }
208}
209
210#[cfg(test)]
211mod tests {
212 use super::*;
213
214 #[test_log::test(tokio::test)]
215 async fn codec() {
216 let v = '$'
217 .encode_utf8(&mut [0; 1])
218 .as_bytes()
219 .read_char_utf8()
220 .await
221 .expect("failed to read `$`");
222 assert_eq!(v, '$');
223
224 let v = '@'
225 .encode_utf8(&mut [0; 1])
226 .as_bytes()
227 .read_char_utf8()
228 .await
229 .expect("failed to read `@`");
230 assert_eq!(v, '@');
231
232 let v = 'И'
233 .encode_utf8(&mut [0; 2])
234 .as_bytes()
235 .read_char_utf8()
236 .await
237 .expect("failed to read `И`");
238 assert_eq!(v, 'И');
239
240 let v = 'ह'
241 .encode_utf8(&mut [0; 3])
242 .as_bytes()
243 .read_char_utf8()
244 .await
245 .expect("failed to read `ह`");
246 assert_eq!(v, 'ह');
247
248 let v = '€'
249 .encode_utf8(&mut [0; 3])
250 .as_bytes()
251 .read_char_utf8()
252 .await
253 .expect("failed to read `€`");
254 assert_eq!(v, '€');
255
256 let v = '한'
257 .encode_utf8(&mut [0; 3])
258 .as_bytes()
259 .read_char_utf8()
260 .await
261 .expect("failed to read `한`");
262 assert_eq!(v, '한');
263
264 let v = '𐍈'
265 .encode_utf8(&mut [0; 4])
266 .as_bytes()
267 .read_char_utf8()
268 .await
269 .expect("failed to read `𐍈`");
270 assert_eq!(v, '𐍈');
271 }
272
273 #[test_log::test]
274 fn truncated_eof() {
275 let mut src = BytesMut::from("€".as_bytes()[..2].as_ref());
276 let err = Utf8Codec
277 .decode_eof(&mut src)
278 .expect_err("truncated char must fail");
279 assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
280
281 assert!(Utf8Codec
282 .decode_eof(&mut BytesMut::new())
283 .unwrap()
284 .is_none());
285 }
286
287 #[test_log::test(tokio::test)]
288 async fn truncated_invalid() {
289 for mut src in [b"\xe2\x28".as_slice(), b"\xf0\x9f\x28"] {
290 let err = Utf8Codec
291 .decode_eof(&mut BytesMut::from(src))
292 .expect_err("invalid char must fail");
293 assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput);
294
295 let err = src
296 .read_char_utf8()
297 .await
298 .expect_err("invalid char must fail");
299 assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput);
300 }
301 }
302
303 #[test_log::test(tokio::test)]
304 async fn overlong() {
305 for mut src in [
306 b"\xc0\xaf".as_slice(),
307 b"\xc1\xbf",
308 b"\xe0\x80\xaf",
309 b"\xe0\x9f\xbf",
310 b"\xf0\x80\x80\xaf",
311 b"\xf0\x8f\xbf\xbf",
312 b"\xc0",
313 b"\xe0\x80",
314 b"\xf0\x80",
315 ] {
316 let err = Utf8Codec
317 .decode_eof(&mut BytesMut::from(src))
318 .expect_err("overlong encoding must fail");
319 assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput);
320
321 let err = src
322 .read_char_utf8()
323 .await
324 .expect_err("overlong encoding must fail");
325 assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput);
326 }
327
328 for c in ['\u{80}', '\u{800}', '\u{10000}'] {
329 let mut buf = [0; 4];
330 let mut src = c.encode_utf8(&mut buf).as_bytes();
331 assert_eq!(
332 Utf8Codec.decode(&mut BytesMut::from(src)).unwrap().unwrap(),
333 c
334 );
335 assert_eq!(src.read_char_utf8().await.unwrap(), c);
336 }
337 }
338}