timeout and proper chunk handling
This commit is contained in:
parent
42b742feb1
commit
401fd144cd
|
@ -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))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Reference in a new issue