uniauth/
util.rs

1//! Utility functions and traits
2
3use std::{
4	io::{Error, ErrorKind, Result},
5	marker::{Send, Unpin}
6};
7
8use async_trait::async_trait;
9pub use byteorder::{LittleEndian as LE, ReadBytesExt, WriteBytesExt};
10use tokio::io::{AsyncReadExt, AsyncWriteExt};
11
12/// Helper function to make a packet vec
13#[inline]
14pub fn make_packet<F>(n: usize, f: F) -> Result<Vec<u8>>
15	where F: Fn(&mut Vec<u8>) -> Result<()>
16{
17	let mut packet = Vec::with_capacity(n);
18	f(&mut packet)?;
19	Ok(packet)
20}
21
22pub trait WriteExtraExt : WriteBytesExt {
23	/// Writes a utf-8 string prepended by a byte with the length, in bytes, of the string
24	fn write_str(&mut self, s: &str) -> Result<()> {
25		self.write_b8(s.as_bytes())
26	}
27	/// Writes a Vec of bytes prepended with 1 byte length
28	fn write_b8(&mut self, bytes: &[u8]) -> Result<()>;
29	/// Writes a Vec of bytes prepended with 2 bytes length
30	fn write_b16(&mut self, bytes: &[u8]) -> Result<()>;
31	/// Writes a Vec of bytes prepended with 4 bytes length
32	fn write_b32(&mut self, bytes: &[u8], max: usize) -> Result<()>;
33}
34
35impl <W: WriteBytesExt> WriteExtraExt for W {
36	fn write_b8(&mut self, bytes: &[u8]) -> Result<()> {
37		if bytes.len() > u8::MAX.into() {
38			return Err(Error::new(ErrorKind::InvalidInput, "Data too long to read back"));
39		}
40
41		self.write_u8(bytes.len() as u8)?;
42		self.write_all(bytes)
43	}
44
45	fn write_b16(&mut self, bytes: &[u8]) -> Result<()> {
46		if bytes.len() > u16::MAX as usize {
47			return Err(Error::new(ErrorKind::InvalidInput, "Data too long to read back"));
48		}
49
50		self.write_u16::<LE>(bytes.len() as u16)?;
51		self.write_all(bytes)
52	}
53
54	fn write_b32(&mut self, bytes: &[u8], max: usize) -> Result<()> {
55		if bytes.len() > max {
56			return Err(Error::new(ErrorKind::InvalidInput, "Data too long to read back"));
57		}
58
59		self.write_u32::<LE>(bytes.len() as u32)?;
60		self.write_all(bytes)
61	}
62}
63
64#[async_trait]
65pub trait AsyncWriteExtraExt : AsyncWriteExt {
66	/// Writes a utf-8 string prepended by a byte with the length, in bytes, of the string
67	async fn write_str(&mut self, s: &str) -> Result<()> {
68		self.write_b8(s.as_bytes()).await
69	}
70
71	/// Writes a Vec of bytes prepended with 1 byte length
72	async fn write_b8(&mut self, bytes: &[u8]) -> Result<()>;
73	/// Writes a Vec of bytes prepended with 2 bytes length
74	async fn write_b16(&mut self, bytes: &[u8]) -> Result<()>;
75	/// Writes a Vec of bytes prepended with 4 bytes length
76	async fn write_b32(&mut self, bytes: &[u8], max: usize) -> Result<()>;
77}
78
79#[async_trait]
80impl <W: AsyncWriteExt + Send + Unpin> AsyncWriteExtraExt for W {
81	async fn write_b8(&mut self, bytes: &[u8]) -> Result<()> {
82		if bytes.len() > u8::MAX.into() {
83			return Err(Error::new(ErrorKind::InvalidInput, "Data too long to read back"));
84		}
85
86		self.write_u8(bytes.len() as u8).await?;
87		self.write_all(bytes).await
88	}
89
90	async fn write_b16(&mut self, bytes: &[u8]) -> Result<()> {
91		if bytes.len() > u16::MAX as usize {
92			return Err(Error::new(ErrorKind::InvalidInput, "Data too long to read back"));
93		}
94
95		self.write_u16_le(bytes.len() as u16).await?;
96		self.write_all(bytes).await
97	}
98
99	async fn write_b32(&mut self, bytes: &[u8], max: usize) -> Result<()> {
100		if bytes.len() > max {
101			return Err(Error::new(ErrorKind::InvalidInput, "Data too long to read back"));
102		}
103
104		self.write_u32_le(bytes.len() as u32).await?;
105		self.write_all(bytes).await
106	}
107}
108
109pub trait ReadExtraExt : ReadBytesExt {
110	/// Reads a 0-255 byte utf-8 encoded string
111	fn read_str(&mut self) -> Result<String> {
112		let bytes = self.read_b8()?;
113		String::from_utf8(bytes)
114			.map_err(|_| Error::new(ErrorKind::InvalidData, "String contains non-utf8 bytes"))
115	}
116
117	/// Reads a Vec of bytes prepended with 1 byte length
118	fn read_b8(&mut self) -> Result<Vec<u8>>;
119	/// Reads a Vec of bytes prepended with 2 bytes length
120	fn read_b16(&mut self) -> Result<Vec<u8>>;
121	/// Reads a Vec of bytes prepended with 4 bytes length
122	fn read_b32(&mut self, max: usize) -> Result<Vec<u8>>;
123}
124
125impl <R: ReadBytesExt> ReadExtraExt for R {
126	fn read_b8(&mut self) -> Result<Vec<u8>> {
127		let len = self.read_u8()? as usize;
128		let mut bytes = vec![0; len];
129		self.read_exact(&mut bytes)?;
130		Ok(bytes)
131	}
132
133	fn read_b16(&mut self) -> Result<Vec<u8>> {
134		let len = self.read_u16::<LE>()? as usize;
135		let mut bytes = vec![0; len];
136		self.read_exact(&mut bytes)?;
137		Ok(bytes)
138	}
139
140	fn read_b32(&mut self, max: usize) -> Result<Vec<u8>> {
141		let len = self.read_u32::<LE>()? as usize;
142		if len > max {
143			return Err(Error::new(ErrorKind::InvalidData, "Data too long"));
144		}
145
146		let mut bytes = vec![0; len];
147		self.read_exact(&mut bytes)?;
148		Ok(bytes)
149	}
150}
151
152#[async_trait]
153pub trait AsyncReadExtraExt : AsyncReadExt {
154	/// Reads a 0-255 byte utf-8 encoded string
155	async fn read_str(&mut self) -> Result<String> {
156		let bytes = self.read_b8().await?;
157		String::from_utf8(bytes)
158			.map_err(|_| Error::new(ErrorKind::InvalidData, "String contains non-utf8 bytes"))
159	}
160
161	/// Reads a Vec of bytes prepended with 1 byte length
162	async fn read_b8(&mut self) -> Result<Vec<u8>>;
163	/// Reads a Vec of bytes prepended with 2 bytes length
164	async fn read_b16(&mut self) -> Result<Vec<u8>>;
165	/// Reads a Vec of bytes prepended with 4 bytes length
166	async fn read_b32(&mut self, max: usize) -> Result<Vec<u8>>;
167}
168
169#[async_trait]
170impl <R: AsyncReadExt + Send + Unpin> AsyncReadExtraExt for R {
171	async fn read_b8(&mut self) -> Result<Vec<u8>> {
172		let len = self.read_u8().await? as usize;
173		let mut bytes = vec![0; len];
174		self.read_exact(&mut bytes).await?;
175		Ok(bytes)
176	}
177
178	async fn read_b16(&mut self) -> Result<Vec<u8>> {
179		let len = self.read_u16_le().await? as usize;
180		let mut bytes = vec![0; len];
181		self.read_exact(&mut bytes).await?;
182		Ok(bytes)
183	}
184
185	async fn read_b32(&mut self, max: usize) -> Result<Vec<u8>> {
186		let len = self.read_u32_le().await? as usize;
187		if len > max {
188			return Err(Error::new(ErrorKind::InvalidData, "Data too long"));
189		}
190
191		let mut bytes = vec![0; len];
192		self.read_exact(&mut bytes).await?;
193		Ok(bytes)
194	}
195}
196
uniauth/util.rs

uniauth/
util.rs
