Line data Source code
1 : use crate::proxy::ErrorSource;
2 : use crate::{
3 : cancellation::CancellationHandlerMain,
4 : config::ProxyConfig,
5 : context::RequestMonitoring,
6 : error::{io_error, ReportableError},
7 : metrics::Metrics,
8 : proxy::{handle_client, ClientMode},
9 : rate_limiter::EndpointRateLimiter,
10 : };
11 : use anyhow::Context as _;
12 : use bytes::{Buf, BufMut, Bytes, BytesMut};
13 : use framed_websockets::{Frame, OpCode, WebSocketServer};
14 : use futures::{Sink, Stream};
15 : use hyper1::upgrade::OnUpgrade;
16 : use hyper_util::rt::TokioIo;
17 : use pin_project_lite::pin_project;
18 :
19 : use std::{
20 : pin::Pin,
21 : sync::Arc,
22 : task::{ready, Context, Poll},
23 : };
24 : use tokio::io::{self, AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf};
25 : use tracing::warn;
26 :
27 : pin_project! {
28 : /// This is a wrapper around a [`WebSocketStream`] that
29 : /// implements [`AsyncRead`] and [`AsyncWrite`].
30 : pub struct WebSocketRw<S> {
31 : #[pin]
32 : stream: WebSocketServer<S>,
33 : recv: Bytes,
34 : send: BytesMut,
35 : }
36 : }
37 :
38 : impl<S> WebSocketRw<S> {
39 2 : pub fn new(stream: WebSocketServer<S>) -> Self {
40 2 : Self {
41 2 : stream,
42 2 : recv: Bytes::new(),
43 2 : send: BytesMut::new(),
44 2 : }
45 2 : }
46 : }
47 :
48 : impl<S: AsyncRead + AsyncWrite + Unpin> AsyncWrite for WebSocketRw<S> {
49 2 : fn poll_write(
50 2 : self: Pin<&mut Self>,
51 2 : cx: &mut Context<'_>,
52 2 : buf: &[u8],
53 2 : ) -> Poll<io::Result<usize>> {
54 2 : let this = self.project();
55 2 : let mut stream = this.stream;
56 :
57 2 : ready!(stream.as_mut().poll_ready(cx).map_err(io_error))?;
58 :
59 2 : this.send.put(buf);
60 2 : match stream.as_mut().start_send(Frame::binary(this.send.split())) {
61 2 : Ok(()) => Poll::Ready(Ok(buf.len())),
62 0 : Err(e) => Poll::Ready(Err(io_error(e))),
63 : }
64 2 : }
65 :
66 2 : fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
67 2 : let stream = self.project().stream;
68 2 : stream.poll_flush(cx).map_err(io_error)
69 2 : }
70 :
71 0 : fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
72 0 : let stream = self.project().stream;
73 0 : stream.poll_close(cx).map_err(io_error)
74 0 : }
75 : }
76 :
77 : impl<S: AsyncRead + AsyncWrite + Unpin> AsyncRead for WebSocketRw<S> {
78 6 : fn poll_read(
79 6 : mut self: Pin<&mut Self>,
80 6 : cx: &mut Context<'_>,
81 6 : buf: &mut ReadBuf<'_>,
82 6 : ) -> Poll<io::Result<()>> {
83 6 : let bytes = ready!(self.as_mut().poll_fill_buf(cx))?;
84 4 : let len = std::cmp::min(bytes.len(), buf.remaining());
85 4 : buf.put_slice(&bytes[..len]);
86 4 : self.consume(len);
87 4 : Poll::Ready(Ok(()))
88 6 : }
89 : }
90 :
91 : impl<S: AsyncRead + AsyncWrite + Unpin> AsyncBufRead for WebSocketRw<S> {
92 6 : fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
93 6 : // Please refer to poll_fill_buf's documentation.
94 6 : const EOF: Poll<io::Result<&[u8]>> = Poll::Ready(Ok(&[]));
95 6 :
96 6 : let mut this = self.project();
97 8 : loop {
98 8 : if !this.recv.chunk().is_empty() {
99 2 : let chunk = (*this.recv).chunk();
100 2 : return Poll::Ready(Ok(chunk));
101 6 : }
102 :
103 6 : let res = ready!(this.stream.as_mut().poll_next(cx));
104 4 : match res.transpose().map_err(io_error)? {
105 4 : Some(message) => match message.opcode {
106 0 : OpCode::Ping => {}
107 0 : OpCode::Pong => {}
108 : OpCode::Text => {
109 : // We expect to see only binary messages.
110 0 : let error = "unexpected text message in the websocket";
111 0 : warn!(length = message.payload.len(), error);
112 0 : return Poll::Ready(Err(io_error(error)));
113 : }
114 : OpCode::Binary | OpCode::Continuation => {
115 2 : debug_assert!(this.recv.is_empty());
116 2 : *this.recv = message.payload.freeze();
117 : }
118 2 : OpCode::Close => return EOF,
119 : },
120 0 : None => return EOF,
121 : }
122 : }
123 6 : }
124 :
125 4 : fn consume(self: Pin<&mut Self>, amount: usize) {
126 4 : self.project().recv.advance(amount);
127 4 : }
128 : }
129 :
130 0 : pub async fn serve_websocket(
131 0 : config: &'static ProxyConfig,
132 0 : mut ctx: RequestMonitoring,
133 0 : websocket: OnUpgrade,
134 0 : cancellation_handler: Arc<CancellationHandlerMain>,
135 0 : endpoint_rate_limiter: Arc<EndpointRateLimiter>,
136 0 : hostname: Option<String>,
137 0 : ) -> anyhow::Result<()> {
138 0 : let websocket = websocket.await?;
139 0 : let websocket = WebSocketServer::after_handshake(TokioIo::new(websocket));
140 0 :
141 0 : let conn_gauge = Metrics::get()
142 0 : .proxy
143 0 : .client_connections
144 0 : .guard(crate::metrics::Protocol::Ws);
145 :
146 0 : let res = Box::pin(handle_client(
147 0 : config,
148 0 : &mut ctx,
149 0 : cancellation_handler,
150 0 : WebSocketRw::new(websocket),
151 0 : ClientMode::Websockets { hostname },
152 0 : endpoint_rate_limiter,
153 0 : conn_gauge,
154 0 : ))
155 0 : .await;
156 :
157 0 : match res {
158 0 : Err(e) => {
159 0 : // todo: log and push to ctx the error kind
160 0 : ctx.set_error_kind(e.get_error_kind());
161 0 : Err(e.into())
162 : }
163 : Ok(None) => {
164 0 : ctx.set_success();
165 0 : Ok(())
166 : }
167 0 : Ok(Some(p)) => {
168 0 : ctx.set_success();
169 0 : ctx.log_connect();
170 0 : match p.proxy_pass().await {
171 0 : Ok(()) => Ok(()),
172 0 : Err(ErrorSource::Client(err)) => Err(err).context("client"),
173 0 : Err(ErrorSource::Compute(err)) => Err(err).context("compute"),
174 : }
175 : }
176 : }
177 0 : }
178 :
179 : #[cfg(test)]
180 : mod tests {
181 : use std::pin::pin;
182 :
183 : use framed_websockets::WebSocketServer;
184 : use futures::{SinkExt, StreamExt};
185 : use tokio::{
186 : io::{duplex, AsyncReadExt, AsyncWriteExt},
187 : task::JoinSet,
188 : };
189 : use tokio_tungstenite::{
190 : tungstenite::{protocol::Role, Message},
191 : WebSocketStream,
192 : };
193 :
194 : use super::WebSocketRw;
195 :
196 : #[tokio::test]
197 2 : async fn websocket_stream_wrapper_happy_path() {
198 2 : let (stream1, stream2) = duplex(1024);
199 2 :
200 2 : let mut js = JoinSet::new();
201 2 :
202 2 : js.spawn(async move {
203 2 : let mut client = WebSocketStream::from_raw_socket(stream1, Role::Client, None).await;
204 2 :
205 2 : client
206 2 : .send(Message::Binary(b"hello world".to_vec()))
207 2 : .await
208 2 : .unwrap();
209 2 :
210 2 : let message = client.next().await.unwrap().unwrap();
211 2 : assert_eq!(message, Message::Binary(b"websockets are cool".to_vec()));
212 2 :
213 2 : client.close(None).await.unwrap();
214 2 : });
215 2 :
216 2 : js.spawn(async move {
217 2 : let mut rw = pin!(WebSocketRw::new(WebSocketServer::after_handshake(stream2)));
218 2 :
219 2 : let mut buf = vec![0; 1024];
220 2 : let n = rw.read(&mut buf).await.unwrap();
221 2 : assert_eq!(&buf[..n], b"hello world");
222 2 :
223 2 : rw.write_all(b"websockets are cool").await.unwrap();
224 2 : rw.flush().await.unwrap();
225 2 :
226 2 : let n = rw.read_to_end(&mut buf).await.unwrap();
227 2 : assert_eq!(n, 0);
228 2 : });
229 2 :
230 2 : js.join_next().await.unwrap().unwrap();
231 2 : js.join_next().await.unwrap().unwrap();
232 2 : }
233 : }
|
