1//! socks proxy
2
3use anyhow::{Context as _, Result};
4use http_client::Url;
5use tokio::net::TcpStream;
6use tokio_socks::{
7 IntoTargetAddr, TargetAddr,
8 tcp::{Socks4Stream, Socks5Stream},
9};
10
11use super::AsyncReadWrite;
12
13/// Identification to a Socks V4 Proxy
14pub(super) struct Socks4Identification<'a> {
15 user_id: &'a str,
16}
17
18/// Authorization to a Socks V5 Proxy
19pub(super) struct Socks5Authorization<'a> {
20 username: &'a str,
21 password: &'a str,
22}
23
24/// Socks Proxy Protocol Version
25///
26/// V4 allows identification using a user_id
27/// V5 allows authorization using a username and password
28pub(super) enum SocksVersion<'a> {
29 V4 {
30 local_dns: bool,
31 identification: Option<Socks4Identification<'a>>,
32 },
33 V5 {
34 local_dns: bool,
35 authorization: Option<Socks5Authorization<'a>>,
36 },
37}
38
39pub(super) fn parse_socks_proxy<'t>(scheme: &str, proxy: &'t Url) -> SocksVersion<'t> {
40 if scheme.starts_with("socks4") {
41 let identification = match proxy.username() {
42 "" => None,
43 username => Some(Socks4Identification { user_id: username }),
44 };
45 SocksVersion::V4 {
46 local_dns: scheme != "socks4a",
47 identification,
48 }
49 } else {
50 let authorization = proxy.password().map(|password| Socks5Authorization {
51 username: proxy.username(),
52 password,
53 });
54 SocksVersion::V5 {
55 local_dns: scheme != "socks5h",
56 authorization,
57 }
58 }
59}
60
61pub(super) async fn connect_socks_proxy_stream(
62 stream: TcpStream,
63 socks_version: SocksVersion<'_>,
64 rpc_host: (&str, u16),
65) -> Result<Box<dyn AsyncReadWrite>> {
66 let rpc_host = rpc_host
67 .into_target_addr()
68 .context("Failed to parse target addr")?;
69
70 let local_dns = match &socks_version {
71 SocksVersion::V4 { local_dns, .. } => local_dns,
72 SocksVersion::V5 { local_dns, .. } => local_dns,
73 };
74 let rpc_host = match (rpc_host, local_dns) {
75 (TargetAddr::Domain(domain, port), true) => {
76 let ip_addr = tokio::net::lookup_host((domain.as_ref(), port))
77 .await
78 .with_context(|| format!("Failed to lookup domain {}", domain))?
79 .next()
80 .ok_or_else(|| anyhow::anyhow!("Failed to lookup domain {}", domain))?;
81 TargetAddr::Ip(ip_addr)
82 }
83 (rpc_host, _) => rpc_host,
84 };
85
86 match socks_version {
87 SocksVersion::V4 {
88 identification: None,
89 ..
90 } => {
91 let socks = Socks4Stream::connect_with_socket(stream, rpc_host)
92 .await
93 .context("error connecting to socks")?;
94 Ok(Box::new(socks))
95 }
96 SocksVersion::V4 {
97 identification: Some(Socks4Identification { user_id }),
98 ..
99 } => {
100 let socks = Socks4Stream::connect_with_userid_and_socket(stream, rpc_host, user_id)
101 .await
102 .context("error connecting to socks")?;
103 Ok(Box::new(socks))
104 }
105 SocksVersion::V5 {
106 authorization: None,
107 ..
108 } => {
109 let socks = Socks5Stream::connect_with_socket(stream, rpc_host)
110 .await
111 .context("error connecting to socks")?;
112 Ok(Box::new(socks))
113 }
114 SocksVersion::V5 {
115 authorization: Some(Socks5Authorization { username, password }),
116 ..
117 } => {
118 let socks = Socks5Stream::connect_with_password_and_socket(
119 stream, rpc_host, username, password,
120 )
121 .await
122 .context("error connecting to socks")?;
123 Ok(Box::new(socks))
124 }
125 }
126}
127
128#[cfg(test)]
129mod tests {
130 use url::Url;
131
132 use super::*;
133
134 #[test]
135 fn parse_socks4() {
136 let proxy = Url::parse("socks4://proxy.example.com:1080").unwrap();
137 let scheme = proxy.scheme();
138
139 let version = parse_socks_proxy(scheme, &proxy);
140 assert!(matches!(
141 version,
142 SocksVersion::V4 {
143 local_dns: true,
144 identification: None
145 }
146 ))
147 }
148
149 #[test]
150 fn parse_socks4_with_identification() {
151 let proxy = Url::parse("socks4://userid@proxy.example.com:1080").unwrap();
152 let scheme = proxy.scheme();
153
154 let version = parse_socks_proxy(scheme, &proxy);
155 assert!(matches!(
156 version,
157 SocksVersion::V4 {
158 local_dns: true,
159 identification: Some(Socks4Identification { user_id: "userid" })
160 }
161 ))
162 }
163
164 #[test]
165 fn parse_socks4_with_remote_dns() {
166 let proxy = Url::parse("socks4a://proxy.example.com:1080").unwrap();
167 let scheme = proxy.scheme();
168
169 let version = parse_socks_proxy(scheme, &proxy);
170 assert!(matches!(
171 version,
172 SocksVersion::V4 {
173 local_dns: false,
174 identification: None
175 }
176 ))
177 }
178
179 #[test]
180 fn parse_socks5() {
181 let proxy = Url::parse("socks5://proxy.example.com:1080").unwrap();
182 let scheme = proxy.scheme();
183
184 let version = parse_socks_proxy(scheme, &proxy);
185 assert!(matches!(
186 version,
187 SocksVersion::V5 {
188 local_dns: true,
189 authorization: None
190 }
191 ))
192 }
193
194 #[test]
195 fn parse_socks5_with_authorization() {
196 let proxy = Url::parse("socks5://username:password@proxy.example.com:1080").unwrap();
197 let scheme = proxy.scheme();
198
199 let version = parse_socks_proxy(scheme, &proxy);
200 assert!(matches!(
201 version,
202 SocksVersion::V5 {
203 local_dns: true,
204 authorization: Some(Socks5Authorization {
205 username: "username",
206 password: "password"
207 })
208 }
209 ))
210 }
211
212 #[test]
213 fn parse_socks5_with_remote_dns() {
214 let proxy = Url::parse("socks5h://proxy.example.com:1080").unwrap();
215 let scheme = proxy.scheme();
216
217 let version = parse_socks_proxy(scheme, &proxy);
218 assert!(matches!(
219 version,
220 SocksVersion::V5 {
221 local_dns: false,
222 authorization: None
223 }
224 ))
225 }
226}