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