typeduck-codex-utils-rustls-provider 0.17.0

Tokio binding for Tungstenite, the Lightweight stream-based WebSocket implementation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215

#![cfg(feature = "proxy")]

use std::env;

use futures_util::{SinkExt, StreamExt};
use tokio::{
    io::{self, AsyncReadExt, AsyncWriteExt},
    net::{TcpListener, TcpStream},
    time::{timeout, Duration},
};
use tokio_tungstenite::{accept_async, connect_async, tungstenite::Message};

#[tokio::test]
async fn proxy_http_and_socks5() {
    run_proxy_test("HTTP_PROXY", "http").await;
    run_proxy_test("ALL_PROXY", "socks5").await;
}

#[tokio::test]
async fn proxy_http_and_socks5_real() {
    // Optional real-proxy checks:
    // - REAL_HTTP_PROXY / REAL_SOCKS5_PROXY are used to avoid clobbering any
    //   existing HTTP_PROXY / ALL_PROXY settings on the developer machine.
    // - The test copies values into HTTP_PROXY/ALL_PROXY only while it runs.
    let http_proxy = env::var("REAL_HTTP_PROXY").ok();
    let socks_proxy = env::var("REAL_SOCKS5_PROXY").ok();

    if let Some(http_proxy) = http_proxy {
        run_proxy_test_with_url("HTTP_PROXY", &http_proxy).await;
    }
    if let Some(socks_proxy) = socks_proxy {
        run_proxy_test_with_url("ALL_PROXY", &socks_proxy).await;
    }
}

async fn run_proxy_test(proxy_env_key: &str, proxy_scheme: &str) {
    let (target_port, target_handle) = spawn_ws_echo_server().await;
    let target_addr = format!("127.0.0.1:{target_port}");

    let (proxy_port, proxy_handle) = spawn_proxy(proxy_env_key, &target_addr).await;

    let prev_http_proxy = env::var("HTTP_PROXY").ok();
    let prev_https_proxy = env::var("HTTPS_PROXY").ok();
    let prev_all_proxy = env::var("ALL_PROXY").ok();
    let prev_no_proxy = env::var("NO_PROXY").ok();

    env::remove_var("HTTP_PROXY");
    env::remove_var("HTTPS_PROXY");
    env::remove_var("ALL_PROXY");
    env::remove_var("NO_PROXY");

    let proxy_url = format!("{proxy_scheme}://127.0.0.1:{proxy_port}");
    env::set_var(proxy_env_key, proxy_url);

    let url = format!("ws://{target_addr}");
    let (mut socket, _response) = connect_async(url).await.expect("proxy connect");
    socket.send(Message::Text("hello".into())).await.expect("send");
    let msg = socket.next().await.expect("read").expect("message");
    assert_eq!(msg, Message::Text("hello".into()));
    let _ = socket.close(None).await;

    restore_env("HTTP_PROXY", prev_http_proxy);
    restore_env("HTTPS_PROXY", prev_https_proxy);
    restore_env("ALL_PROXY", prev_all_proxy);
    restore_env("NO_PROXY", prev_no_proxy);

    proxy_handle.await.expect("proxy task");
    target_handle.await.expect("target task");
}

async fn run_proxy_test_with_url(proxy_env_key: &str, proxy_url: &str) {
    let (target_port, target_handle) = spawn_ws_echo_server().await;
    let target_addr = format!("127.0.0.1:{target_port}");

    let prev_http_proxy = env::var("HTTP_PROXY").ok();
    let prev_https_proxy = env::var("HTTPS_PROXY").ok();
    let prev_all_proxy = env::var("ALL_PROXY").ok();
    let prev_no_proxy = env::var("NO_PROXY").ok();

    env::remove_var("HTTP_PROXY");
    env::remove_var("HTTPS_PROXY");
    env::remove_var("ALL_PROXY");
    env::remove_var("NO_PROXY");

    env::set_var(proxy_env_key, proxy_url);

    let url = format!("ws://{target_addr}");
    let (mut socket, _response) = connect_async(url).await.expect("proxy connect");
    socket.send(Message::Text("hello".into())).await.expect("send");
    let msg = socket.next().await.expect("read").expect("message");
    assert_eq!(msg, Message::Text("hello".into()));
    let _ = socket.close(None).await;

    restore_env("HTTP_PROXY", prev_http_proxy);
    restore_env("HTTPS_PROXY", prev_https_proxy);
    restore_env("ALL_PROXY", prev_all_proxy);
    restore_env("NO_PROXY", prev_no_proxy);

    target_handle.await.expect("target task");
}

