elph_ai/api/
websocket_connect.rs

1//! WebSocket TCP/TLS connect with optional HTTP(S) proxy tunnel (mirroring pi-ai Codex).
2
3use std::pin::Pin;
4use std::sync::{Arc, OnceLock};
5use std::task::{Context, Poll};
6
7use anyhow::{Context as AnyhowContext, Result, anyhow, bail};
8use rustls::ClientConfig;
9use rustls::RootCertStore;
10use rustls::pki_types::ServerName;
11use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf};
12
13use tokio::net::TcpStream;
14use tokio_rustls::TlsConnector;
15use tokio_rustls::client::TlsStream;
16use tokio_tungstenite::WebSocketStream;
17use tokio_tungstenite::client_async;
18use tokio_tungstenite::tungstenite::client::IntoClientRequest;
19use tokio_tungstenite::tungstenite::http::HeaderValue;
20use url::Url;
21
22use crate::api::http_proxy::{resolve_http_proxy_url_for_target, websocket_proxy_lookup_url};
23use crate::types::ProviderEnv;
24
25pub type WsStream = WebSocketStream<CodexWsIo>;
26
27/// IO layer for Codex WebSockets (direct TLS, plain, or nested TLS through HTTPS proxy).
28pub enum CodexWsIo {
29    Plain(TcpStream),
30    Tls(Box<TlsStream<TcpStream>>),
31    TlsOverProxyTls(Box<TlsStream<TlsStream<TcpStream>>>),
32}
33
34impl AsyncRead for CodexWsIo {
35    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
36        match self.get_mut() {
37            CodexWsIo::Plain(stream) => Pin::new(stream).poll_read(cx, buf),
38            CodexWsIo::Tls(stream) => Pin::new(stream).poll_read(cx, buf),
39            CodexWsIo::TlsOverProxyTls(stream) => Pin::new(stream).poll_read(cx, buf),
40        }
41    }
42}
43
44impl AsyncWrite for CodexWsIo {
45    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<std::io::Result<usize>> {
46        match self.get_mut() {
47            CodexWsIo::Plain(stream) => Pin::new(stream).poll_write(cx, buf),
48            CodexWsIo::Tls(stream) => Pin::new(stream).poll_write(cx, buf),
49            CodexWsIo::TlsOverProxyTls(stream) => Pin::new(stream).poll_write(cx, buf),
50        }
51    }
52
53    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
54        match self.get_mut() {
55            CodexWsIo::Plain(stream) => Pin::new(stream).poll_flush(cx),
56            CodexWsIo::Tls(stream) => Pin::new(stream).poll_flush(cx),
57            CodexWsIo::TlsOverProxyTls(stream) => Pin::new(stream).poll_flush(cx),
58        }
59    }
60
61    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
62        match self.get_mut() {
63            CodexWsIo::Plain(stream) => Pin::new(stream).poll_shutdown(cx),
64            CodexWsIo::Tls(stream) => Pin::new(stream).poll_shutdown(cx),
65            CodexWsIo::TlsOverProxyTls(stream) => Pin::new(stream).poll_shutdown(cx),
66        }
67    }
68}
69
70fn ensure_crypto_provider() {
71    static INSTALLED: OnceLock<()> = OnceLock::new();
72    INSTALLED.get_or_init(|| {
73        let _ = rustls::crypto::ring::default_provider().install_default();
74    });
75}
76
77fn shared_tls_config() -> Arc<ClientConfig> {
78    static CONFIG: OnceLock<Arc<ClientConfig>> = OnceLock::new();
79    CONFIG
80        .get_or_init(|| {
81            ensure_crypto_provider();
82            let mut roots = RootCertStore::empty();
83            roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
84            Arc::new(
85                ClientConfig::builder()
86                    .with_root_certificates(roots)
87                    .with_no_client_auth(),
88            )
89        })
90        .clone()
91}
92
93fn parse_websocket_endpoint(ws_url: &str) -> Result<(bool, String, u16)> {
94    let url = Url::parse(ws_url).with_context(|| format!("invalid WebSocket URL: {ws_url}"))?;
95    let host = url
96        .host_str()
97        .ok_or_else(|| anyhow!("WebSocket URL missing host: {ws_url}"))?
98        .to_string();
99    let secure = matches!(url.scheme(), "wss" | "https");
100    let port = url.port().unwrap_or(if secure { 443 } else { 80 });
101    Ok((secure, host, port))
102}
103
104fn proxy_endpoint(proxy_url: &Url) -> Result<(String, u16)> {
105    let host = proxy_url
106        .host_str()
107        .ok_or_else(|| anyhow!("proxy URL missing host"))?
108        .to_string();
109    let port = proxy_url
110        .port()
111        .unwrap_or(if proxy_url.scheme() == "https" { 443 } else { 80 });
112    Ok((host, port))
113}
114
115async fn read_until_headers_end(stream: &mut (impl AsyncReadExt + Unpin)) -> Result<String> {
116    let mut buf = Vec::new();
117    let mut chunk = [0u8; 1024];
118    loop {
119        let read = stream.read(&mut chunk).await?;
120        if read == 0 {
121            bail!("proxy closed before CONNECT response");
122        }
123        buf.extend_from_slice(&chunk[..read]);
124        if buf.windows(4).any(|window| window == b"\r\n\r\n") {
125            break;
126        }
127        if buf.len() > 16_384 {
128            bail!("proxy CONNECT response too large");
129        }
130    }
131    Ok(String::from_utf8_lossy(&buf).into_owned())
132}
133
134fn ensure_connect_success(response: &str) -> Result<()> {
135    let status_line = response.lines().next().unwrap_or_default();
136    if status_line.contains(" 200 ") {
137        return Ok(());
138    }
139    bail!("proxy CONNECT failed: {status_line}")
140}
141
142async fn send_http_connect(
143    stream: &mut (impl AsyncReadExt + AsyncWriteExt + Unpin),
144    target_host: &str,
145    target_port: u16,
146) -> Result<()> {
147    let request = format!("CONNECT {target_host}:{target_port} HTTP/1.1\r\nHost: {target_host}:{target_port}\r\n\r\n");
148    stream.write_all(request.as_bytes()).await?;
149    let response = read_until_headers_end(stream).await?;
150    ensure_connect_success(&response)
151}
152
153async fn tls_handshake<S>(server_name: &str, stream: S) -> Result<TlsStream<S>>
154where
155    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
156{
157    let dns_name = ServerName::try_from(server_name.to_string()).map_err(|_| anyhow!("invalid TLS server name"))?;
158    let connector = TlsConnector::from(shared_tls_config());
159    connector
160        .connect(dns_name, stream)
161        .await
162        .map_err(|error| anyhow!("TLS handshake failed for {server_name}: {error}"))
163}
164
165async fn open_stream(ws_url: &str, env: Option<&ProviderEnv>) -> Result<CodexWsIo> {
166    let (secure, host, port) = parse_websocket_endpoint(ws_url)?;
167    let lookup_url = websocket_proxy_lookup_url(ws_url);
168    let proxy = resolve_http_proxy_url_for_target(&lookup_url, env)?;
169
170    let Some(proxy_url) = proxy else {
171        let tcp = TcpStream::connect((host.as_str(), port))
172            .await
173            .with_context(|| format!("failed to connect to {host}:{port}"))?;
174        return if secure {
175            Ok(CodexWsIo::Tls(Box::new(tls_handshake(&host, tcp).await?)))
176        } else {
177            Ok(CodexWsIo::Plain(tcp))
178        };
179    };
180
181    let (proxy_host, proxy_port) = proxy_endpoint(&proxy_url)?;
182    let tcp = TcpStream::connect((proxy_host.as_str(), proxy_port))
183        .await
184        .with_context(|| format!("failed to connect to proxy {proxy_host}:{proxy_port}"))?;
185
186    if proxy_url.scheme() == "https" {
187        let mut proxy_tls = tls_handshake(proxy_host.as_str(), tcp).await?;
188        send_http_connect(&mut proxy_tls, &host, port).await?;
189        if secure {
190            let target_tls = tls_handshake(&host, proxy_tls).await?;
191            Ok(CodexWsIo::TlsOverProxyTls(Box::new(target_tls)))
192        } else {
193            bail!("HTTPS proxy with non-secure WebSocket target is unsupported");
194        }
195    } else {
196        let mut stream = tcp;
197        send_http_connect(&mut stream, &host, port).await?;
198        if secure {
199            Ok(CodexWsIo::Tls(Box::new(tls_handshake(&host, stream).await?)))
200        } else {
201            Ok(CodexWsIo::Plain(stream))
202        }
203    }
204}
205
206/// Open a WebSocket connection, optionally routing through HTTP proxy env vars.
207pub async fn connect_websocket_with_proxy(
208    ws_url: &str,
209    headers: &std::collections::HashMap<String, String>,
210    timeout_ms: u64,
211    env: Option<&ProviderEnv>,
212) -> Result<WsStream> {
213    let mut request = ws_url.into_client_request()?;
214    for (k, v) in headers {
215        if k.eq_ignore_ascii_case("accept") {
216            continue;
217        }
218        request.headers_mut().insert(
219            http::HeaderName::from_bytes(k.as_bytes()).map_err(|e| anyhow!("invalid header name {k}: {e}"))?,
220            HeaderValue::from_str(v).map_err(|e| anyhow!("invalid header {k}: {e}"))?,
221        );
222    }
223
224    let timeout = std::time::Duration::from_millis(timeout_ms);
225    let connect = tokio::time::timeout(timeout, async {
226        let stream = open_stream(ws_url, env).await?;
227        client_async(request, stream)
228            .await
229            .map(|(socket, _)| socket)
230            .map_err(Into::into)
231    });
232    connect.await.map_err(|_| anyhow!("WebSocket connect timeout"))?
233}
234
235#[cfg(test)]
236mod tests {
237    use super::*;
238
239    #[test]
240    fn parses_wss_endpoint() {
241        let (secure, host, port) = parse_websocket_endpoint("wss://chatgpt.com/backend-api/codex/responses").unwrap();
242        assert!(secure);
243        assert_eq!(host, "chatgpt.com");
244        assert_eq!(port, 443);
245    }
246
247    #[test]
248    fn parses_ws_custom_port() {
249        let (secure, host, port) = parse_websocket_endpoint("ws://localhost:9001/ws").unwrap();
250        assert!(!secure);
251        assert_eq!(host, "localhost");
252        assert_eq!(port, 9001);
253    }
254
255    #[test]
256    fn connect_response_must_be_200() {
257        ensure_connect_success("HTTP/1.1 200 Connection Established\r\n\r\n").unwrap();
258        assert!(ensure_connect_success("HTTP/1.1 403 Forbidden\r\n\r\n").is_err());
259    }
260}
elph_ai/api/websocket_connect.rs

elph_ai/api/
websocket_connect.rs
