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::{
|
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))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
Loading…
Reference in a new issue