Git-Mirrors/veilid
1
0
Fork 0
mirror of https://gitlab.com/veilid/veilid.git synced 2024-09-20 08:06:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix asyncpeekstream

Browse Source
This commit is contained in:
John Smith 2022-08-24 13:49:05 -04:00
parent cd99a4210b
commit ce19a1bfbf
4 changed files with 76 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
veilid-core/src/network_manager/native/protocol/tcp.rs
View File

@ -68,7 +68,6 @@ impl RawTcpNetworkConnection {
let mut header = [0u8; 4];
network_result_try!(stream.read_exact(&mut header).await.into_network_result()?);
if header[0] != b'V' || header[1] != b'L' {
bail_io_error_other!("received invalid TCP frame header");
}
@ -78,7 +77,8 @@ impl RawTcpNetworkConnection {
}
let mut out: Vec<u8> = vec![0u8; len];
network_result_try!(stream.read_exact(&mut out).await.into_network_result()?);
let nrout = stream.read_exact(&mut out).await.into_network_result()?;
network_result_try!(nrout);
Ok(NetworkResult::Value(out))
}

8
veilid-core/src/network_manager/native/protocol/ws.rs
View File

@ -81,6 +81,14 @@ where
Ok(v) => NetworkResult::value(v),
Err(e) => err_to_network_result(e),
};
if !out.is_value() {
tracing::Span::current().record("network_result", &tracing::field::display(&out));
return Ok(out);
}
let out = match self.stream.clone().flush().await {
Ok(v) => NetworkResult::value(v),
Err(e) => err_to_network_result(e),
};
tracing::Span::current().record("network_result", &tracing::field::display(&out));
Ok(out)
}

46
veilid-core/src/tests/native/test_async_peek_stream.rs
View File

@ -294,6 +294,51 @@ pub async fn test_peek_some_read_peek_some_read_all_read() {
assert_eq!(inbuf3, outbuf[peeksize1..outbuf.len()].to_vec());
}
pub async fn test_peek_exact_read_peek_exact_read_all_read() {
info!("test_peek_exact_read_peek_exact_read_all_read");
let (mut a, mut c) = make_async_peek_stream_loopback().await;
// write everything
let outbuf = MESSAGE.to_vec();
a.write_all(&outbuf).await.unwrap();
// peek partially
let mut peekbuf1: Vec<u8> = Vec::new();
peekbuf1.resize(outbuf.len() / 4, 0u8);
let peeksize1 = c.peek_exact(&mut peekbuf1).await.unwrap();
assert_eq!(peeksize1, peekbuf1.len());
// read partially
let mut inbuf1: Vec<u8> = Vec::new();
inbuf1.resize(peeksize1 - 1, 0u8);
c.read_exact(&mut inbuf1).await.unwrap();
// peek partially
let mut peekbuf2: Vec<u8> = Vec::new();
peekbuf2.resize(2, 0u8);
let peeksize2 = c.peek_exact(&mut peekbuf2).await.unwrap();
assert_eq!(peeksize2, peekbuf2.len());
// read partially
let mut inbuf2: Vec<u8> = Vec::new();
inbuf2.resize(1, 0u8);
c.read_exact(&mut inbuf2).await.unwrap();
// read remaining
let mut inbuf3: Vec<u8> = Vec::new();
inbuf3.resize(outbuf.len() - peeksize1, 0u8);
c.read_exact(&mut inbuf3).await.unwrap();
assert_eq!(peekbuf1, outbuf[0..peeksize1].to_vec());
assert_eq!(inbuf1, outbuf[0..peeksize1 - 1].to_vec());
assert_eq!(
peekbuf2[0..peeksize2].to_vec(),
outbuf[peeksize1 - 1..peeksize1 + 1].to_vec()
);
assert_eq!(inbuf2, peekbuf2[0..1].to_vec());
assert_eq!(inbuf3, outbuf[peeksize1..outbuf.len()].to_vec());
}
pub async fn test_all() {
test_nothing().await;
test_no_peek().await;
@ -303,4 +348,5 @@ pub async fn test_all() {
test_peek_some_read_peek_some_read().await;
test_peek_some_read_peek_all_read().await;
test_peek_some_read_peek_some_read_all_read().await;
test_peek_exact_read_peek_exact_read_all_read().await;
}

33
veilid-core/src/xx/async_peek_stream.rs
View File

@ -1,5 +1,4 @@
use super::*;
use futures_util::AsyncReadExt;
use std::io;
use task::{Context, Poll};
@ -32,10 +31,10 @@ impl Future for Peek<'_> {
if buf_len > inner.peekbuf_len {
//
inner.peekbuf.resize(buf_len, 0u8);
let mut read_future = inner
.stream
.read(&mut inner.peekbuf.as_mut_slice()[inner.peekbuf_len..buf_len]);
let read_len = match Pin::new(&mut read_future).poll(cx) {
let read_len = match Pin::new(&mut inner.stream).poll_read(
cx,
&mut inner.peekbuf.as_mut_slice()[inner.peekbuf_len..buf_len],
) {
Poll::Pending => {
inner.peekbuf.resize(inner.peekbuf_len, 0u8);
return Poll::Pending;
@ -60,6 +59,7 @@ impl Future for Peek<'_> {
pub struct PeekExact<'a> {
aps: AsyncPeekStream,
buf: &'a mut [u8],
cur_read: usize,
}
impl Unpin for PeekExact<'_> {}
@ -78,10 +78,10 @@ impl Future for PeekExact<'_> {
if buf_len > inner.peekbuf_len {
//
inner.peekbuf.resize(buf_len, 0u8);
let mut read_future = inner
.stream
.read_exact(&mut inner.peekbuf.as_mut_slice()[inner.peekbuf_len..buf_len]);
match Pin::new(&mut read_future).poll(cx) {
let read_len = match Pin::new(&mut inner.stream).poll_read(
cx,
&mut inner.peekbuf.as_mut_slice()[inner.peekbuf_len..buf_len],
) {
Poll::Pending => {
inner.peekbuf.resize(inner.peekbuf_len, 0u8);
return Poll::Pending;
@ -89,13 +89,19 @@ impl Future for PeekExact<'_> {
Poll::Ready(Err(e)) => {
return Poll::Ready(Err(e));
}
Poll::Ready(Ok(())) => (),
Poll::Ready(Ok(v)) => v,
};
inner.peekbuf_len = buf_len;
inner.peekbuf_len += read_len;
inner.peekbuf.resize(inner.peekbuf_len, 0u8);
copy_len = inner.peekbuf_len;
}
this.buf[..copy_len].copy_from_slice(&inner.peekbuf[..copy_len]);
Poll::Ready(Ok(copy_len))
this.buf[this.cur_read..copy_len].copy_from_slice(&inner.peekbuf[this.cur_read..copy_len]);
this.cur_read = copy_len;
if this.cur_read == buf_len {
Poll::Ready(Ok(buf_len))
} else {
Poll::Pending
}
}
}
/////////
@ -139,6 +145,7 @@ impl AsyncPeekStream {
PeekExact::<'a> {
aps: self.clone(),
buf,
cur_read: 0,
}
}
}