Skip to main content

ib_shell_item/hook/prop/
system.rs

1#![allow(non_upper_case_globals)]
2use std::{cell::SyncUnsafeCell, ffi::c_void};
3
4use bon::Builder;
5use ib_hook::inline::InlineHook;
6use serde::{Deserialize, Serialize};
7use tracing::trace;
8use widestring::U16CStr;
9use windows::{
10    Win32::{
11        Foundation::{PROPERTYKEY, S_OK},
12        Storage::EnhancedStorage::PKEY_Size,
13        System::Com::StructuredStorage::PROPVARIANT,
14        UI::Shell::PropertiesSystem::{IPropertySystem, PDFF_ALWAYSKB, PROPDESC_FORMAT_FLAGS},
15    },
16    core::{HRESULT, Interface, PWSTR},
17};
18
19use crate::{hook::prop::magic, prop::system::PropertySystem, string::prefix_u16cstr_ptr};
20
21#[derive(Default, Serialize, Deserialize, Clone, Builder, Debug)]
22pub struct PropertySystemHookConfig {
23    #[builder(default)]
24    size_no_alwayskb: bool,
25    #[builder(default)]
26    size_max_bar: bool,
27}
28
29type FormatForDisplayAlloc = unsafe extern "system" fn(
30    *mut c_void,
31    *const PROPERTYKEY,
32    *const PROPVARIANT,
33    PROPDESC_FORMAT_FLAGS,
34    *mut PWSTR,
35) -> HRESULT;
36
37pub(crate) struct PropertySystemHook {
38    format_for_display_alloc: Option<InlineHook<FormatForDisplayAlloc>>,
39    config: PropertySystemHookConfig,
40}
41
42impl PropertySystemHook {
43    pub fn new(config: PropertySystemHookConfig) -> anyhow::Result<Self> {
44        let system = IPropertySystem::new()?;
45
46        let format_for_display_alloc = if config.size_max_bar || config.size_no_alwayskb {
47            let format_for_display_alloc = system.vtable().FormatForDisplayAlloc;
48            InlineHook::new_enabled(format_for_display_alloc, format_for_display_alloc_detour).ok()
49        } else {
50            None
51        };
52
53        Ok(Self {
54            format_for_display_alloc,
55            config,
56        })
57    }
58}
59
60pub(crate) static HOOK: SyncUnsafeCell<Option<PropertySystemHook>> = SyncUnsafeCell::new(None);
61
62pub fn apply(config: Option<PropertySystemHookConfig>) -> anyhow::Result<()> {
63    let hook = unsafe { &mut *HOOK.get() };
64    *hook = match config {
65        Some(config) => Some(PropertySystemHook::new(config)?),
66        None => None,
67    };
68    Ok(())
69}
70
71#[allow(non_snake_case)]
72unsafe extern "system" fn format_for_display_alloc_detour(
73    this: *mut c_void,
74    key: *const PROPERTYKEY,
75    propvar: *const PROPVARIANT,
76    pdff: PROPDESC_FORMAT_FLAGS,
77    ppszDisplay: *mut PWSTR,
78) -> HRESULT {
79    trace!(?key, ?propvar, ?pdff);
80    let hook = unsafe { &*HOOK.get() }.as_ref().unwrap();
81    let trampoline = hook.format_for_display_alloc.as_ref().unwrap().trampoline();
82    let real = || unsafe { trampoline(this, key, propvar, pdff, ppszDisplay) };
83
84    // let system = unsafe { IPropertySystem::from_raw_borrowed(&this) }.unwrap();
85    let pkey = unsafe { &*key };
86    match *pkey {
87        PKEY_Size => {
88            if hook.config.size_no_alwayskb || hook.config.size_max_bar {
89                let r = unsafe {
90                    trampoline(
91                        this,
92                        key,
93                        propvar,
94                        if hook.config.size_no_alwayskb {
95                            pdff & !PDFF_ALWAYSKB
96                        } else {
97                            pdff
98                        },
99                        ppszDisplay,
100                    )
101                };
102                if r.is_err() {
103                    return r;
104                }
105            }
106            if hook.config.size_max_bar
107                && let Some(size) = TryInto::<u64>::try_into(unsafe { &*propvar }).ok()
108                && let Some(max) = magic::ui8_read_u64(propvar)
109            {
110                let s = unsafe { U16CStr::from_ptr_str((*ppszDisplay).0) };
111
112                /*
113                // If too wide, the size column will be truncated from unit to numbers...
114                let bar_width = 25;
115
116                let n = if max == 0 {
117                    1
118                } else {
119                    std::cmp::min((size * bar_width / max) + 1, bar_width)
120                };
121                // Fortunately, the default font is not monospaced so making bar is easy.
122                // 1px per char.
123                // Unfortunately, we also can't overlap bar with size string easily.
124                // let bar = "\u{258F}".repeat(n as usize).to_string() + " ";
125                let mut bar = widestring::utf16str!("\u{258F}").repeat(n as usize);
126                */
127                let mut bar = crate::string::bar::StringBar::builder()
128                    .value(size)
129                    .max(max)
130                    // .width(bar_width)
131                    .min_bar(true)
132                    .build()
133                    .to_utf16_string();
134                bar.push(' ');
135                // bar.push(crate::string::bar::HAIR_SPACE);
136                // bar.insert(0, ' ');
137                // debug!(?bar, ?s);
138
139                // let s = format!("{bar} {:.0} KB", size / 1024);
140                // unsafe { *ppszDisplay = str_to_co_task(&s) };
141                // let prefix: Vec<u16> = bar.encode_utf16().collect();
142                let prefix = bar.as_slice();
143                unsafe { *ppszDisplay = prefix_u16cstr_ptr(s, &prefix) };
144                // unsafe { *ppszDisplay = suffix_u16cstr_ptr(s, &prefix) };
145            };
146            return S_OK;
147        }
148        _ => (),
149    }
150
151    real()
152}