Skip to main content

async_wsocket/native/
mod.rs

1// Copyright (c) 2022-2024 Yuki Kishimoto
2// Distributed under the MIT software license
3
4//! Native
5
6#[cfg(feature = "socks")]
7use std::net::SocketAddr;
8
9use tokio::io::{AsyncRead, AsyncWrite};
10use tokio::net::TcpStream;
11use tokio_tungstenite::tungstenite::protocol::Role;
12pub use tokio_tungstenite::tungstenite::Message;
13pub use tokio_tungstenite::WebSocketStream;
14use url::Url;
15
16mod error;
17#[cfg(feature = "socks")]
18mod socks;
19
20pub use self::error::Error;
21#[cfg(feature = "socks")]
22use self::socks::TcpSocks5Stream;
23use crate::socket::WebSocket;
24use crate::ConnectionMode;
25
26pub async fn connect(url: &Url, mode: &ConnectionMode) -> Result<WebSocket, Error> {
27    match mode {
28        ConnectionMode::Direct => connect_direct(url).await,
29        #[cfg(feature = "socks")]
30        ConnectionMode::Proxy(proxy) => connect_proxy(url, *proxy).await,
31    }
32}
33
34async fn connect_direct(url: &Url) -> Result<WebSocket, Error> {
35    let host: &str = url.host_str().ok_or_else(Error::empty_host)?;
36    let port: u16 = url
37        .port_or_known_default()
38        .ok_or_else(Error::invalid_port)?;
39
40    let host: String = format!("{}:{}", host, port);
41
42    let tcp_stream: TcpStream = tokio_happy_eyeballs::connect(host).await?;
43
44    // NOT REMOVE `Box::pin`!
45    // Use `Box::pin` to fix stack overflow on windows targets due to large `Future`
46    let (stream, _) = Box::pin(tokio_tungstenite::client_async_tls(
47        url.as_str(),
48        tcp_stream,
49    ))
50    .await?;
51    Ok(WebSocket::tokio(Box::new(stream)))
52}
53
54#[cfg(feature = "socks")]
55async fn connect_proxy(url: &Url, proxy: SocketAddr) -> Result<WebSocket, Error> {
56    let host: &str = url.host_str().ok_or_else(Error::empty_host)?;
57    let port: u16 = url
58        .port_or_known_default()
59        .ok_or_else(Error::invalid_port)?;
60    let addr: String = format!("{host}:{port}");
61
62    let conn: TcpStream = TcpSocks5Stream::connect(proxy, addr).await?;
63    // NOT REMOVE `Box::pin`!
64    // Use `Box::pin` to fix stack overflow on windows targets due to large `Future`
65    let (stream, _) = Box::pin(tokio_tungstenite::client_async_tls(url.as_str(), conn)).await?;
66    Ok(WebSocket::tokio(Box::new(stream)))
67}
68
69#[inline]
70pub async fn accept<S>(raw_stream: S) -> Result<WebSocketStream<S>, Error>
71where
72    S: AsyncRead + AsyncWrite + Unpin,
73{
74    Ok(tokio_tungstenite::accept_async(raw_stream).await?)
75}
76
77/// Take an already upgraded websocket connection
78///
79/// Useful for when using [hyper] or [warp] or any other HTTP server
80#[inline]
81pub async fn take_upgraded<S>(raw_stream: S) -> WebSocketStream<S>
82where
83    S: AsyncRead + AsyncWrite + Unpin,
84{
85    WebSocketStream::from_raw_socket(raw_stream, Role::Server, None).await
86}