use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicU8, AtomicUsize, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, ReadBuf};
pub const MAX_PACKET_BYTES: usize = 67_108_864;
pub const FULL_PACKET_PAYLOAD: usize = 0x00FF_FFFF;
#[derive(Debug, Default)]
pub struct PacketLimitTrip {
tripped: AtomicBool,
sequence: AtomicU8,
declared: AtomicUsize,
}
impl PacketLimitTrip {
pub fn tripped(&self) -> bool {
self.tripped.load(Ordering::SeqCst)
}
pub fn reply_sequence(&self) -> u8 {
self.sequence.load(Ordering::SeqCst).wrapping_add(1)
}
pub fn declared_bytes(&self) -> usize {
self.declared.load(Ordering::SeqCst)
}
fn record(&self, sequence: u8, declared: usize) {
self.sequence.store(sequence, Ordering::SeqCst);
self.declared.store(declared, Ordering::SeqCst);
self.tripped.store(true, Ordering::SeqCst);
}
}
pub struct PacketLimitReader<R> {
inner: R,
limit: usize,
trip: Arc<PacketLimitTrip>,
header: [u8; 4],
header_len: usize,
payload_remaining: usize,
logical_declared: usize,
}
impl<R> PacketLimitReader<R> {
pub fn new(inner: R, trip: Arc<PacketLimitTrip>) -> Self {
Self::with_limit(inner, MAX_PACKET_BYTES, trip)
}
pub fn with_limit(inner: R, limit: usize, trip: Arc<PacketLimitTrip>) -> Self {
Self {
inner,
limit,
trip,
header: [0u8; 4],
header_len: 0,
payload_remaining: 0,
logical_declared: 0,
}
}
pub fn into_inner(self) -> R {
self.inner
}
fn observe(&mut self, bytes: &[u8]) -> Result<(), ()> {
let mut i = 0usize;
while i < bytes.len() {
if self.payload_remaining > 0 {
let take = self.payload_remaining.min(bytes.len() - i);
self.payload_remaining -= take;
i += take;
continue;
}
let need = 4 - self.header_len;
let take = need.min(bytes.len() - i);
self.header[self.header_len..self.header_len + take]
.copy_from_slice(&bytes[i..i + take]);
self.header_len += take;
i += take;
if self.header_len < 4 {
break;
}
self.header_len = 0;
let declared =
u32::from_le_bytes([self.header[0], self.header[1], self.header[2], 0]) as usize;
let sequence = self.header[3];
let logical = self.logical_declared.saturating_add(declared);
if logical > self.limit {
self.trip.record(sequence, logical);
return Err(());
}
self.payload_remaining = declared;
self.logical_declared = if declared == FULL_PACKET_PAYLOAD {
logical
} else {
0
};
}
Ok(())
}
}
impl<R: AsyncRead + Unpin> AsyncRead for PacketLimitReader<R> {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
let this = self.get_mut();
let before = buf.filled().len();
match Pin::new(&mut this.inner).poll_read(cx, buf) {
Poll::Ready(Ok(())) => {
let filled = buf.filled().len();
if filled > before {
let fresh = buf.filled()[before..filled].to_vec();
if this.observe(&fresh).is_err() {
return Poll::Ready(Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"MySQL packet exceeds max_allowed_packet",
)));
}
}
Poll::Ready(Ok(()))
}
other => other,
}
}
}
pub fn packet_too_large_err(sequence: u8) -> Vec<u8> {
const MESSAGE: &[u8] = b"Got a packet bigger than 'max_allowed_packet' bytes";
const ER_NET_PACKET_TOO_LARGE: u16 = 1153;
let mut payload = Vec::with_capacity(9 + MESSAGE.len());
payload.push(0xFF);
payload.extend_from_slice(&ER_NET_PACKET_TOO_LARGE.to_le_bytes());
payload.push(b'#');
payload.extend_from_slice(b"08S01");
payload.extend_from_slice(MESSAGE);
let len = payload.len();
let mut packet = Vec::with_capacity(4 + len);
packet.push((len & 0xFF) as u8);
packet.push(((len >> 8) & 0xFF) as u8);
packet.push(((len >> 16) & 0xFF) as u8);
packet.push(sequence);
packet.extend_from_slice(&payload);
packet
}