timeout and proper chunk handling

This commit is contained in:
Alex Janka 2024-03-07 09:24:20 +11:00
parent 42b742feb1
commit 401fd144cd

View file

@ -1,4 +1,8 @@
use std::{collections::HashMap, time::Duration};
use std::{
collections::{HashMap, VecDeque},
mem,
time::Duration,
};
use chacha20poly1305::{
aead::generic_array::GenericArray, AeadInPlace, ChaCha20Poly1305, KeyInit, Nonce,
@ -228,30 +232,30 @@ impl AccessorySocket {
{
if transfer_encoding.value == b"chunked" {
loop {
packet.append(&mut self.get_next().await?);
let chunked = ChunkedTransfer::clone_from(&packet).collect::<Vec<Vec<u8>>>();
let utf8_decoded = String::from_utf8(packet.clone())?;
if let Some(last) = chunked.chunks(2).last() {
if last.len() == 2 && last[0] == b"0" {
break;
}
}
let split = utf8_decoded
.split_terminator("\r\n")
.map(String::from)
.collect::<Vec<_>>();
if let Some(last) = split.chunks(2).last() {
if last.len() == 2 && last[0] == "0" {
match tokio::time::timeout(Duration::from_secs(2), self.get_next()).await {
Ok(next) => {
packet.append(&mut next?);
}
Err(_) => {
log::error!("timed out");
break;
}
}
}
let utf8_decoded = String::from_utf8(std::mem::take(&mut packet))?;
let mut chunked =
ChunkedTransfer::from(std::mem::take(&mut packet)).collect::<Vec<Vec<u8>>>();
let split = utf8_decoded
.split_terminator("\r\n")
.map(String::from)
.collect::<Vec<_>>();
for chunk in split.chunks_exact(2) {
packet.extend_from_slice(chunk[1].as_bytes())
for chunk in chunked.chunks_exact_mut(2) {
packet.append(&mut chunk[1]);
}
}
}
@ -329,3 +333,52 @@ pub enum DiscoveryError {
#[error("not found")]
NotFound,
}
struct ChunkedTransfer {
data: Vec<u8>,
}
impl From<Vec<u8>> for ChunkedTransfer {
fn from(value: Vec<u8>) -> Self {
Self { data: value }
}
}
impl ChunkedTransfer {
fn clone_from(value: &[u8]) -> Self {
Self {
data: value.to_owned(),
}
}
}
impl Iterator for ChunkedTransfer {
type Item = Vec<u8>;
fn next(&mut self) -> Option<Self::Item> {
if self.data.is_empty() {
None
} else {
let mut crlf_index = None;
for (i, val) in self.data.iter().enumerate() {
if let Some(next) = self.data.get(i + 1) {
if *val == b'\r' && *next == b'\n' {
crlf_index = Some(i);
break;
}
}
}
if let Some(i) = crlf_index {
let remainder = self.data.split_off(i);
let mut remainder = VecDeque::from(remainder);
let _ = remainder.pop_front().map(|v| v as char);
let _ = remainder.pop_front().map(|v| v as char);
Some(mem::replace(&mut self.data, Vec::from(remainder)))
} else {
Some(mem::take(&mut self.data))
}
}
}
}