1use std::{
2 borrow::Cow,
3 fmt,
4 pin::Pin,
5 sync::Arc,
6 task::{Context, Poll},
7};
8
9use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64_STANDARD};
10use futures_util::{SinkExt, Stream, StreamExt, stream::FusedStream};
11use reqwest::{Method, StatusCode};
12use rustls::{
13 ClientConfig, DigitallySignedStruct, RootCertStore, SignatureScheme,
14 client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier},
15 pki_types::{CertificateDer, ServerName, UnixTime},
16};
17use tokio::{
18 io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt},
19 net::{TcpStream, lookup_host},
20 time::timeout,
21};
22use tokio_rustls::TlsConnector;
23use tokio_socks::{TargetAddr, tcp::Socks5Stream};
24use tokio_tungstenite::{
25 Connector, MaybeTlsStream, WebSocketStream, client_async_tls_with_config,
26 tungstenite::{
27 Message,
28 client::IntoClientRequest,
29 protocol::{WebSocketConfig, frame::CloseFrame, frame::coding::CloseCode},
30 },
31};
32use url::{Host, Url};
33
34use super::PanelEvent;
35use crate::{Client, Error, ProxyConfig, ProxyError, ProxyScheme, Result};
36
37const MAX_MESSAGE_SIZE: usize = 10 * 1024 * 1024;
38const CLIENT_CLOSE_REASON: &str = "xui-rs disconnect";
39const MAX_CONNECT_RESPONSE_SIZE: usize = 16 * 1024;
40
41trait AsyncIo: AsyncRead + AsyncWrite {}
42
43impl<T> AsyncIo for T where T: AsyncRead + AsyncWrite + ?Sized {}
44
45type BoxedIo = Box<dyn AsyncIo + Send + Unpin>;
46type Socket = WebSocketStream<MaybeTlsStream<BoxedIo>>;
47
48#[derive(Clone, Copy, Debug)]
50pub struct EventsApi<'client> {
51 client: &'client Client,
52}
53
54impl<'client> EventsApi<'client> {
55 pub(crate) const fn new(client: &'client Client) -> Self {
56 Self { client }
57 }
58
59 pub async fn connect(self) -> Result<EventStream> {
72 EventStream::connect(self.client.clone()).await
73 }
74}
75
76#[derive(Clone, Debug, Eq, PartialEq)]
78pub struct WebSocketClose {
79 pub code: u16,
81 pub reason: String,
83}
84
85pub struct EventStream {
91 client: Client,
92 endpoint: Url,
93 socket: Option<Socket>,
94 close: Option<WebSocketClose>,
95 terminated: bool,
96}
97
98impl EventStream {
99 async fn connect(client: Client) -> Result<Self> {
100 let (endpoint, socket) = connect_socket(&client).await?;
101 Ok(Self {
102 client,
103 endpoint,
104 socket: Some(socket),
105 close: None,
106 terminated: false,
107 })
108 }
109
110 pub async fn next_event(&mut self) -> Result<Option<PanelEvent>> {
121 self.next().await.transpose()
122 }
123
124 pub const fn close_info(&self) -> Option<&WebSocketClose> {
126 self.close.as_ref()
127 }
128
129 pub const fn is_closed(&self) -> bool {
131 self.terminated
132 }
133
134 pub async fn reconnect(&mut self) -> Result<()> {
146 if !self.terminated {
147 return Err(Error::Configuration(
148 "cannot reconnect an active WebSocket event stream".to_owned(),
149 ));
150 }
151 let (endpoint, socket) = connect_socket(&self.client).await?;
152 self.endpoint = endpoint;
153 self.socket = Some(socket);
154 self.close = None;
155 self.terminated = false;
156 Ok(())
157 }
158
159 pub async fn close(&mut self) -> Result<()> {
166 let Some(mut socket) = self.socket.take() else {
167 self.terminated = true;
168 return Ok(());
169 };
170 let close = WebSocketClose {
171 code: 1000,
172 reason: CLIENT_CLOSE_REASON.to_owned(),
173 };
174 let frame = CloseFrame {
175 code: CloseCode::Normal,
176 reason: CLIENT_CLOSE_REASON.into(),
177 };
178 let result = socket.send(Message::Close(Some(frame))).await;
179 self.close = Some(close);
180 self.terminated = true;
181 match result {
182 Ok(())
183 | Err(
184 tokio_tungstenite::tungstenite::Error::ConnectionClosed
185 | tokio_tungstenite::tungstenite::Error::AlreadyClosed,
186 ) => Ok(()),
187 Err(source) => Err(socket_error(self.endpoint.clone(), source)),
188 }
189 }
190}
191
192impl fmt::Debug for EventStream {
193 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
194 formatter
195 .debug_struct("EventStream")
196 .field("client", &"[REDACTED]")
197 .field("panel_origin", &origin(&self.endpoint))
198 .field("endpoint_path", &"[REDACTED]")
199 .field(
200 "socket",
201 &self.socket.as_ref().map(|_| "[CONNECTED; REDACTED]"),
202 )
203 .field("closed", &self.terminated)
204 .field("close", &self.close)
205 .finish()
206 }
207}
208
209impl Stream for EventStream {
210 type Item = Result<PanelEvent>;
211
212 fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
213 loop {
214 if self.terminated {
215 return Poll::Ready(None);
216 }
217 let poll = if let Some(socket) = self.socket.as_mut() {
218 Pin::new(socket).poll_next(context)
219 } else {
220 self.terminated = true;
221 return Poll::Ready(None);
222 };
223 match poll {
224 Poll::Pending => return Poll::Pending,
225 Poll::Ready(Some(Ok(Message::Text(text)))) => {
226 return Poll::Ready(Some(PanelEvent::decode(text.as_str())));
227 }
228 Poll::Ready(Some(Ok(Message::Ping(_) | Message::Pong(_)))) => {}
229 Poll::Ready(Some(Ok(Message::Close(frame)))) => {
230 self.close = Some(frame.map_or(
231 WebSocketClose {
232 code: 1005,
233 reason: String::new(),
234 },
235 |frame| WebSocketClose {
236 code: frame.code.into(),
237 reason: frame.reason.to_string(),
238 },
239 ));
240 self.socket = None;
241 self.terminated = true;
242 return Poll::Ready(None);
243 }
244 Poll::Ready(Some(Ok(Message::Binary(_)))) => {
245 self.socket = None;
246 self.terminated = true;
247 return Poll::Ready(Some(Err(Error::UnexpectedWebSocketFrame {
248 url: Box::new(self.endpoint.clone()),
249 kind: "binary data",
250 })));
251 }
252 Poll::Ready(Some(Ok(Message::Frame(_)))) => {
253 self.socket = None;
254 self.terminated = true;
255 return Poll::Ready(Some(Err(Error::UnexpectedWebSocketFrame {
256 url: Box::new(self.endpoint.clone()),
257 kind: "raw",
258 })));
259 }
260 Poll::Ready(
261 Some(Err(
262 tokio_tungstenite::tungstenite::Error::ConnectionClosed
263 | tokio_tungstenite::tungstenite::Error::AlreadyClosed,
264 ))
265 | None,
266 ) => {
267 self.socket = None;
268 self.terminated = true;
269 return Poll::Ready(None);
270 }
271 Poll::Ready(Some(Err(source))) => {
272 self.socket = None;
273 self.terminated = true;
274 return Poll::Ready(Some(Err(socket_error(self.endpoint.clone(), source))));
275 }
276 }
277 }
278 }
279}
280
281impl FusedStream for EventStream {
282 fn is_terminated(&self) -> bool {
283 self.terminated
284 }
285}
286
287async fn connect_socket(client: &Client) -> Result<(Url, Socket)> {
288 let http_url = client.endpoint("ws")?;
289 let mut endpoint = http_url.clone();
290 endpoint
291 .set_scheme(match http_url.scheme() {
292 "http" => "ws",
293 "https" => "wss",
294 _ => {
295 return Err(Error::Configuration(
296 "panel URL must use HTTP or HTTPS".to_owned(),
297 ));
298 }
299 })
300 .map_err(|()| Error::Configuration("could not construct WebSocket URL".to_owned()))?;
301
302 let mut request = endpoint
303 .as_str()
304 .into_client_request()
305 .map_err(|source| socket_error(endpoint.clone(), source))?;
306 let user_agent = client.user_agent().parse().map_err(|error| {
307 Error::Configuration(format!("invalid WebSocket user agent header: {error}"))
308 })?;
309 request.headers_mut().insert(
310 tokio_tungstenite::tungstenite::http::header::USER_AGENT,
311 user_agent,
312 );
313 if let Some(cookie) = client.cookie_header(&http_url) {
314 request
315 .headers_mut()
316 .insert(tokio_tungstenite::tungstenite::http::header::COOKIE, cookie);
317 }
318
319 let config = WebSocketConfig::default()
320 .max_message_size(Some(MAX_MESSAGE_SIZE))
321 .max_frame_size(Some(MAX_MESSAGE_SIZE));
322 let duration = client.connect_timeout();
323 let connect = async {
324 let transport = open_transport(client, &endpoint).await?;
325 let connector = client
326 .accepts_invalid_certs()
327 .then(insecure_rustls_connector);
328 client_async_tls_with_config(request, transport, Some(config), connector)
329 .await
330 .map_err(|source| handshake_error(endpoint.clone(), source))
331 };
332 let (socket, _) =
333 timeout(duration, connect)
334 .await
335 .map_err(|_| Error::WebSocketConnectTimeout {
336 url: Box::new(endpoint.clone()),
337 timeout: duration,
338 })??;
339 Ok((endpoint, socket))
340}
341
342async fn open_transport(client: &Client, endpoint: &Url) -> Result<BoxedIo> {
343 let Some(proxy) = client.proxy() else {
344 return TcpStream::connect(target_address(endpoint)?)
345 .await
346 .map(|stream| Box::new(stream) as BoxedIo)
347 .map_err(|source| {
348 socket_error(
349 endpoint.clone(),
350 tokio_tungstenite::tungstenite::Error::Io(source),
351 )
352 });
353 };
354
355 let result = match proxy.scheme() {
356 ProxyScheme::Http | ProxyScheme::Https => open_http_tunnel(proxy, endpoint, client).await,
357 ProxyScheme::Socks5 | ProxyScheme::Socks5h => open_socks_tunnel(proxy, endpoint).await,
358 };
359 result.map_err(|source| Error::Proxy {
360 scheme: proxy.scheme(),
361 url: Box::new(endpoint.clone()),
362 source: Box::new(source),
363 })
364}
365
366async fn open_http_tunnel(
367 proxy: &ProxyConfig,
368 endpoint: &Url,
369 client: &Client,
370) -> std::result::Result<BoxedIo, ProxyError> {
371 let tcp = TcpStream::connect((proxy.host(), proxy.port()))
372 .await
373 .map_err(ProxyError::Io)?;
374 let mut transport: BoxedIo = if proxy.scheme() == ProxyScheme::Https {
375 let config = proxy_tls_config(client.accepts_invalid_certs())?;
376 let server_name = ServerName::try_from(proxy.host().to_owned()).map_err(|_| {
377 ProxyError::TlsConfiguration("proxy host is not a valid TLS server name".to_owned())
378 })?;
379 let tls = TlsConnector::from(config)
380 .connect(server_name, tcp)
381 .await
382 .map_err(ProxyError::Io)?;
383 Box::new(tls)
384 } else {
385 Box::new(tcp)
386 };
387
388 let authority = target_authority(endpoint)?;
389 let authentication = proxy
390 .credentials()
391 .map_or_else(String::new, |(username, password)| {
392 let encoded = BASE64_STANDARD.encode(format!("{username}:{password}"));
393 format!("Proxy-Authorization: Basic {encoded}\r\n")
394 });
395 let request = format!(
396 "CONNECT {authority} HTTP/1.1\r\nHost: {authority}\r\nUser-Agent: {}\r\n{authentication}\r\n",
397 client.user_agent()
398 );
399 transport
400 .write_all(request.as_bytes())
401 .await
402 .map_err(ProxyError::Io)?;
403 transport.flush().await.map_err(ProxyError::Io)?;
404 read_connect_response(&mut transport).await?;
405 Ok(transport)
406}
407
408async fn read_connect_response(transport: &mut BoxedIo) -> std::result::Result<(), ProxyError> {
409 let mut response = Vec::with_capacity(1024);
410 while !response.windows(4).any(|window| window == b"\r\n\r\n") {
411 if response.len() == MAX_CONNECT_RESPONSE_SIZE {
412 return Err(ProxyError::HttpResponseTooLarge {
413 limit: MAX_CONNECT_RESPONSE_SIZE,
414 });
415 }
416 let remaining = MAX_CONNECT_RESPONSE_SIZE - response.len();
417 let mut buffer = [0_u8; 1024];
418 let capacity = remaining.min(buffer.len());
419 let read = transport
420 .read(&mut buffer[..capacity])
421 .await
422 .map_err(ProxyError::Io)?;
423 if read == 0 {
424 return Err(ProxyError::InvalidHttpResponse);
425 }
426 response.extend_from_slice(&buffer[..read]);
427 }
428
429 let mut headers = [httparse::EMPTY_HEADER; 32];
430 let mut parsed = httparse::Response::new(&mut headers);
431 match parsed.parse(&response) {
432 Ok(httparse::Status::Complete(_)) => {}
433 Ok(httparse::Status::Partial) | Err(_) => return Err(ProxyError::InvalidHttpResponse),
434 }
435 let status = parsed
436 .code
437 .and_then(|code| StatusCode::from_u16(code).ok())
438 .ok_or(ProxyError::InvalidHttpResponse)?;
439 if !status.is_success() {
440 return Err(ProxyError::HttpStatus(status));
441 }
442 Ok(())
443}
444
445async fn open_socks_tunnel(
446 proxy: &ProxyConfig,
447 endpoint: &Url,
448) -> std::result::Result<BoxedIo, ProxyError> {
449 let host = endpoint.host_str().ok_or(ProxyError::InvalidHttpResponse)?;
450 let port = endpoint
451 .port_or_known_default()
452 .ok_or(ProxyError::InvalidHttpResponse)?;
453 if proxy.scheme().resolves_remotely() {
454 let target = TargetAddr::Domain(Cow::Owned(host.to_owned()), port);
455 return connect_socks(proxy, target)
456 .await
457 .map(|stream| Box::new(stream) as BoxedIo);
458 }
459
460 let addresses = lookup_host((host, port)).await.map_err(ProxyError::Io)?;
461 let mut last_error = None;
462 for address in addresses {
463 match connect_socks(proxy, TargetAddr::Ip(address)).await {
464 Ok(stream) => return Ok(Box::new(stream)),
465 Err(error) => last_error = Some(error),
466 }
467 }
468 Err(last_error.unwrap_or_else(|| {
469 ProxyError::Io(std::io::Error::new(
470 std::io::ErrorKind::NotFound,
471 "target hostname did not resolve",
472 ))
473 }))
474}
475
476async fn connect_socks(
477 proxy: &ProxyConfig,
478 target: TargetAddr<'_>,
479) -> std::result::Result<Socks5Stream<TcpStream>, ProxyError> {
480 let stream = match proxy.credentials() {
481 Some((username, password)) => {
482 Socks5Stream::connect_with_password(
483 (proxy.host(), proxy.port()),
484 target,
485 username,
486 password,
487 )
488 .await
489 }
490 None => Socks5Stream::connect((proxy.host(), proxy.port()), target).await,
491 }
492 .map_err(|source| ProxyError::Socks5(Box::new(source)))?;
493 Ok(stream)
494}
495
496fn target_address(endpoint: &Url) -> Result<(&str, u16)> {
497 let host = endpoint
498 .host_str()
499 .ok_or_else(|| Error::Configuration("WebSocket URL must contain a host".to_owned()))?;
500 let port = endpoint.port_or_known_default().ok_or_else(|| {
501 Error::Configuration("WebSocket URL must contain or imply a port".to_owned())
502 })?;
503 Ok((host, port))
504}
505
506fn target_authority(endpoint: &Url) -> std::result::Result<String, ProxyError> {
507 let host = endpoint.host().ok_or(ProxyError::InvalidHttpResponse)?;
508 let port = endpoint
509 .port_or_known_default()
510 .ok_or(ProxyError::InvalidHttpResponse)?;
511 Ok(match host {
512 Host::Ipv6(address) => format!("[{address}]:{port}"),
513 Host::Ipv4(address) => format!("{address}:{port}"),
514 Host::Domain(domain) => format!("{domain}:{port}"),
515 })
516}
517
518fn proxy_tls_config(
519 accept_invalid_certs: bool,
520) -> std::result::Result<Arc<ClientConfig>, ProxyError> {
521 if accept_invalid_certs {
522 return Ok(match insecure_rustls_connector() {
523 Connector::Rustls(config) => config,
524 _ => unreachable!("insecure connector is always rustls"),
525 });
526 }
527
528 let loaded = rustls_native_certs::load_native_certs();
529 if loaded.certs.is_empty() {
530 return Err(ProxyError::TlsConfiguration(
531 "no native root CA certificates were found".to_owned(),
532 ));
533 }
534 let mut roots = RootCertStore::empty();
535 roots.add_parsable_certificates(loaded.certs);
536 Ok(Arc::new(
537 ClientConfig::builder()
538 .with_root_certificates(roots)
539 .with_no_client_auth(),
540 ))
541}
542
543fn handshake_error(endpoint: Url, source: tokio_tungstenite::tungstenite::Error) -> Error {
544 if let tokio_tungstenite::tungstenite::Error::Http(response) = &source {
545 return match response.status() {
546 StatusCode::UNAUTHORIZED => Error::Unauthorized {
547 method: Method::GET,
548 url: Box::new(endpoint),
549 },
550 StatusCode::FORBIDDEN => Error::Forbidden {
551 method: Method::GET,
552 url: Box::new(endpoint),
553 },
554 status => Error::HttpStatus {
555 method: Method::GET,
556 url: Box::new(endpoint),
557 status,
558 },
559 };
560 }
561 socket_error(endpoint, source)
562}
563
564fn socket_error(endpoint: Url, source: tokio_tungstenite::tungstenite::Error) -> Error {
565 Error::WebSocket {
566 url: Box::new(endpoint),
567 source: Box::new(source),
568 }
569}
570
571fn insecure_rustls_connector() -> Connector {
572 let config = rustls::ClientConfig::builder()
573 .dangerous()
574 .with_custom_certificate_verifier(Arc::new(NoCertificateVerification))
575 .with_no_client_auth();
576 Connector::Rustls(Arc::new(config))
577}
578
579#[derive(Debug)]
580struct NoCertificateVerification;
581
582impl ServerCertVerifier for NoCertificateVerification {
583 fn verify_server_cert(
584 &self,
585 _end_entity: &CertificateDer<'_>,
586 _intermediates: &[CertificateDer<'_>],
587 _server_name: &ServerName<'_>,
588 _ocsp_response: &[u8],
589 _now: UnixTime,
590 ) -> std::result::Result<ServerCertVerified, rustls::Error> {
591 Ok(ServerCertVerified::assertion())
592 }
593
594 fn verify_tls12_signature(
595 &self,
596 _message: &[u8],
597 _certificate: &CertificateDer<'_>,
598 _signature: &DigitallySignedStruct,
599 ) -> std::result::Result<HandshakeSignatureValid, rustls::Error> {
600 Ok(HandshakeSignatureValid::assertion())
601 }
602
603 fn verify_tls13_signature(
604 &self,
605 _message: &[u8],
606 _certificate: &CertificateDer<'_>,
607 _signature: &DigitallySignedStruct,
608 ) -> std::result::Result<HandshakeSignatureValid, rustls::Error> {
609 Ok(HandshakeSignatureValid::assertion())
610 }
611
612 fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
613 vec![
614 SignatureScheme::RSA_PKCS1_SHA1,
615 SignatureScheme::ECDSA_SHA1_Legacy,
616 SignatureScheme::RSA_PKCS1_SHA256,
617 SignatureScheme::ECDSA_NISTP256_SHA256,
618 SignatureScheme::RSA_PKCS1_SHA384,
619 SignatureScheme::ECDSA_NISTP384_SHA384,
620 SignatureScheme::RSA_PKCS1_SHA512,
621 SignatureScheme::ECDSA_NISTP521_SHA512,
622 SignatureScheme::RSA_PSS_SHA256,
623 SignatureScheme::RSA_PSS_SHA384,
624 SignatureScheme::RSA_PSS_SHA512,
625 SignatureScheme::ED25519,
626 SignatureScheme::ED448,
627 ]
628 }
629}
630
631fn origin(url: &Url) -> String {
632 let mut value = url.clone();
633 value.set_path("");
634 value.set_query(None);
635 value.set_fragment(None);
636 value.to_string().trim_end_matches('/').to_owned()
637}
638
639#[cfg(test)]
640mod tests {
641 use super::*;
642 use futures_util::stream::FusedStream;
643 use tokio::io::{DuplexStream, duplex};
644 use tokio_tungstenite::tungstenite::protocol::Role;
645
646 type TestServer = WebSocketStream<DuplexStream>;
647
648 async fn stream_pair() -> (EventStream, TestServer) {
649 let (client_io, server_io) = duplex(64 * 1024);
650 let transport = MaybeTlsStream::Plain(Box::new(client_io) as BoxedIo);
651 let socket = WebSocketStream::from_raw_socket(transport, Role::Client, None).await;
652 let server = WebSocketStream::from_raw_socket(server_io, Role::Server, None).await;
653 let endpoint = Url::parse("ws://panel.example.com/secret/ws").unwrap();
654 let client = Client::new("http://panel.example.com/secret").unwrap();
655 (
656 EventStream {
657 client,
658 endpoint,
659 socket: Some(socket),
660 close: None,
661 terminated: false,
662 },
663 server,
664 )
665 }
666
667 #[tokio::test]
668 async fn stream_skips_control_frames_and_records_peer_close() {
669 let (mut events, mut server) = stream_pair().await;
670 server
671 .send(Message::Ping(vec![1, 2, 3].into()))
672 .await
673 .unwrap();
674 server
675 .send(Message::Text(
676 r#"{"type":"future","payload":{"value":1},"time":7}"#.into(),
677 ))
678 .await
679 .unwrap();
680 let event = events.next_event().await.unwrap().unwrap();
681 assert_eq!(event.timestamp_ms, 7);
682 assert!(matches!(event.kind, crate::PanelEventKind::Unknown { .. }));
683
684 server.send(Message::Close(None)).await.unwrap();
685 assert!(events.next_event().await.unwrap().is_none());
686 assert!(events.is_closed());
687 assert!(events.is_terminated());
688 assert_eq!(
689 events.close_info(),
690 Some(&WebSocketClose {
691 code: 1005,
692 reason: String::new(),
693 })
694 );
695 assert!(events.next_event().await.unwrap().is_none());
696 events.close().await.unwrap();
697 }
698
699 #[tokio::test]
700 async fn binary_frames_fail_closed_and_redact_stream_debug() {
701 let (mut events, mut server) = stream_pair().await;
702 let debug = format!("{events:?}");
703 assert!(debug.contains("panel_origin: \"ws://panel.example.com\""));
704 assert!(!debug.contains("/secret/ws"));
705
706 server
707 .send(Message::Binary(vec![1, 2, 3].into()))
708 .await
709 .unwrap();
710 let error = events.next_event().await.unwrap_err();
711 assert!(matches!(
712 error,
713 Error::UnexpectedWebSocketFrame {
714 kind: "binary data",
715 ..
716 }
717 ));
718 assert!(events.is_closed());
719 assert!(events.is_terminated());
720 }
721
722 #[tokio::test]
723 async fn explicit_close_sends_a_normal_redacted_close_frame() {
724 let (mut events, mut server) = stream_pair().await;
725 events.close().await.unwrap();
726
727 assert!(events.is_closed());
728 assert_eq!(
729 events.close_info(),
730 Some(&WebSocketClose {
731 code: 1000,
732 reason: CLIENT_CLOSE_REASON.to_owned(),
733 })
734 );
735 let message = server.next().await.unwrap().unwrap();
736 let Message::Close(Some(frame)) = message else {
737 panic!("expected a close frame, got {message:?}");
738 };
739 assert_eq!(frame.code, CloseCode::Normal);
740 assert_eq!(frame.reason, CLIENT_CLOSE_REASON);
741 events.close().await.unwrap();
742 }
743
744 #[tokio::test]
745 async fn peer_close_details_and_abrupt_errors_are_observable() {
746 let (mut events, mut server) = stream_pair().await;
747 server
748 .send(Message::Close(Some(CloseFrame {
749 code: CloseCode::Away,
750 reason: "maintenance".into(),
751 })))
752 .await
753 .unwrap();
754 assert!(events.next_event().await.unwrap().is_none());
755 assert_eq!(
756 events.close_info(),
757 Some(&WebSocketClose {
758 code: 1001,
759 reason: "maintenance".to_owned(),
760 })
761 );
762
763 let (mut events, server) = stream_pair().await;
764 drop(server);
765 let result = events.next_event().await;
766 assert!(result.is_err() || result.unwrap().is_none());
767 assert!(events.is_closed());
768 }
769
770 #[tokio::test]
771 async fn missing_socket_terminates_without_panicking() {
772 let (mut events, _server) = stream_pair().await;
773 events.socket = None;
774 assert!(events.next_event().await.unwrap().is_none());
775 assert!(events.is_closed());
776 assert!(events.is_terminated());
777 }
778
779 async fn parse_connect_response(bytes: &[u8]) -> std::result::Result<(), ProxyError> {
780 let (client, mut server) = duplex(MAX_CONNECT_RESPONSE_SIZE + 1024);
781 let bytes = bytes.to_vec();
782 let writer = tokio::spawn(async move {
783 server.write_all(&bytes).await.unwrap();
784 server.shutdown().await.unwrap();
785 });
786 let mut transport: BoxedIo = Box::new(client);
787 let result = read_connect_response(&mut transport).await;
788 writer.await.unwrap();
789 result
790 }
791
792 #[tokio::test]
793 async fn connect_response_parser_bounds_and_validates_proxy_headers() {
794 parse_connect_response(b"HTTP/1.1 200 Connection established\r\nHeader: value\r\n\r\n")
795 .await
796 .unwrap();
797 assert!(matches!(
798 parse_connect_response(b"HTTP/1.1 407 Proxy Authentication Required\r\n\r\n").await,
799 Err(ProxyError::HttpStatus(
800 StatusCode::PROXY_AUTHENTICATION_REQUIRED
801 ))
802 ));
803 assert!(matches!(
804 parse_connect_response(b"not-http\r\n\r\n").await,
805 Err(ProxyError::InvalidHttpResponse)
806 ));
807 assert!(matches!(
808 parse_connect_response(b"HTTP/1.1 200").await,
809 Err(ProxyError::InvalidHttpResponse)
810 ));
811 let oversized = vec![b'x'; MAX_CONNECT_RESPONSE_SIZE];
812 assert!(matches!(
813 parse_connect_response(&oversized).await,
814 Err(ProxyError::HttpResponseTooLarge { limit }) if limit == MAX_CONNECT_RESPONSE_SIZE
815 ));
816 }
817
818 #[test]
819 fn websocket_url_and_error_helpers_cover_all_address_families() {
820 let domain = Url::parse("wss://panel.example.com/ws").unwrap();
821 let ipv4 = Url::parse("ws://127.0.0.1:8080/ws").unwrap();
822 let ipv6 = Url::parse("ws://[::1]:8080/ws").unwrap();
823 assert_eq!(target_address(&domain).unwrap(), ("panel.example.com", 443));
824 assert_eq!(target_authority(&domain).unwrap(), "panel.example.com:443");
825 assert_eq!(target_authority(&ipv4).unwrap(), "127.0.0.1:8080");
826 assert_eq!(target_authority(&ipv6).unwrap(), "[::1]:8080");
827 assert!(target_address(&Url::parse("custom:/path").unwrap()).is_err());
828 assert!(target_address(&Url::parse("custom://host/path").unwrap()).is_err());
829
830 for (status, expected) in [
831 (StatusCode::UNAUTHORIZED, crate::ErrorKind::Unauthorized),
832 (StatusCode::FORBIDDEN, crate::ErrorKind::Forbidden),
833 (StatusCode::BAD_GATEWAY, crate::ErrorKind::HttpStatus),
834 ] {
835 let response = tokio_tungstenite::tungstenite::http::Response::builder()
836 .status(status)
837 .body(None)
838 .unwrap();
839 let error = handshake_error(
840 domain.clone(),
841 tokio_tungstenite::tungstenite::Error::Http(Box::new(response)),
842 );
843 assert_eq!(error.kind(), expected);
844 }
845 let error = socket_error(
846 domain.clone(),
847 tokio_tungstenite::tungstenite::Error::Io(std::io::Error::other("closed")),
848 );
849 assert_eq!(error.kind(), crate::ErrorKind::WebSocket);
850 assert_eq!(
851 origin(&Url::parse("wss://panel.example.com/secret?x=1#fragment").unwrap()),
852 "wss://panel.example.com"
853 );
854 }
855
856 #[test]
857 fn custom_certificate_verifier_is_explicit_and_rustls_only() {
858 let verifier = NoCertificateVerification;
859 let certificate = CertificateDer::from(vec![1, 2, 3]);
860 let server_name = ServerName::try_from("panel.example.com").unwrap();
861 verifier
862 .verify_server_cert(
863 &certificate,
864 &[],
865 &server_name,
866 &[],
867 UnixTime::since_unix_epoch(std::time::Duration::ZERO),
868 )
869 .unwrap();
870 assert_eq!(verifier.supported_verify_schemes().len(), 13);
871 assert!(matches!(insecure_rustls_connector(), Connector::Rustls(_)));
872 assert!(proxy_tls_config(true).is_ok());
873 assert!(proxy_tls_config(false).is_ok());
874 }
875}