bim_core/clients/
base.rs

1use std::io::{Read, Write};
2use std::net::{SocketAddr, TcpStream};
3use std::sync::Arc;
4use std::thread;
5use std::time::Duration;
6
7#[cfg(debug_assertions)]
8use log::debug;
9
10use rustls::{OwnedTrustAnchor, RootCertStore};
11use url::Url;
12
13use crate::utils::SpeedTestResult;
14
15pub trait GenericStream: Read + Write {}
16
17impl<T: Read + Write> GenericStream for T {}
18
19pub fn make_connection(
20    address: &SocketAddr,
21    url: &Url,
22) -> Result<Box<dyn GenericStream>, String> {
23    let ssl = if url.scheme() == "https" { true } else { false };
24    let mut retry = 3;
25
26    while retry > 0 {
27        if let Ok(stream) = TcpStream::connect_timeout(&address, Duration::from_micros(3_000_000)) {
28            #[cfg(debug_assertions)]
29            debug!("TCP connected");
30
31            let _r = stream.set_write_timeout(Some(Duration::from_secs(3)));
32            let _r = stream.set_read_timeout(Some(Duration::from_secs(3)));
33            if !ssl {
34                return Ok(Box::new(stream));
35            }
36
37            let mut root_store = RootCertStore::empty();
38            root_store.add_server_trust_anchors(webpki_roots::TLS_SERVER_ROOTS.0.iter().map(
39                |ta| {
40                    OwnedTrustAnchor::from_subject_spki_name_constraints(
41                        ta.subject,
42                        ta.spki,
43                        ta.name_constraints,
44                    )
45                },
46            ));
47            let config = rustls::ClientConfig::builder()
48                .with_safe_defaults()
49                .with_root_certificates(root_store)
50                .with_no_client_auth();
51
52            let server_name = url.host_str().unwrap().try_into().unwrap();
53            let conn = rustls::ClientConnection::new(Arc::new(config), server_name).unwrap();
54            let tls = rustls::StreamOwned::new(conn, stream);
55
56            #[cfg(debug_assertions)]
57            debug!("SSL connected");
58
59            return Ok(Box::new(tls));
60        }
61
62        retry -= 1;
63    }
64    return Err(String::from("连接失败"));
65}
66
67pub trait Client {
68    fn result(&self) -> SpeedTestResult;
69
70    fn ping(&mut self) -> bool;
71
72    fn upload(&mut self) -> bool;
73
74    fn download(&mut self) -> bool;
75
76    fn run(&mut self) -> bool {
77        let r = self.ping();
78        if r {
79            thread::sleep(Duration::from_secs(2));
80            self.upload();
81            thread::sleep(Duration::from_secs(3));
82            self.download();
83            return true;
84        }
85        false
86    }
87}