1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
use std::error::Error;
use std::io;
use tokio::net::TcpStream;
use tokio_rustls::server::TlsStream;
use tokio::io::{split, AsyncWriteExt, AsyncReadExt};
pub enum ConnectionType{
Plain(Option<TcpStream>),
TLS(Option<TlsStream<TcpStream>>)
}
/// # Connection
///
/// This struct is a helpful struct to handle the nitty gritty of
/// connections, such as reading and writing to the stream
pub struct Connection{
/// The connection type - plain or TLS
connection_type: ConnectionType,
/// The read buffer size
read_buffer_size: usize,
/// We need to store the TLS stream so we don't run into copy/clone issues
tls_stream: Option<TlsStream<TcpStream>>,
}
impl Connection {
/// # New
///
/// Create a new connection handler from a `TcpStream`
pub fn new(mut connection_type: ConnectionType, read_buffer_size: usize) -> Self {
// Check if we need to use TLS
let tls_stream = match connection_type {
ConnectionType::TLS(ref mut tls_stream) => {
let stream = tls_stream.take().unwrap();
Some(stream)
},
_ => None,
};
Connection{
connection_type,
read_buffer_size,
tls_stream,
}
}
/// # Read To String
///
/// Read the `TcpStream` to a `String`
pub async fn read_to_string(&mut self) -> Result<String, &'static str> {
let mut string = unsafe {
String::from_utf8_unchecked(
self.read_to_vec()
.await
.expect("Error reading vec from stream"),
)
};
//.expect("Error decoding stream to utf-8");
trim_newline(&mut string);
Ok(string)
}
/// # Read To Vec
///
/// Read the `TcpStream` to a `Vec<u8>`. Returns a `Result` as we cannot guarantee a successful read.
async fn read_to_vec(&mut self) -> Result<Vec<u8>, Box<dyn Error>> {
let mut buffer = Vec::with_capacity(self.read_buffer_size);
// We need to check if we are using TLS or not, as we need to read differently
// (TLS uses a different read function, and does not support try_read)
match self.connection_type {
ConnectionType::Plain(ref mut stream) => {
let stream = stream.as_mut().unwrap();
// We loop while we're waiting for a read
loop {
stream.readable().await?; // we await for the stream to be readable
// Try to read data, this may still fail with `WouldBlock`
// if the readiness event is a false positive.
match stream.try_read_buf(&mut buffer) {
Ok(0) => break, // No data recieved
Ok(_) => {
break; // we recieved some data, just break the loop and return it
}
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
continue; // The IO is busy and would block - just try again
}
Err(e) => {
return Err(e.into()); // Some other error - quit
}
}
}
let buffer: Vec<u8> = buffer
.iter()
.filter(|x| **x != 0x0u8)
.map(|x| *x)
.collect(); // Remove all empty bytes (to avoid trailing whitespaces)
return Ok(buffer) // return our vector;
},
ConnectionType::TLS(_) => {
// Unfortunately, the TLS stream doesn't implement `AsyncRead`
// so we have to use the `split` combinator to split the stream
// This is a bit of a hack, but it works. We also lose
// async/await support, but that's ok for now.
let tls_stream = self.tls_stream.take().unwrap();
let (mut reader, writer) = split(tls_stream);
// Try to read data, this may still fail with `WouldBlock`
// if the readiness event is a false positive.
reader.read_buf(&mut buffer).await.unwrap();
let buffer: Vec<u8> = buffer
.iter()
.filter(|x| **x != 0x0u8)
.map(|x| *x)
.collect(); // Remove all empty bytes (to avoid trailing whitespaces)
// Reconnect the stream
self.tls_stream = Some(reader.unsplit(writer));
return Ok(buffer) // return our vector;
}
};
}
/// # Write String
///
/// Write a `String` value to the `TcpStream`. Returns a `Result` as we cannot guarantee a successful write.
pub async fn write_string(&mut self, data: String) -> Result<(), Box<dyn Error>> {
match self.connection_type{
ConnectionType::Plain(ref mut stream) => {
let stream = stream.as_mut().unwrap();
loop {
// Wait for the socket to be writable
stream.writable().await?;
// See `read_to_vec` for more explaination what happens here
//
// Try to write data, this may still fail with `WouldBlock`
// if the readiness event is a false positive.
match stream.try_write(data.as_bytes()) {
Ok(_) => {
break;
}
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
continue;
}
Err(e) => {
return Err(e.into());
}
}
}
},
ConnectionType::TLS(_) =>{
// Unfortunately, the TLS stream doesn't implement `AsyncWrite`
// so we have to use the `sink` combinator to split the stream
// This is a bit of a hack, but it works. We also lose
// async/await support, but that's ok for now.
let tls_stream = self.tls_stream.take().unwrap();
let (reader, mut writer) = split(tls_stream);
// Try to write data, this may still fail with `WouldBlock`
// if the readiness event is a false positive.
writer.write_all(data.as_bytes()).await?;
// Reconnect the stream
self.tls_stream = Some(reader.unsplit(writer));
}
}
Ok(())
}
/// # Write Bytes
///
/// Write bytes to the TCP stream, useful for sending data for things such as images or downloadable binary files
pub async fn write_bytes(&mut self, data: Vec<u8>) -> Result<(), Box<dyn Error>> {
match self.connection_type{
ConnectionType::Plain(ref mut stream) => {
let stream = stream.as_mut().unwrap();
loop {
// Wait for the socket to be writable
stream.writable().await?;
// See `read_to_vec` for more explaination what happens here
//
// Try to write data, this may still fail with `WouldBlock`
// if the readiness event is a false positive.
match stream.try_write(&data) {
Ok(_) => {
break;
}
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
continue;
}
Err(e) => {
return Err(e.into());
}
}
}
},
ConnectionType::TLS(_) => {
// Unfortunately, the TLS stream doesn't implement `AsyncWrite`
// so we have to use the `sink` combinator to split the stream
// This is a bit of a hack, but it works. We also lose
// async/await support, but that's ok for now.
let tls_stream = self.tls_stream.take().unwrap();
let (reader, mut writer) = split(tls_stream);
// Try to write data, this may still fail with `WouldBlock`
// if the readiness event is a false positive.
writer.write_all(&data).await?;
// Reconnect the stream
self.tls_stream = Some(reader.unsplit(writer));
}
};
Ok(())
}
}
/// Trim the ends of the `String` we got from the `TcpStream` so we don't waste buffer space with whitespace
fn trim_newline(s: &mut String) {
if s.ends_with('\n') {
s.pop();
if s.ends_with('\r') {
s.pop();
}
}
}