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::{ use chacha20poly1305::{
aead::generic_array::GenericArray, AeadInPlace, ChaCha20Poly1305, KeyInit, Nonce, aead::generic_array::GenericArray, AeadInPlace, ChaCha20Poly1305, KeyInit, Nonce,
@ -228,30 +232,30 @@ impl AccessorySocket {
{ {
if transfer_encoding.value == b"chunked" { if transfer_encoding.value == b"chunked" {
loop { 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 match tokio::time::timeout(Duration::from_secs(2), self.get_next()).await {
.split_terminator("\r\n") Ok(next) => {
.map(String::from) packet.append(&mut next?);
.collect::<Vec<_>>(); }
if let Some(last) = split.chunks(2).last() { Err(_) => {
if last.len() == 2 && last[0] == "0" { log::error!("timed out");
break; 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 for chunk in chunked.chunks_exact_mut(2) {
.split_terminator("\r\n") packet.append(&mut chunk[1]);
.map(String::from)
.collect::<Vec<_>>();
for chunk in split.chunks_exact(2) {
packet.extend_from_slice(chunk[1].as_bytes())
} }
} }
} }
@ -329,3 +333,52 @@ pub enum DiscoveryError {
#[error("not found")] #[error("not found")]
NotFound, 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))
}
}
}
}