tenant.openagents/omega

1#![cfg(target_os = "windows")]
2
3use std::{os::windows::ffi::OsStringExt, path::PathBuf};
4
5use windows::{
6    Win32::{
7        Foundation::{
8            CLASS_E_CLASSNOTAVAILABLE, E_FAIL, E_INVALIDARG, E_NOTIMPL, ERROR_INSUFFICIENT_BUFFER,
9            GetLastError, HINSTANCE, MAX_PATH,
10        },
11        Globalization::u_strlen,
12        System::{
13            Com::{IBindCtx, IClassFactory, IClassFactory_Impl},
14            LibraryLoader::GetModuleFileNameW,
15            SystemServices::DLL_PROCESS_ATTACH,
16        },
17        UI::Shell::{
18            ECF_DEFAULT, ECS_ENABLED, IEnumExplorerCommand, IExplorerCommand,
19            IExplorerCommand_Impl, IShellItemArray, SHStrDupW, SIGDN_FILESYSPATH,
20        },
21    },
22    core::{BOOL, GUID, HRESULT, HSTRING, Interface, Ref, Result, implement},
23};
24
25static mut DLL_INSTANCE: HINSTANCE = HINSTANCE(std::ptr::null_mut());
26
27#[unsafe(no_mangle)]
28extern "system" fn DllMain(
29    hinstdll: HINSTANCE,
30    fdwreason: u32,
31    _lpvreserved: *mut core::ffi::c_void,
32) -> bool {
33    if fdwreason == DLL_PROCESS_ATTACH {
34        unsafe { DLL_INSTANCE = hinstdll };
35    }
36
37    true
38}
39
40#[implement(IExplorerCommand)]
41struct ExplorerCommandInjector;
42
43#[allow(non_snake_case)]
44impl IExplorerCommand_Impl for ExplorerCommandInjector_Impl {
45    fn GetTitle(&self, _: Ref<IShellItemArray>) -> Result<windows_core::PWSTR> {
46        let command_description =
47            retrieve_command_description().unwrap_or(HSTRING::from("Open with Omega"));
48        unsafe { SHStrDupW(&command_description) }
49    }
50
51    fn GetIcon(&self, _: Ref<IShellItemArray>) -> Result<windows_core::PWSTR> {
52        let Some(zed_exe) = get_zed_exe_path() else {
53            return Err(E_FAIL.into());
54        };
55        unsafe { SHStrDupW(&HSTRING::from(zed_exe)) }
56    }
57
58    fn GetToolTip(&self, _: Ref<IShellItemArray>) -> Result<windows_core::PWSTR> {
59        Err(E_NOTIMPL.into())
60    }
61
62    fn GetCanonicalName(&self) -> Result<windows_core::GUID> {
63        Ok(GUID::zeroed())
64    }
65
66    fn GetState(&self, _: Ref<IShellItemArray>, _: BOOL) -> Result<u32> {
67        Ok(ECS_ENABLED.0 as _)
68    }
69
70    fn Invoke(&self, psiitemarray: Ref<IShellItemArray>, _: Ref<IBindCtx>) -> Result<()> {
71        let items = psiitemarray.ok()?;
72        let Some(zed_exe) = get_zed_exe_path() else {
73            return Ok(());
74        };
75
76        let count = unsafe { items.GetCount()? };
77        for idx in 0..count {
78            let item = unsafe { items.GetItemAt(idx)? };
79            let item_path = unsafe { item.GetDisplayName(SIGDN_FILESYSPATH)?.to_string()? };
80            #[allow(clippy::disallowed_methods, reason = "no async context in sight..")]
81            std::process::Command::new(&zed_exe)
82                .arg(&item_path)
83                .spawn()
84                .map_err(|_| E_INVALIDARG)?;
85        }
86
87        Ok(())
88    }
89
90    fn GetFlags(&self) -> Result<u32> {
91        Ok(ECF_DEFAULT.0 as _)
92    }
93
94    fn EnumSubCommands(&self) -> Result<IEnumExplorerCommand> {
95        Err(E_NOTIMPL.into())
96    }
97}
98
99#[implement(IClassFactory)]
100struct ExplorerCommandInjectorFactory;
101
102impl IClassFactory_Impl for ExplorerCommandInjectorFactory_Impl {
103    fn CreateInstance(
104        &self,
105        punkouter: Ref<windows_core::IUnknown>,
106        riid: *const windows_core::GUID,
107        ppvobject: *mut *mut core::ffi::c_void,
108    ) -> Result<()> {
109        if ppvobject.is_null() || riid.is_null() {
110            return Err(windows::Win32::Foundation::E_POINTER.into());
111        }
112
113        unsafe {
114            *ppvobject = std::ptr::null_mut();
115        }
116
117        if punkouter.is_none() {
118            let factory: IExplorerCommand = ExplorerCommandInjector {}.into();
119            unsafe { factory.query(riid, ppvobject).ok() }
120        } else {
121            Err(E_INVALIDARG.into())
122        }
123    }
124
125    fn LockServer(&self, _: BOOL) -> Result<()> {
126        Ok(())
127    }
128}
129
130#[cfg(all(feature = "stable", not(feature = "preview"), not(feature = "nightly")))]
131const MODULE_ID: GUID = GUID::from_u128(0x6a1f6b13_3b82_48a1_9e06_7bb0a6d0bffd);
132#[cfg(all(feature = "preview", not(feature = "stable"), not(feature = "nightly")))]
133const MODULE_ID: GUID = GUID::from_u128(0xaf8e85ea_fb20_4db2_93cf_56513c1ec697);
134#[cfg(all(feature = "nightly", not(feature = "stable"), not(feature = "preview")))]
135const MODULE_ID: GUID = GUID::from_u128(0x266f2cfe_1653_42af_b55c_fe3590c83871);
136
137// Make cargo clippy happy
138#[cfg(all(feature = "nightly", feature = "stable", feature = "preview"))]
139const MODULE_ID: GUID = GUID::from_u128(0x685f4d49_6718_4c55_b271_ebb5c6a48d6f);
140
141#[unsafe(no_mangle)]
142extern "system" fn DllGetClassObject(
143    class_id: *const GUID,
144    iid: *const GUID,
145    out: *mut *mut std::ffi::c_void,
146) -> HRESULT {
147    if out.is_null() || class_id.is_null() || iid.is_null() {
148        return E_INVALIDARG;
149    }
150
151    unsafe {
152        *out = std::ptr::null_mut();
153    }
154    let class_id = unsafe { *class_id };
155    if class_id == MODULE_ID {
156        let instance: IClassFactory = ExplorerCommandInjectorFactory {}.into();
157        unsafe { instance.query(iid, out) }
158    } else {
159        CLASS_E_CLASSNOTAVAILABLE
160    }
161}
162
163fn get_zed_install_folder() -> Option<PathBuf> {
164    let mut buf = vec![0u16; MAX_PATH as usize];
165    unsafe { GetModuleFileNameW(Some(DLL_INSTANCE.into()), &mut buf) };
166
167    while unsafe { GetLastError() } == ERROR_INSUFFICIENT_BUFFER {
168        buf = vec![0u16; buf.len() * 2];
169        unsafe { GetModuleFileNameW(Some(DLL_INSTANCE.into()), &mut buf) };
170    }
171    let len = unsafe { u_strlen(buf.as_ptr()) };
172    let path: PathBuf = std::ffi::OsString::from_wide(&buf[..len as usize])
173        .into_string()
174        .ok()?
175        .into();
176    Some(path.parent()?.parent()?.to_path_buf())
177}
178
179#[inline]
180fn get_zed_exe_path() -> Option<String> {
181    get_zed_install_folder().map(|path| path.join("Zed.exe").to_string_lossy().into_owned())
182}
183
184#[inline]
185fn retrieve_command_description() -> Result<HSTRING> {
186    #[cfg(all(feature = "stable", not(feature = "preview"), not(feature = "nightly")))]
187    const REG_PATH: &str = "Software\\Classes\\ZedEditorContextMenu";
188    #[cfg(all(feature = "preview", not(feature = "stable"), not(feature = "nightly")))]
189    const REG_PATH: &str = "Software\\Classes\\ZedEditorPreviewContextMenu";
190    #[cfg(all(feature = "nightly", not(feature = "stable"), not(feature = "preview")))]
191    const REG_PATH: &str = "Software\\Classes\\ZedEditorNightlyContextMenu";
192
193    // Make cargo clippy happy
194    #[cfg(all(feature = "nightly", feature = "stable", feature = "preview"))]
195    const REG_PATH: &str = "Software\\Classes\\ZedEditorClippyContextMenu";
196
197    let key = windows_registry::CURRENT_USER.open(REG_PATH)?;
198    key.get_hstring("Title")
199}
200