Trtyr's Blog

1
use std::net::{IpAddr, SocketAddr};
2
use std::time::Duration;
3

4
#[tokio::main]
5
async fn main() {
6
    // 初始化数据
7
    // 这里用 unwrap 是因为 IP 是硬编码的，如果这都挂了，说明写代码的人脑子进水了
8
    // 如果实际工具你要是敢用 unwrap，掐死你！
9
    let host = "127.0.0.1".parse::<IpAddr>().unwrap();
10
    let posts: Vec<u16> = (1..=1000).collect(); // 测试端口
11

12
    // 初始化扫描器
13
    // 扫描器应该是一个单独的对象
14
    // 结构体属性有这些
15
    // - address(IpAddr): 这是要目标地址
16
    // - posts(Vec<u16>): 这是要处理的目标端口，这里同理，也应该是引用
17
    // - batch_size(usize): 因为我们设计了一个任务池，这里要给一个池子大小
18
    // - timeout: 扫描器进行扫描的超时时间
19
    let scanner = Scanner::new(
20
        host,
21
        posts,
22
        1000,                        // 池子大小设置为 1000
23
        Duration::from_millis(1500), // 超时设置为 1.5
24
    );
25

26
    // 我们现在拿到了一个扫描器对象，这个对象应该自己有一些异步方法
27
    // 我们创建一个 run 方法，这个 run 方法会返回所有开放的 IP:Port，也就是 Vec<SocketAddr>
28
    let results = scanner.run().await;
29

30
    // 现在我们拿到了 Vec<SocketAddr>，直接遍历取出就行
31
    for socket in results {
32
        println!("{}", socket);
33
    }
34
}
35

36
struct Scanner {
37
    host: IpAddr,
38
    ports: Vec<u16>,
39
    batch_size: usize,
40
    timeout: Duration,
41
}
42

43
impl Scanner {
44
    fn new(host: IpAddr, ports: Vec<u16>, batch_size: usize, timeout: Duration) -> Self {
45
        Self {
46
            host,
47
            ports,
48
            batch_size,
49
            timeout,
50
        }
51
    }
52

53
    async fn run(&self) -> Vec<SocketAddr> {
54
        let result: Vec<SocketAddr> = Vec::new();
55

56
        result
57
    }
58
}

1
async fn run(&self) -> Vec<SocketAddr> {
2
        let mut result: Vec<SocketAddr> = Vec::new(); // 创建一个空的缓冲区，放最后的结果
3
        let mut tasks = FuturesUnordered::new(); // 创建任务池
4
        let mut port_iterator = self.ports.iter(); // 创建一个 ports 迭代器
5

6
        // 进行预填充，先放 1000 个任务对象
7
        for _ in 0..self.batch_size {
8
            // 迭代器的 next() 方法，返回的是一个 Some(T)
9
            if let Some(&port) = port_iterator.next() {
10
                let timeout = self.timeout; // 超时时间
11
                // 拿到 port 后，构建 socket 对象(ScoketAddr)
12
                let socket = SocketAddr::new(self.host, port);
13
                // 将动作封装成一个动作包（scan_worker），push进 task 里
14
                tasks.push(Scanner::scan_worker(socket, timeout));
15
            } else {
16
                // 如果端口还没填满 batch_size 就没了，直接跳出
17
                break;
18
            }
19
        }
20

21
        // 执行，跑完一个，造一个，放进去一个
22
        // tasks.next().await 会等待任意一个 future 完成
23
        // scan_result 是 Option<SocketAddr> 类型
24
        while let Some(scan_result) = tasks.next().await {
25
            // 把 open_socket 给拿出来
26
            if let Some(open_socket) = scan_result {
27
                // 这就是结果了，直接 push 进 result 里
28
                result.push(open_socket);
29
            }
30

31
            // 完成了一个，task 就有了一个空位，从迭代器里生成一个放进去
32
            if let Some(&port) = port_iterator.next() {
33
                let socket = SocketAddr::new(self.host, port);
34
                let timeout = self.timeout;
35
                tasks.push(Scanner::scan_worker(socket, timeout));
36
            }
37
        }
38

39
        result
40
    }

1
async fn scan_worker(socket: SocketAddr, timeout: Duration) -> Option<SocketAddr> {
2
        if Scanner::socket_scan(socket, timeout).await {
3
            Some(socket)
4
        } else {
5
            None
6
        }
7
    }

1
// 端口扫描的核心组件
2
    async fn socket_scan(socket: SocketAddr, timeout: Duration) -> bool {
3
        // tokio::time::timeout 接收一个 Duration 和一个 Future
4
        // 它会把 Future 包一层，如果超时了就返回 Err(Elapsed)
5
        match tokio::time::timeout(timeout, TcpStream::connect(socket)).await {
6
            Ok(Ok(_)) => true,
7
            _ => false,
8
        }
9
    }