1use std::net::SocketAddr;
2use std::sync::atomic::{AtomicU64, Ordering};
3use std::sync::Arc;
4
5use tokio::io::{AsyncReadExt, AsyncWriteExt};
6use tokio::net::{TcpListener, UdpSocket};
7use tokio_rustls::TlsAcceptor;
8
9pub struct TcpEchoServer {
10 addr: SocketAddr,
11 connection_count: Arc<AtomicU64>,
12 handle: tokio::task::JoinHandle<()>,
13}
14
15impl TcpEchoServer {
16 pub async fn start() -> Self {
17 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
18 let addr = listener.local_addr().unwrap();
19 let connection_count = Arc::new(AtomicU64::new(0));
20 let cc = connection_count.clone();
21
22 let handle = tokio::spawn(async move {
23 loop {
24 let (mut stream, _) = match listener.accept().await {
25 Ok(s) => s,
26 Err(_) => break,
27 };
28 cc.fetch_add(1, Ordering::Relaxed);
29 tokio::spawn(async move {
30 let mut buf = [0u8; 4096];
31 loop {
32 match stream.read(&mut buf).await {
33 Ok(0) => break,
34 Ok(n) => {
35 if stream.write_all(&buf[..n]).await.is_err() {
36 break;
37 }
38 }
39 Err(_) => break,
40 }
41 }
42 });
43 }
44 });
45
46 Self {
47 addr,
48 connection_count,
49 handle,
50 }
51 }
52
53 pub fn addr(&self) -> SocketAddr {
54 self.addr
55 }
56
57 pub fn connection_count(&self) -> &AtomicU64 {
58 &self.connection_count
59 }
60}
61
62impl Drop for TcpEchoServer {
63 fn drop(&mut self) {
64 self.handle.abort();
65 }
66}
67
68pub struct UdpEchoServer {
69 addr: SocketAddr,
70 packet_count: Arc<AtomicU64>,
71 handle: tokio::task::JoinHandle<()>,
72}
73
74impl UdpEchoServer {
75 pub async fn start() -> Self {
76 let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
77 let addr = socket.local_addr().unwrap();
78 let packet_count = Arc::new(AtomicU64::new(0));
79 let pc = packet_count.clone();
80
81 let handle = tokio::spawn(async move {
82 let mut buf = [0u8; 65535];
83 loop {
84 let (n, peer) = match socket.recv_from(&mut buf).await {
85 Ok(v) => v,
86 Err(_) => break,
87 };
88 pc.fetch_add(1, Ordering::Relaxed);
89 let _ = socket.send_to(&buf[..n], peer).await;
90 }
91 });
92
93 Self {
94 addr,
95 packet_count,
96 handle,
97 }
98 }
99
100 pub fn addr(&self) -> SocketAddr {
101 self.addr
102 }
103
104 pub fn packet_count(&self) -> &AtomicU64 {
105 &self.packet_count
106 }
107}
108
109impl Drop for UdpEchoServer {
110 fn drop(&mut self) {
111 self.handle.abort();
112 }
113}
114
115pub struct HttpOriginServer {
116 addr: SocketAddr,
117 request_count: Arc<AtomicU64>,
118 handle: tokio::task::JoinHandle<()>,
119}
120
121impl HttpOriginServer {
122 pub async fn start() -> Self {
123 Self::start_with_body(b"hello from origin").await
124 }
125
126 pub async fn start_with_body(body: &[u8]) -> Self {
127 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
128 let addr = listener.local_addr().unwrap();
129 let request_count = Arc::new(AtomicU64::new(0));
130 let rc = request_count.clone();
131 let body = body.to_vec();
132
133 let handle = tokio::spawn(async move {
134 loop {
135 let (mut stream, _) = match listener.accept().await {
136 Ok(s) => s,
137 Err(_) => break,
138 };
139 rc.fetch_add(1, Ordering::Relaxed);
140 let body = body.clone();
141 tokio::spawn(async move {
142 let mut request = Vec::new();
143 let mut buf = [0u8; 4096];
144 loop {
145 match stream.read(&mut buf).await {
146 Ok(0) => return,
147 Ok(n) => {
148 request.extend_from_slice(&buf[..n]);
149 if request.windows(4).any(|w| w == b"\r\n\r\n") {
150 break;
151 }
152 }
153 Err(_) => return,
154 }
155 }
156 let response = format!(
157 "HTTP/1.1 200 OK\r\n\
158 Content-Length: {}\r\n\
159 Connection: close\r\n\
160 \r\n",
161 body.len()
162 );
163 let _ = stream.write_all(response.as_bytes()).await;
164 let _ = stream.write_all(&body).await;
165 let _ = stream.shutdown().await;
166 });
167 }
168 });
169
170 Self {
171 addr,
172 request_count,
173 handle,
174 }
175 }
176
177 pub fn addr(&self) -> SocketAddr {
178 self.addr
179 }
180
181 pub fn request_count(&self) -> &AtomicU64 {
182 &self.request_count
183 }
184}
185
186impl Drop for HttpOriginServer {
187 fn drop(&mut self) {
188 self.handle.abort();
189 }
190}
191
192pub struct HttpConnectUpstream {
193 addr: SocketAddr,
194 connection_count: Arc<AtomicU64>,
195 handle: tokio::task::JoinHandle<()>,
196}
197
198impl HttpConnectUpstream {
199 pub async fn start() -> Self {
200 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
201 let addr = listener.local_addr().unwrap();
202 let connection_count = Arc::new(AtomicU64::new(0));
203 let cc = connection_count.clone();
204
205 let handle = tokio::spawn(async move {
206 loop {
207 let (mut stream, _) = match listener.accept().await {
208 Ok(s) => s,
209 Err(_) => break,
210 };
211 cc.fetch_add(1, Ordering::Relaxed);
212 tokio::spawn(async move {
213 let mut request = Vec::new();
214 let mut buf = [0u8; 4096];
215 loop {
216 match stream.read(&mut buf).await {
217 Ok(0) => return,
218 Ok(n) => {
219 request.extend_from_slice(&buf[..n]);
220 if request.windows(4).any(|w| w == b"\r\n\r\n") {
221 break;
222 }
223 }
224 Err(_) => return,
225 }
226 }
227
228 let request_str = String::from_utf8_lossy(&request);
229 let first_line = request_str.lines().next().unwrap_or("");
230 let parts: Vec<&str> = first_line.split_whitespace().collect();
231
232 if parts.len() < 2 || parts[0] != "CONNECT" {
233 let _ = stream.write_all(b"HTTP/1.1 400 Bad Request\r\n\r\n").await;
234 return;
235 }
236
237 let target_addr = parts[1].to_string();
238
239 let target = match tokio::net::TcpStream::connect(&target_addr).await {
240 Ok(t) => t,
241 Err(_) => {
242 let _ = stream.write_all(b"HTTP/1.1 502 Bad Gateway\r\n\r\n").await;
243 return;
244 }
245 };
246
247 if stream
248 .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n")
249 .await
250 .is_err()
251 {
252 return;
253 }
254
255 let (mut cr, mut cw) = stream.into_split();
256 let (mut tr, mut tw) = target.into_split();
257
258 let c2t = tokio::spawn(async move {
259 let _ = tokio::io::copy(&mut cr, &mut tw).await;
260 let _ = tw.shutdown().await;
261 });
262 let t2c = tokio::spawn(async move {
263 let _ = tokio::io::copy(&mut tr, &mut cw).await;
264 let _ = cw.shutdown().await;
265 });
266 let _ = tokio::join!(c2t, t2c);
267 });
268 }
269 });
270
271 Self {
272 addr,
273 connection_count,
274 handle,
275 }
276 }
277
278 pub fn addr(&self) -> SocketAddr {
279 self.addr
280 }
281
282 pub fn connection_count(&self) -> &AtomicU64 {
283 &self.connection_count
284 }
285}
286
287impl Drop for HttpConnectUpstream {
288 fn drop(&mut self) {
289 self.handle.abort();
290 }
291}
292
293pub struct Socks5Upstream {
294 addr: SocketAddr,
295 connection_count: Arc<AtomicU64>,
296 handle: tokio::task::JoinHandle<()>,
297}
298
299impl Socks5Upstream {
300 pub async fn start() -> Self {
301 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
302 let addr = listener.local_addr().unwrap();
303 let connection_count = Arc::new(AtomicU64::new(0));
304 let cc = connection_count.clone();
305
306 let handle =
307 tokio::spawn(async move {
308 loop {
309 let (mut stream, _) = match listener.accept().await {
310 Ok(s) => s,
311 Err(_) => break,
312 };
313 cc.fetch_add(1, Ordering::Relaxed);
314 tokio::spawn(async move {
315 let mut header = [0u8; 2];
316 if stream.read_exact(&mut header).await.is_err() {
317 return;
318 }
319 let nmethods = header[1] as usize;
320 let mut methods = vec![0u8; nmethods];
321 if stream.read_exact(&mut methods).await.is_err() {
322 return;
323 }
324 if stream.write_all(&[0x05, 0x00]).await.is_err() {
325 return;
326 }
327
328 let mut req = [0u8; 4];
329 if stream.read_exact(&mut req).await.is_err() {
330 return;
331 }
332 let atyp = req[3];
333 let target_addr =
334 match atyp {
335 0x01 => {
336 let mut addr = [0u8; 4];
337 if stream.read_exact(&mut addr).await.is_err() {
338 return;
339 }
340 let port = stream.read_u16().await.unwrap_or(0);
341 format!(
342 "{}.{}.{}.{}:{}",
343 addr[0], addr[1], addr[2], addr[3], port
344 )
345 }
346 0x03 => {
347 let len = stream.read_u8().await.unwrap_or(0) as usize;
348 let mut domain = vec![0u8; len];
349 if stream.read_exact(&mut domain).await.is_err() {
350 return;
351 }
352 let port = stream.read_u16().await.unwrap_or(0);
353 let domain = String::from_utf8_lossy(&domain);
354 format!("{}:{}", domain, port)
355 }
356 0x04 => {
357 let mut addr = [0u8; 16];
358 if stream.read_exact(&mut addr).await.is_err() {
359 return;
360 }
361 let port = stream.read_u16().await.unwrap_or(0);
362 format!(
363 "[{:02x}:{:02x}:{:02x}:{:02x}:{:02x}:{:02x}:{:02x}:{:02x}]:{}",
364 addr[0], addr[1], addr[2], addr[3], addr[4], addr[5], addr[6],
365 addr[7], port
366 )
367 }
368 _ => return,
369 };
370
371 let target = match tokio::net::TcpStream::connect(&target_addr).await {
372 Ok(t) => t,
373 Err(_) => {
374 let _ = stream
375 .write_all(&[0x05, 0x01, 0x00, 0x01, 0, 0, 0, 0, 0, 0])
376 .await;
377 return;
378 }
379 };
380
381 if stream
382 .write_all(&[0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0])
383 .await
384 .is_err()
385 {
386 return;
387 }
388
389 let (mut cr, mut cw) = stream.into_split();
390 let (mut tr, mut tw) = target.into_split();
391
392 let c2t = tokio::spawn(async move {
393 let _ = tokio::io::copy(&mut cr, &mut tw).await;
394 let _ = tw.shutdown().await;
395 });
396 let t2c = tokio::spawn(async move {
397 let _ = tokio::io::copy(&mut tr, &mut cw).await;
398 let _ = cw.shutdown().await;
399 });
400 let _ = tokio::join!(c2t, t2c);
401 });
402 }
403 });
404
405 Self {
406 addr,
407 connection_count,
408 handle,
409 }
410 }
411
412 pub fn addr(&self) -> SocketAddr {
413 self.addr
414 }
415
416 pub fn connection_count(&self) -> &AtomicU64 {
417 &self.connection_count
418 }
419}
420
421impl Drop for Socks5Upstream {
422 fn drop(&mut self) {
423 self.handle.abort();
424 }
425}
426
427pub struct Socks4Upstream {
428 addr: SocketAddr,
429 connection_count: Arc<AtomicU64>,
430 handle: tokio::task::JoinHandle<()>,
431}
432
433impl Socks4Upstream {
434 pub async fn start() -> Self {
435 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
436 let addr = listener.local_addr().unwrap();
437 let connection_count = Arc::new(AtomicU64::new(0));
438 let cc = connection_count.clone();
439
440 let handle = tokio::spawn(async move {
441 loop {
442 let (mut stream, _) = match listener.accept().await {
443 Ok(s) => s,
444 Err(_) => break,
445 };
446 cc.fetch_add(1, Ordering::Relaxed);
447 tokio::spawn(async move {
448 let mut header = [0u8; 8];
449 if stream.read_exact(&mut header).await.is_err() {
450 return;
451 }
452 let version = header[0];
453 let cmd = header[1];
454 if version != 0x04 || cmd != 0x01 {
455 return;
456 }
457 let port = u16::from_be_bytes([header[2], header[3]]);
458 let ip = [header[4], header[5], header[6], header[7]];
459
460 let mut userid = Vec::new();
461 loop {
462 let mut byte = [0u8; 1];
463 if stream.read_exact(&mut byte).await.is_err() {
464 return;
465 }
466 if byte[0] == 0 {
467 break;
468 }
469 userid.push(byte[0]);
470 }
471 let _ = userid;
472
473 let target_addr = format!("{}.{}.{}.{}:{}", ip[0], ip[1], ip[2], ip[3], port);
474
475 let target = match tokio::net::TcpStream::connect(&target_addr).await {
476 Ok(t) => t,
477 Err(_) => {
478 let _ = stream.write_all(&[0x00, 0x5B, 0, 0, 0, 0, 0, 0]).await;
479 return;
480 }
481 };
482
483 if stream
484 .write_all(&[0x00, 0x5A, 0, 0, 0, 0, 0, 0])
485 .await
486 .is_err()
487 {
488 return;
489 }
490
491 let (mut cr, mut cw) = stream.into_split();
492 let (mut tr, mut tw) = target.into_split();
493
494 let c2t = tokio::spawn(async move {
495 let _ = tokio::io::copy(&mut cr, &mut tw).await;
496 let _ = tw.shutdown().await;
497 });
498 let t2c = tokio::spawn(async move {
499 let _ = tokio::io::copy(&mut tr, &mut cw).await;
500 let _ = cw.shutdown().await;
501 });
502 let _ = tokio::join!(c2t, t2c);
503 });
504 }
505 });
506
507 Self {
508 addr,
509 connection_count,
510 handle,
511 }
512 }
513
514 pub fn addr(&self) -> SocketAddr {
515 self.addr
516 }
517
518 pub fn connection_count(&self) -> &AtomicU64 {
519 &self.connection_count
520 }
521}
522
523impl Drop for Socks4Upstream {
524 fn drop(&mut self) {
525 self.handle.abort();
526 }
527}
528
529pub struct TlsEchoServer {
530 addr: SocketAddr,
531 cert_der: rustls::pki_types::CertificateDer<'static>,
532 handle: tokio::task::JoinHandle<()>,
533}
534
535impl TlsEchoServer {
536 pub async fn start() -> Self {
537 let cert_params = rcgen::CertificateParams::new(vec!["localhost".to_string()]).unwrap();
538 let key_pair = rcgen::KeyPair::generate().unwrap();
539 let cert = cert_params.self_signed(&key_pair).unwrap();
540 let cert_der = cert.der().clone();
541 let key_der = key_pair.serialize_der();
542
543 let server_config = rustls::ServerConfig::builder()
544 .with_no_client_auth()
545 .with_single_cert(
546 vec![cert_der.clone()],
547 rustls::pki_types::PrivatePkcs8KeyDer::from(key_der).into(),
548 )
549 .unwrap();
550 let acceptor = TlsAcceptor::from(Arc::new(server_config));
551
552 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
553 let addr = listener.local_addr().unwrap();
554
555 let handle = tokio::spawn(async move {
556 loop {
557 let (tcp_stream, _) = match listener.accept().await {
558 Ok(s) => s,
559 Err(_) => break,
560 };
561 let acceptor = acceptor.clone();
562 tokio::spawn(async move {
563 let mut stream = match acceptor.accept(tcp_stream).await {
564 Ok(s) => s,
565 Err(_) => return,
566 };
567 let mut buf = [0u8; 4096];
568 loop {
569 match stream.read(&mut buf).await {
570 Ok(0) => break,
571 Ok(n) => {
572 if stream.write_all(&buf[..n]).await.is_err() {
573 break;
574 }
575 }
576 Err(_) => break,
577 }
578 }
579 });
580 }
581 });
582
583 Self {
584 addr,
585 cert_der,
586 handle,
587 }
588 }
589
590 pub fn addr(&self) -> SocketAddr {
591 self.addr
592 }
593
594 pub fn root_store(&self) -> rustls::RootCertStore {
595 let mut store = rustls::RootCertStore::empty();
596 store.add(self.cert_der.clone()).unwrap();
597 store
598 }
599}
600
601impl Drop for TlsEchoServer {
602 fn drop(&mut self) {
603 self.handle.abort();
604 }
605}
606
607pub struct RefusedServer;
608
609impl RefusedServer {
610 pub fn addr() -> SocketAddr {
611 "127.0.0.1:1".parse().unwrap()
612 }
613}