irc_proto/
line.rs

1//! Implementation of line-delimiting codec for Tokio.
2
3use std::io;
4
5use bytes::BytesMut;
6#[cfg(feature = "encoding")]
7use encoding::label::encoding_from_whatwg_label;
8#[cfg(feature = "encoding")]
9use encoding::{DecoderTrap, EncoderTrap, EncodingRef};
10use tokio_util::codec::{Decoder, Encoder};
11
12use crate::error;
13
14/// A line-based codec parameterized by an encoding.
15pub struct LineCodec {
16    #[cfg(feature = "encoding")]
17    encoding: EncodingRef,
18    next_index: usize,
19}
20
21impl LineCodec {
22    /// Creates a new instance of LineCodec from the specified encoding.
23    pub fn new(label: &str) -> error::Result<LineCodec> {
24        Ok(LineCodec {
25            #[cfg(feature = "encoding")]
26            encoding: match encoding_from_whatwg_label(label) {
27                Some(x) => x,
28                None => {
29                    return Err(error::ProtocolError::Io(io::Error::new(
30                        io::ErrorKind::InvalidInput,
31                        &format!("Attempted to use unknown codec {}.", label)[..],
32                    )));
33                }
34            },
35            next_index: 0,
36        })
37    }
38}
39
40impl Decoder for LineCodec {
41    type Item = String;
42    type Error = error::ProtocolError;
43
44    fn decode(&mut self, src: &mut BytesMut) -> error::Result<Option<String>> {
45        if let Some(offset) = src[self.next_index..].iter().position(|b| *b == b'\n') {
46            // Remove the next frame from the buffer.
47            let line = src.split_to(self.next_index + offset + 1);
48
49            // Set the search start index back to 0 since we found a newline.
50            self.next_index = 0;
51
52            #[cfg(feature = "encoding")]
53            {
54                // Decode the line using the codec's encoding.
55                match self.encoding.decode(line.as_ref(), DecoderTrap::Replace) {
56                    Ok(data) => Ok(Some(data)),
57                    Err(data) => Err(io::Error::new(
58                        io::ErrorKind::InvalidInput,
59                        &format!("Failed to decode {} as {}.", data, self.encoding.name())[..],
60                    )
61                    .into()),
62                }
63            }
64
65            #[cfg(not(feature = "encoding"))]
66            {
67                match String::from_utf8(line.to_vec()) {
68                    Ok(data) => Ok(Some(data)),
69                    Err(data) => Err(io::Error::new(
70                        io::ErrorKind::InvalidInput,
71                        &format!("Failed to decode {} as UTF-8.", data)[..],
72                    )
73                    .into()),
74                }
75            }
76        } else {
77            // Set the search start index to the current length since we know that none of the
78            // characters we've already looked at are newlines.
79            self.next_index = src.len();
80            Ok(None)
81        }
82    }
83}
84
85impl Encoder<String> for LineCodec {
86    type Error = error::ProtocolError;
87
88    fn encode(&mut self, msg: String, dst: &mut BytesMut) -> error::Result<()> {
89        #[cfg(feature = "encoding")]
90        {
91            // Encode the message using the codec's encoding.
92            let data: error::Result<Vec<u8>> = self
93                .encoding
94                .encode(&msg, EncoderTrap::Replace)
95                .map_err(|data| {
96                    io::Error::new(
97                        io::ErrorKind::InvalidInput,
98                        &format!("Failed to encode {} as {}.", data, self.encoding.name())[..],
99                    )
100                    .into()
101                });
102            // Write the encoded message to the output buffer.
103            dst.extend(&data?);
104        }
105
106        #[cfg(not(feature = "encoding"))]
107        {
108            dst.extend(msg.into_bytes());
109        }
110
111        Ok(())
112    }
113}