EvanZhu0721 / 公开仓库

SBMS

镜像已同步

SBMS complete the calculation rules for the default logical desktop topology relationship and the physical dimensions of the displays in the Windows multi-screen collaboration system.

RustGPL-3.0main来源 commit a37d6c4b2026/10/15 07:25:26 同步

上游 GitHub

src/launch_broker.rs10.3 KB

1use std::error::Error;

2use std::ffi::OsString;

3use std::fmt::{Display, Formatter};

4use std::mem::size_of;

5use std::os::windows::process::CommandExt;

6use std::path::PathBuf;

7use std::process::{Command, Stdio};

8use std::thread;

9use std::time::{Duration, Instant};

11use windows::Win32::Foundation::{CloseHandle, HANDLE, HWND};

12use windows::Win32::Security::{GetTokenInformation, TOKEN_ELEVATION, TOKEN_QUERY, TokenElevation};

13use windows::Win32::System::Threading::{

14 GetCurrentProcess, OpenProcess, OpenProcessToken, PROCESS_QUERY_LIMITED_INFORMATION,

15};

16use windows::Win32::UI::WindowsAndMessaging::{

17 GetWindowThreadProcessId, MB_ICONERROR, MB_OK, MessageBoxW,

18};

19use windows::core::{PCWSTR, w};

21use crate::control::TrayInstance;

22use crate::diagnostics::{self, Level};

23use crate::win32_tray::{find_host_window, request_open};

25pub const REGISTERED_TASK_NAME: &str = r"\SBMS\Tray-7EB4D7A8-16A9-4B6F-82E3-31A77BC81B6A";

26const CREATE_NO_WINDOW: u32 = 0x0800_0000;

27const START_TIMEOUT: Duration = Duration::from_secs(10);

28const POLL_INTERVAL: Duration = Duration::from_millis(100);

30pub enum LaunchDisposition {

31 RunHere(TrayInstance),

32 Exit,

33}

35enum HostClaim<H, C> {

36 Existing(H),

37 Claimed(C),

38}

40#[derive(Debug)]

41pub struct LaunchError(String);

43impl Display for LaunchError {

44 fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {

45 formatter.write_str(&self.0)

46 }

47}

49impl Error for LaunchError {}

51pub fn route_tray(open_requested: bool) -> Result<LaunchDisposition, LaunchError> {

52 let elevated = current_process_is_elevated()?;

53 if let Some(hwnd) = find_host_window() {

54 return route_to_existing(hwnd, open_requested);

55 }

57 if elevated {

58 return claim_or_route_elevated(open_requested);

59 }

61 diagnostics::log(

62 Level::Info,

63 "launch",

64 "task-handoff",

65 None,

66 "starting the registered elevated tray task",

67 );

68 start_registered_task()?;

70 let deadline = Instant::now() + START_TIMEOUT;

71 loop {

72 if let Some(hwnd) = find_host_window() {

73 return route_to_existing(hwnd, open_requested);

74 }

75 if Instant::now() >= deadline {

76 return Err(LaunchError(

77 "The elevated SBMS tray did not become ready. Repair or reinstall SBMS, then try again."

78 .into(),

79 ));

80 }

81 thread::sleep(POLL_INTERVAL);

82 }

83}

85pub fn tray_open_requested(

86 arguments: impl IntoIterator<Item = OsString>,

87) -> Result<bool, LaunchError> {

88 let mut arguments = arguments.into_iter();

89 let open_requested = match arguments.next() {

90 None => true,

91 Some(argument) if argument == "--open" => true,

92 Some(argument) if argument == "--background" => false,

93 Some(argument) => {

94 return Err(LaunchError(format!(

95 "unsupported SBMS tray argument: {}",

96 argument.to_string_lossy()

97 )));

98 }

99 };

100 if let Some(argument) = arguments.next() {

101 return Err(LaunchError(format!(

102 "unsupported SBMS tray argument: {}",

103 argument.to_string_lossy()

104 )));

105 }

106 Ok(open_requested)

107}

108

