|
|
@@ -1,3 +1,5 @@
|
|
|
1
|
+use std::pin::Pin;
|
|
|
2
|
+
|
|
1
|
3
|
use async_tungstenite::tungstenite::{Error as WsError, Message};
|
|
2
|
4
|
use async_tungstenite::WebSocketStream;
|
|
3
|
5
|
use futures::prelude::*;
|
|
|
@@ -5,17 +7,17 @@ use futures::sink::Sink;
|
|
5
|
7
|
use futures::stream::Stream;
|
|
6
|
8
|
use futures::task::Context;
|
|
7
|
9
|
use futures::task::Poll;
|
|
8
|
|
-use std::pin::Pin;
|
|
|
10
|
+use pin_project::pin_project;
|
|
9
|
11
|
|
|
|
12
|
+#[pin_project]
|
|
10
|
13
|
pub struct WebSocketConnection<S> {
|
|
11
|
|
- stream: Pin<Box<WebSocketStream<S>>>,
|
|
|
14
|
+ #[pin]
|
|
|
15
|
+ stream: WebSocketStream<S>,
|
|
12
|
16
|
}
|
|
13
|
17
|
|
|
14
|
18
|
impl<S> WebSocketConnection<S> {
|
|
15
|
|
- pub fn new(s: WebSocketStream<S>) -> Self {
|
|
16
|
|
- Self {
|
|
17
|
|
- stream: Box::pin(s),
|
|
18
|
|
- }
|
|
|
19
|
+ pub fn new(stream: WebSocketStream<S>) -> Self {
|
|
|
20
|
+ Self { stream }
|
|
19
|
21
|
}
|
|
20
|
22
|
}
|
|
21
|
23
|
|
|
|
@@ -26,31 +28,19 @@ where
|
|
26
|
28
|
type Error = WsError;
|
|
27
|
29
|
|
|
28
|
30
|
fn poll_ready(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
|
|
29
|
|
- // Safe, as we will not move self_.
|
|
30
|
|
- let self_ = unsafe { self.get_unchecked_mut() };
|
|
31
|
|
-
|
|
32
|
|
- Pin::as_mut(&mut self_.stream).poll_ready(cx)
|
|
|
31
|
+ self.project().stream.poll_ready(cx)
|
|
33
|
32
|
}
|
|
34
|
33
|
|
|
35
|
34
|
fn start_send(self: Pin<&mut Self>, item: Vec<u8>) -> Result<(), Self::Error> {
|
|
36
|
|
- // Safe, as we will not move self_.
|
|
37
|
|
- let self_ = unsafe { self.get_unchecked_mut() };
|
|
38
|
|
-
|
|
39
|
|
- Pin::as_mut(&mut self_.stream).start_send(Message::Binary(item))
|
|
|
35
|
+ self.project().stream.start_send(Message::Binary(item))
|
|
40
|
36
|
}
|
|
41
|
37
|
|
|
42
|
38
|
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
|
|
43
|
|
- // Safe, as we will not move self_.
|
|
44
|
|
- let self_ = unsafe { self.get_unchecked_mut() };
|
|
45
|
|
-
|
|
46
|
|
- Pin::as_mut(&mut self_.stream).poll_flush(cx)
|
|
|
39
|
+ self.project().stream.poll_flush(cx)
|
|
47
|
40
|
}
|
|
48
|
41
|
|
|
49
|
42
|
fn poll_close(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
|
|
50
|
|
- // Safe, as we will not move self_.
|
|
51
|
|
- let self_ = unsafe { self.get_unchecked_mut() };
|
|
52
|
|
-
|
|
53
|
|
- Pin::as_mut(&mut self_.stream).poll_close(cx)
|
|
|
43
|
+ self.project().stream.poll_close(cx)
|
|
54
|
44
|
}
|
|
55
|
45
|
}
|
|
56
|
46
|
|
|
|
@@ -64,18 +54,17 @@ where
|
|
64
|
54
|
self: Pin<&mut Self>,
|
|
65
|
55
|
cx: &mut Context,
|
|
66
|
56
|
) -> Poll<Option<<WebSocketConnection<S> as Stream>::Item>> {
|
|
67
|
|
- // Safe, as we will not move self_.
|
|
68
|
|
- let self_ = unsafe { self.get_unchecked_mut() };
|
|
69
|
|
-
|
|
70
|
|
- match Pin::as_mut(&mut self_.stream).poll_next(cx) {
|
|
71
|
|
- Poll::Ready(Some(Ok(Message::Binary(data)))) => Poll::Ready(Some(Ok(data))),
|
|
72
|
|
- Poll::Ready(Some(Ok(_))) => unsafe {
|
|
73
|
|
- // Ignore other message types, try again.
|
|
74
|
|
- Pin::new_unchecked(self_).poll_next(cx)
|
|
75
|
|
- },
|
|
76
|
|
- Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
|
|
77
|
|
- Poll::Ready(None) => Poll::Ready(None),
|
|
78
|
|
- Poll::Pending => Poll::Pending,
|
|
|
57
|
+ let mut this = self.project();
|
|
|
58
|
+ loop {
|
|
|
59
|
+ match this.stream.as_mut().poll_next(cx) {
|
|
|
60
|
+ Poll::Ready(Some(Ok(Message::Binary(data)))) => return Poll::Ready(Some(Ok(data))),
|
|
|
61
|
+ Poll::Ready(Some(Ok(_))) => {
|
|
|
62
|
+ // Ignore other message types, try again.
|
|
|
63
|
+ }
|
|
|
64
|
+ Poll::Ready(Some(Err(e))) => return Poll::Ready(Some(Err(e))),
|
|
|
65
|
+ Poll::Ready(None) => return Poll::Ready(None),
|
|
|
66
|
+ Poll::Pending => return Poll::Pending,
|
|
|
67
|
+ }
|
|
79
|
68
|
}
|
|
80
|
69
|
}
|
|
81
|
70
|
}
|