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

use windows::Win32::{
    Foundation::HANDLE,
    System::{
        Diagnostics::ToolHelp::{
            CreateToolhelp32Snapshot, Thread32First, Thread32Next, TH32CS_SNAPTHREAD, THREADENTRY32,
        },
        Threading::THREAD_ACCESS_RIGHTS,
    },
};
use crate::{size_of, FaitheError};
use super::OwnedThread;

/// Represents single running thread in a process.
#[derive(Debug, Clone)]
pub struct ThreadEntry {
    /// Id of the process this thread is running in.
    pub process_id: u32,
    /// Id of the thread.
    pub thread_id: u32,
    /// Priority of the thread.
    pub base_priority: i32,
}

impl ThreadEntry {
    /// Opens thread
    pub fn open(
        &self,
        inherit_handle: bool,
        desired_access: THREAD_ACCESS_RIGHTS,
    ) -> crate::Result<OwnedThread> {
        OwnedThread::open(self.thread_id, inherit_handle, desired_access)
    }
}

impl From<THREADENTRY32> for ThreadEntry {
    fn from(te: THREADENTRY32) -> Self {
        Self {
            process_id: te.th32OwnerProcessID,
            thread_id: te.th32ThreadID,
            base_priority: te.tpBasePri,
        }
    }
}

/// Iterator over running threads in the process.
pub struct ThreadIterator {
    snapshot: HANDLE,
    entry: THREADENTRY32,
    process_id: u32,
    should_return: bool,
}

impl ThreadIterator {
    /// Creates new iterator over threads in process with id `process_id`.
    pub fn new(process_id: u32) -> crate::Result<Self> {
        unsafe {
            let snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, process_id)
                .map_err(|_| FaitheError::last_error())?;

            let entry = THREADENTRY32 {
                dwSize: size_of!(THREADENTRY32) as _,
                ..Default::default()
            };

            let mut this = Self {
                should_return: false,
                process_id,
                snapshot,
                entry,
            };

            if Thread32First(snapshot, &mut this.entry) == false {
                Err(FaitheError::last_error())
            } else {
                Ok(this)
            }
        }
    }
}

impl Iterator for ThreadIterator {
    type Item = ThreadEntry;

    fn next(&mut self) -> Option<Self::Item> {
        if self.should_return {
            None
        } else {
            let mut this: ThreadEntry = self.entry.into();

            unsafe {
                self.should_return = Thread32Next(self.snapshot, &mut self.entry) == false;

                while this.process_id != self.process_id {
                    this = self.entry.into();
                    self.should_return = Thread32Next(self.snapshot, &mut self.entry) == false;
                    if self.should_return {
                        return None;
                    }
                }
            }

            Some(this)
        }
    }
}