109pub fn show_launch_error(error: &dyn Error) {

110 let message = wide(&format!("SBMS could not start.\n\n{error}"));

111 unsafe {

112 let _ = MessageBoxW(

113 None,

114 PCWSTR(message.as_ptr()),

115 w!("SBMS"),

116 MB_OK | MB_ICONERROR,

117 );

118 }

119}

120

121fn route_to_existing(hwnd: HWND, open_requested: bool) -> Result<LaunchDisposition, LaunchError> {

122 if !window_process_is_elevated(hwnd)? {

123 return Err(LaunchError(

124 "A non-elevated SBMS tray is already running. Exit it from the tray menu, then open SBMS again."

125 .into(),

126 ));

127 }

128 if open_requested {

129 request_open(hwnd).map_err(|error| {

130 LaunchError(format!("Could not open the running SBMS tray: {error}"))

131 })?;

132 }

133 Ok(LaunchDisposition::Exit)

134}

135

136fn claim_or_route_elevated(open_requested: bool) -> Result<LaunchDisposition, LaunchError> {

137 let deadline = Instant::now() + START_TIMEOUT;

138 let claim = poll_host_or_claim(

139 find_host_window,

140 || TrayInstance::acquire().map_err(LaunchError),

141 || {

142 if Instant::now() >= deadline {

143 return false;

144 }

145 thread::sleep(POLL_INTERVAL);

146 true

147 },

148 )?;

149 match claim {

150 Some(HostClaim::Existing(hwnd)) => route_to_existing(hwnd, open_requested),

151 Some(HostClaim::Claimed(instance)) => {

152 diagnostics::log(

153 Level::Debug,

154 "launch",

155 "route",

156 None,

157 "elevated tray process claimed the UI host",

158 );

159 Ok(LaunchDisposition::RunHere(instance))

160 }

161 None => Err(LaunchError(

162

            "Another elevated SBMS tray owns the host lock but did not become ready. Repair or reinstall SBMS, then try again."

163 .into(),

164 )),

165 }

166}

167

168fn poll_host_or_claim<H, C, E>(

169 mut find_host: impl FnMut() -> Option<H>,

170 mut try_claim: impl FnMut() -> Result<Option<C>, E>,

171 mut wait_for_retry: impl FnMut() -> bool,

172) -> Result<Option<HostClaim<H, C>>, E> {

173 loop {

174 if let Some(host) = find_host() {

175 return Ok(Some(HostClaim::Existing(host)));

176 }

177 if let Some(claim) = try_claim()? {

178 return Ok(Some(HostClaim::Claimed(claim)));

179 }

180 if !wait_for_retry() {

181 return Ok(None);

182 }

183 }

184}

185

186fn start_registered_task() -> Result<(), LaunchError> {

187 let system_root = std::env::var_os("SystemRoot")

188 .filter(|value| !value.is_empty())

189 .ok_or_else(|| LaunchError("Windows did not provide the SystemRoot path.".into()))?;

190 let executable = PathBuf::from(system_root)

191 .join("System32")

192 .join("schtasks.exe");

193 let status = Command::new(&executable)

194 .args(["/Run", "/TN", REGISTERED_TASK_NAME])

195 .creation_flags(CREATE_NO_WINDOW)

196 .stdin(Stdio::null())

197 .stdout(Stdio::null())

198 .stderr(Stdio::null())

199 .status()

200 .map_err(|error| {

201 LaunchError(format!(

202 "Could not start the registered SBMS task at {}: {error}",

203 executable.display()

204 ))

205 })?;

206 if !status.success() {

207 return Err(LaunchError(format!(

208 "The registered SBMS startup task failed with exit code {}. Repair or reinstall SBMS.",

209 status.code().unwrap_or(-1)

210 )));

211 }

212 Ok(())

213}

214

215fn current_process_is_elevated() -> Result<bool, LaunchError> {

216 token_is_elevated(unsafe { GetCurrentProcess() })

217}

218