fn restore_env(key: &str, value: Option<String>) {
    match value {
        Some(value) => env::set_var(key, value),
        None => env::remove_var(key),
    }
}

async fn spawn_ws_echo_server() -> (u16, tokio::task::JoinHandle<()>) {
    let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind ws server");
    let port = listener.local_addr().expect("addr").port();
    let handle = tokio::spawn(async move {
        let (stream, _) = listener.accept().await.expect("accept");
        let mut ws = accept_async(stream).await.expect("accept ws");
        if let Some(Ok(msg)) = ws.next().await {
            let _ = ws.send(msg).await;
        }
        let _ = ws.close(None).await;
    });
    (port, handle)
}

async fn spawn_proxy(
    proxy_env_key: &str,
    target_addr: &str,
) -> (u16, tokio::task::JoinHandle<()>) {
    let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind proxy");
    let port = listener.local_addr().expect("addr").port();
    let target_addr = target_addr.to_string();
    let proxy_env_key = proxy_env_key.to_string();
    let handle = tokio::spawn(async move {
        let (stream, _) = listener.accept().await.expect("accept");
        if proxy_env_key == "HTTP_PROXY" {
            handle_http_connect(stream, &target_addr).await;
        } else {
            handle_socks5(stream, &target_addr).await;
        }
    });
    (port, handle)
}

async fn handle_http_connect(mut client: TcpStream, target_addr: &str) {
    let mut buf = Vec::new();
    let mut chunk = [0u8; 512];
    loop {
        let read = client.read(&mut chunk).await.expect("read");
        if read == 0 {
            return;
        }
        buf.extend_from_slice(&chunk[..read]);
        if buf.windows(4).any(|w| w == b"\r\n\r\n") {
            break;
        }
    }

    if !std::str::from_utf8(&buf).unwrap_or("").starts_with("CONNECT") {
        return;
    }

    let mut upstream = TcpStream::connect(target_addr).await.expect("connect upstream");
    client
        .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n")
        .await
        .expect("write");

    let _ = timeout(Duration::from_secs(2), io::copy_bidirectional(&mut client, &mut upstream)).await;
}

async fn handle_socks5(mut client: TcpStream, target_addr: &str) {
    let mut header = [0u8; 2];
    client.read_exact(&mut header).await.expect("read greeting");
    let methods_len = header[1] as usize;
    let mut methods = vec![0u8; methods_len];
    client.read_exact(&mut methods).await.expect("read methods");
    client.write_all(&[0x05, 0x00]).await.expect("write method");

    let mut req = [0u8; 4];
    client.read_exact(&mut req).await.expect("read request");
    if req[1] != 0x01 {
        return;
    }

    let _addr = match req[3] {
        0x01 => {
            let mut ip = [0u8; 4];
            client.read_exact(&mut ip).await.expect("read ip");
            std::net::Ipv4Addr::from(ip).to_string()
        }
        0x03 => {
            let mut len = [0u8; 1];
            client.read_exact(&mut len).await.expect("read len");
            let mut host = vec![0u8; len[0] as usize];
            client.read_exact(&mut host).await.expect("read host");
            String::from_utf8_lossy(&host).to_string()
        }
        0x04 => {
            let mut ip = [0u8; 16];
            client.read_exact(&mut ip).await.expect("read ip");
            std::net::Ipv6Addr::from(ip).to_string()
        }
        _ => return,
    };

    let mut port = [0u8; 2];
    client.read_exact(&mut port).await.expect("read port");
    let _port = u16::from_be_bytes(port);

    let mut upstream = TcpStream::connect(target_addr).await.expect("connect upstream");
    client
        .write_all(&[0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0])
        .await
        .expect("write reply");

    let _ = timeout(Duration::from_secs(2), io::copy_bidirectional(&mut client, &mut upstream)).await;
}