telnet_codec/
lib.rs

1use tokio_util::codec::{Encoder, Decoder};
2use bytes::{BytesMut, Buf, BufMut, Bytes};
3use std::{
4    io
5};
6
7const SE: u8 = 240;
8const SB: u8 = 250;
9const WILL: u8 = 251;
10const WONT: u8 = 252;
11const DO: u8 = 253;
12const DONT: u8 = 254;
13const IAC: u8 = 255;
14
15// TelnetEvents are the bread and butter of this Codec.
16#[derive(Clone, Debug)]
17pub enum TelnetEvent {
18    // WILL|WONT|DO|DONT <OPTION>
19    Negotiate(u8, u8),
20
21    // IAC SB <OPTION> <DATA> IAC SE
22    SubNegotiate(u8, Bytes),
23
24    // Raw data. The application will have to figure out what these mean.
25    Data(Bytes),
26
27    // An IAC <command> other than those involved in negotiation and sub-options.
28    Command(u8)
29}
30
31impl From<TelnetEvent> for Bytes {
32    fn from(src: TelnetEvent) -> Self {
33        let mut out = BytesMut::new();
34
35        match src {
36            TelnetEvent::Data(data) => {
37                out.reserve(data.len());
38                out.put(data);
39            },
40            TelnetEvent::Negotiate(comm, op) => {
41                out.reserve(3);
42                out.extend(&[IAC, comm, op]);
43            },
44            TelnetEvent::SubNegotiate(op, data) => {
45                out.reserve(5 + data.len());
46                out.extend(&[IAC, SB, op]);
47                out.extend(data);
48                out.extend(&[IAC, SB]);
49            },
50            TelnetEvent::Command(byte) => {
51                out.reserve(2);
52                out.extend(&[IAC, byte]);
53            }
54        }
55        out.freeze()
56    }
57}
58
59#[derive(Debug)]
60pub struct TelnetCodec {
61    max_buffer: usize,
62}
63
64impl TelnetCodec {
65    pub fn new(max_buffer: usize) -> Self {
66
67        TelnetCodec {
68            max_buffer,
69        }
70    }
71}
72
73impl Default for TelnetCodec {
74    fn default() -> Self {
75        Self::new(1024)
76    }
77}
78
79impl Decoder for TelnetCodec {
80    type Item = TelnetEvent;
81    type Error = io::Error;
82
83    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
84
85        if src.is_empty() {
86            return Ok(None);
87        }
88
89        if src[0] == IAC {
90            if src.len() > 1 {
91                match src[1] {
92                    IAC => {
93                        // This is an escaped IAC. Send it onwards as data.
94                        src.advance(2);
95                        let mut data = BytesMut::with_capacity(1);
96                        data.put_u8(IAC);
97                        return Ok(Some(TelnetEvent::Data(data.freeze())));
98                    },
99                    WILL | WONT | DO | DONT => {
100                        if src.len() > 2 {
101                            let answer = TelnetEvent::Negotiate(src[1], src[2]);
102                            src.advance(3);
103                            return Ok(Some(answer));
104                        } else {
105                            // Not enough bytes for negotiation...yet.
106                            return Ok(None)
107                        }
108                    },
109                    SB => {
110                        // Since the valid signature is IAC SB <option> <data> IAC SE, and data might be empty, we need at least 5 bytes.
111                        if src.len() > 4 {
112                            if let Some(ipos) = src.as_ref().windows(2).position(|b| b[0] == IAC && b[1] == SE) {
113                                // Split off any available up to an IAC and stuff it in the sub data buffer.
114                                let mut data = src.split_to(ipos);
115                                src.advance(2);
116                                let discard = data.split_to(3);
117                                let answer = TelnetEvent::SubNegotiate(discard[2], data.freeze());
118                                return Ok(Some(answer))
119                            } else {
120                                return Ok(None)
121                            }
122                        } else {
123                            // Not enough bytes for sub-negotiation...yet.
124                            return Ok(None)
125                        }
126                    },
127                    _ => {
128                        // Anything that's not the above is a simple IAC Command.
129                        let cmd = src[1];
130                        src.advance(2);
131                        return Ok(Some(TelnetEvent::Command(cmd)))
132                    }
133                }
134            } else {
135                // Need more bytes than a single IAC...
136                return Ok(None)
137            }
138        } else {
139            if let Some(ipos) = src.as_ref().iter().position(|b| b == &IAC) {
140                // Split off any available up to an IAC and stuff it in the sub data buffer.
141                return Ok(Some(TelnetEvent::Data(src.split_to(ipos).freeze())))
142            } else {
143                return Ok(Some(TelnetEvent::Data(src.split_to(src.len()).freeze())))
144            }
145        }
146    }
147}
148
149impl Encoder<TelnetEvent> for TelnetCodec {
150    type Error = io::Error;
151
152    fn encode(&mut self, item: TelnetEvent, dst: &mut BytesMut) -> Result<(), Self::Error> {
153        let out = Bytes::from(item);
154        dst.reserve(out.len());
155        dst.put(out.as_ref());
156        Ok(())
157    }
158}
159
160#[cfg(test)]
161mod tests {
162    #[test]
163    fn it_works() {
164        assert_eq!(2 + 2, 4);
165    }
166}