proxmox-mirrors/proxmox
1
0
Fork 1
mirror of https://git.proxmox.com/git/proxmox synced 2025-08-06 07:56:12 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

tools/websocket: require writer to be Unpin

Browse Source 
so we can safely pin/unpin it

Signed-off-by: Dominik Csapak <d.csapak@proxmox.com>
This commit is contained in:
Dominik Csapak 2020-07-03 11:56:38 +02:00 committed by Dietmar Maurer
parent 556df80841
commit 521b26d6a2
1 changed files with 9 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
proxmox/src/tools/websocket.rs
View File

@ -106,14 +106,14 @@ pub fn create_frame(
Ok(buf)
}
pub struct WebSocketWriter<W: AsyncWrite> {
pub struct WebSocketWriter<W: AsyncWrite + Unpin> {
writer: W,
text: bool,
mask: Option<[u8; 4]>,
frame: Option<(Vec<u8>, usize, usize)>,
}
impl<W: AsyncWrite> WebSocketWriter<W> {
impl<W: AsyncWrite + Unpin> WebSocketWriter<W> {
pub fn new(mask: Option<[u8; 4]>, text: bool, writer: W) -> WebSocketWriter<W> {
WebSocketWriter {
writer: writer,
@ -124,13 +124,13 @@ impl<W: AsyncWrite> WebSocketWriter<W> {
}
}
impl<W: AsyncWrite> AsyncWrite for WebSocketWriter<W> {
impl<W: AsyncWrite + Unpin> AsyncWrite for WebSocketWriter<W> {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context,
buf: &[u8],
) -> Poll<Result<usize, Error>> {
let this = unsafe { Pin::into_inner_unchecked(self) };
let this = Pin::get_mut(self);
let frametype = match this.text {
true => OpCode::Text,
@ -151,8 +151,7 @@ impl<W: AsyncWrite> AsyncWrite for WebSocketWriter<W> {
// we have a frame in any case, so unwrap is ok
let (buf, pos, origsize) = this.frame.as_mut().unwrap();
loop {
//let size = unsafe { Pin::new_unchecked(&mut this.writer) }.poll_write(cx, &buf[*pos..]))?
match unsafe { Pin::new_unchecked(&mut this.writer) }.poll_write(cx, &buf[*pos..]) {
match Pin::new(&mut this.writer).poll_write(cx, &buf[*pos..]) {
Poll::Ready(Ok(size)) => {
*pos += size;
if *pos == buf.len() {
@ -167,11 +166,13 @@ impl<W: AsyncWrite> AsyncWrite for WebSocketWriter<W> {
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Error>> {
unsafe { self.map_unchecked_mut(|x| &mut x.writer).poll_flush(cx) }
let this = Pin::get_mut(self);
Pin::new(&mut this.writer).poll_flush(cx)
}
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Error>> {
unsafe { self.map_unchecked_mut(|x| &mut x.writer).poll_shutdown(cx) }
let this = Pin::get_mut(self);
Pin::new(&mut this.writer).poll_shutdown(cx)
}
}