net_mc/
conn.rs

1use super::packet::*;
2use aes::Aes128;
3use anyhow::Result;
4use cfb8::{
5    cipher::{AsyncStreamCipher, NewCipher},
6    Cfb8,
7};
8use std::net::{Shutdown, SocketAddr, TcpStream};
9use std::time::Duration;
10use std::{
11    convert::TryFrom,
12    io::{self, Write},
13    net::{TcpListener, ToSocketAddrs},
14};
15
16#[derive(Clone, Copy)]
17pub enum ProtocolState {
18    Handshake,
19    Status,
20    Login,
21    Play,
22}
23
24impl ProtocolState {
25    pub fn to_i32(&self) -> i32 {
26        match self {
27            ProtocolState::Handshake => 0,
28            ProtocolState::Status => 1,
29            ProtocolState::Login => 2,
30            ProtocolState::Play => 3,
31        }
32    }
33}
34
35pub struct Listener(pub TcpListener);
36
37impl Listener {
38    pub fn bind(addr: impl ToSocketAddrs) -> Result<Listener> {
39        Ok(Listener(TcpListener::bind(addr)?))
40    }
41    pub fn accept(&mut self) -> Result<Conn> {
42        Ok(Conn::try_from(self.0.accept()?.0)?)
43    }
44}
45
46/// Conn wraps around TcpStream to simplify sending and receiving packets.
47pub struct Conn {
48    pub peer: SocketAddr,
49    pub stream: TcpStream,
50    /// State is set to Handshake on connect but is not handled by Conn.
51    pub state: ProtocolState,
52    cipher: Option<Cipher>,
53    pub writer: io::BufWriter<TcpStream>,
54    pub reader: io::BufReader<TcpStream>,
55    pub threshhold: i32,
56}
57
58struct Cipher {
59    write: Cfb8<Aes128>,
60    read: Cfb8<Aes128>,
61}
62
63impl io::Write for Conn {
64    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
65        match self.cipher.as_mut() {
66            Some(ref mut cipher) => {
67                let mut data = vec![0; buf.len()];
68                data[..buf.len()].clone_from_slice(&buf[..]);
69                cipher.write.encrypt(&mut data);
70
71                self.writer.write(&data)
72            }
73            None => self.writer.write(buf),
74        }
75    }
76    fn flush(&mut self) -> io::Result<()> {
77        self.writer.flush()
78    }
79}
80
81impl io::Read for Conn {
82    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
83        match self.cipher.as_mut() {
84            Option::None => self.reader.read(buf),
85            Option::Some(cipher) => {
86                let ret = self.reader.read(buf)?;
87                cipher.read.decrypt(&mut buf[..ret]);
88
89                Ok(ret)
90            }
91        }
92    }
93}
94
95impl TryFrom<TcpStream> for Conn {
96    type Error = anyhow::Error;
97    fn try_from(stream: TcpStream) -> Result<Self> {
98        let writer = io::BufWriter::new(stream.try_clone()?);
99        let reader = io::BufReader::new(stream.try_clone()?);
100        Ok(Self {
101            peer: stream.peer_addr()?,
102            stream,
103            state: ProtocolState::Handshake,
104            cipher: None,
105            writer,
106            reader,
107            threshhold: -1,
108        })
109    }
110}
111
112impl Conn {
113    pub fn connect(addr: SocketAddr) -> anyhow::Result<Self> {
114        let stream = TcpStream::connect(addr)?;
115        let writer = io::BufWriter::new(stream.try_clone()?);
116        let reader = io::BufReader::new(stream.try_clone()?);
117        Ok(Self {
118            peer: addr,
119            stream,
120            state: ProtocolState::Handshake,
121            cipher: None,
122            writer,
123            reader,
124            threshhold: -1,
125        })
126    }
127
128    pub fn connect_timeout(addr: &SocketAddr, timeout: Duration) -> anyhow::Result<Self> {
129        let stream = TcpStream::connect_timeout(addr, timeout)?;
130        let writer = io::BufWriter::new(stream.try_clone()?);
131        let reader = io::BufReader::new(stream.try_clone()?);
132        Ok(Self {
133            peer: *addr,
134            stream,
135            state: ProtocolState::Handshake,
136            cipher: None,
137            writer,
138            reader,
139            threshhold: -1,
140        })
141    }
142
143    pub fn shutdown(&mut self) -> io::Result<()> {
144        self.stream.shutdown(Shutdown::Both)
145    }
146
147    pub fn send_packet(&mut self, packet: &impl Packet) -> anyhow::Result<()> {
148        packet.encode()?.pack(self, self.threshhold)?;
149        self.flush()?;
150        Ok(())
151    }
152
153    pub fn read_packet(&mut self) -> Result<RawPacket> {
154        RawPacket::unpack(self, self.threshhold)
155    }
156
157    pub fn set_compression_threshhold(&mut self, threshhold: i32) {
158        self.threshhold = threshhold;
159    }
160
161    pub fn enable_encryption(&mut self, key: &[u8]) -> anyhow::Result<()> {
162        let cipher = Cfb8::<Aes128>::new_from_slices(key, key);
163        let write = match cipher {
164            Ok(c) => c,
165            Err(e) => return Err(anyhow::anyhow!("{}", e)),
166        };
167        let cipher = Cfb8::<Aes128>::new_from_slices(key, key);
168        let read = match cipher {
169            Ok(c) => c,
170            Err(e) => return Err(anyhow::anyhow!("{}", e)),
171        };
172        self.cipher = Some(Cipher { write, read });
173
174        Ok(())
175    }
176}
net_mc/conn.rs

net_mc/
conn.rs