219fn window_process_is_elevated(hwnd: HWND) -> Result<bool, LaunchError> {

220 let mut process_id = 0;

221 unsafe {

222 GetWindowThreadProcessId(hwnd, Some(&mut process_id));

223 }

224 if process_id == 0 {

225 return Err(LaunchError(

226 "Could not identify the running SBMS tray process.".into(),

227 ));

228 }

229 let process = unsafe { OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, false, process_id) }

230 .map_err(|error| {

231 LaunchError(format!("Could not inspect the running SBMS tray: {error}"))

232 })?;

233 let result = token_is_elevated(process);

234 unsafe {

235 let _ = CloseHandle(process);

236 }

237 result

238}

239

240fn token_is_elevated(process: HANDLE) -> Result<bool, LaunchError> {

241 let mut token = HANDLE::default();

242 unsafe { OpenProcessToken(process, TOKEN_QUERY, &mut token) }.map_err(|error| {

243 LaunchError(format!("Could not inspect the SBMS process token: {error}"))

244 })?;

245 let mut elevation = TOKEN_ELEVATION::default();

246 let mut returned = 0;

247 let result = unsafe {

248 GetTokenInformation(

249 token,

250 TokenElevation,

251 Some((&mut elevation as *mut TOKEN_ELEVATION).cast()),

252 size_of::<TOKEN_ELEVATION>() as u32,

253 &mut returned,

254 )

255 };

256 unsafe {

257 let _ = CloseHandle(token);

258 }

259 result

260 .map(|()| elevation.TokenIsElevated != 0)

261 .map_err(|error| LaunchError(format!("Could not read the SBMS elevation state: {error}")))

262}

263

264fn wide(value: &str) -> Vec<u16> {

265 let mut encoded: Vec<u16> = value.encode_utf16().collect();

266 encoded.push(0);

267 encoded

268}

269

270#[cfg(test)]

271mod tests {

272 use super::*;

273

274 #[test]

275 fn registered_task_name_is_stable() {

276 assert_eq!(

277 REGISTERED_TASK_NAME,

278 r"\SBMS\Tray-7EB4D7A8-16A9-4B6F-82E3-31A77BC81B6A"

279 );

280 }

281

282 #[test]

283 fn tray_arguments_default_to_open() {

284 assert!(tray_open_requested([]).unwrap());

285 assert!(tray_open_requested(["--open".into()]).unwrap());

286 assert!(!tray_open_requested(["--background".into()]).unwrap());

287 }

288

289 #[test]

290 fn tray_arguments_reject_unknown_or_extra_values() {

291 assert!(tray_open_requested(["--unknown".into()]).is_err());

292 assert!(tray_open_requested(["--open".into(), "extra".into()]).is_err());

293 }

294

295 #[test]

296 fn polling_routes_to_host_when_busy_owner_becomes_ready() {

297 let mut host_checks = 0;

298 let mut claim_attempts = 0;

299 let mut waits = 0;

300 let result = poll_host_or_claim(

301 || {

302 host_checks += 1;

303 (host_checks == 2).then_some("host")

304 },

305 || {

306 claim_attempts += 1;

307 Ok::<_, ()>(None::<&str>)

308 },

309 || {

310 waits += 1;

311 true

312 },

313 )

314 .unwrap();

315

316 assert!(matches!(result, Some(HostClaim::Existing("host"))));

317 assert_eq!(claim_attempts, 1);

318 assert_eq!(waits, 1);

319 }

320

321 #[test]

322 fn polling_claims_host_when_busy_mutex_is_released() {

323 let mut claim_attempts = 0;

324 let mut waits = 0;

325 let result = poll_host_or_claim(

326 || None::<&str>,

327 || {

328 claim_attempts += 1;

329 Ok::<_, ()>((claim_attempts == 2).then_some("claim"))

330 },

331 || {

332 waits += 1;

333 true

334 },

335 )

336 .unwrap();

337

338 assert!(matches!(result, Some(HostClaim::Claimed("claim"))));

339 assert_eq!(waits, 1);

340 }

341

342 #[test]

343 fn wide_adds_exactly_one_terminator() {

344 assert_eq!(wide("SBMS"), [83, 66, 77, 83, 0]);

345 }

346}

347