Inspired by balena.io's drivelist, I create a Rust version. The idea is simple, using Windows' Setup API to read all connected storages list (with mount points/drive letters), using winapi crate
///file main.rs
use core::{slice, ffi};
use std::{ptr::{null_mut, addr_of, null}, mem::{zeroed, size_of, MaybeUninit, align_of, transmute}, str::from_utf8};
use winapi::{um::{setupapi::{SetupDiGetClassDevsA, DIGCF_PRESENT, DIGCF_DEVICEINTERFACE, SP_DEVINFO_DATA, SetupDiEnumDeviceInfo, HDEVINFO, PSP_DEVINFO_DATA, SetupDiGetDeviceRegistryPropertyW, SPDRP_FRIENDLYNAME, SPDRP_REMOVAL_POLICY, SPDRP_ENUMERATOR_NAME, SetupDiDestroyDeviceInfoList, SP_DEVICE_INTERFACE_DATA, SetupDiEnumDeviceInterfaces, SetupDiGetDeviceInterfaceDetailW, PSP_DEVICE_INTERFACE_DETAIL_DATA_W, SP_DEVICE_INTERFACE_DETAIL_DATA_W}, winioctl::{GUID_DEVINTERFACE_DISK, VOLUME_DISK_EXTENTS, IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS, PVOLUME_DISK_EXTENTS}, handleapi::{INVALID_HANDLE_VALUE, CloseHandle}, cfgmgr32::{CM_REMOVAL_POLICY_EXPECT_SURPRISE_REMOVAL, CM_REMOVAL_POLICY_EXPECT_ORDERLY_REMOVAL}, errhandlingapi::GetLastError, fileapi::{CreateFileW, OPEN_EXISTING}, winnt::{FILE_SHARE_READ, FILE_SHARE_WRITE, FILE_ATTRIBUTE_NORMAL}, ioapiset::DeviceIoControl}, shared::{minwindef::MAX_PATH, winerror::{ERROR_NO_MORE_ITEMS, ERROR_INSUFFICIENT_BUFFER}}, ctypes::c_void};
fn get_detail_data(h_dev_info:HDEVINFO,device_info_data:PSP_DEVINFO_DATA)
{
let mut h_device=INVALID_HANDLE_VALUE;
let mut index=0_u32;
unsafe{
loop {
if h_device!= INVALID_HANDLE_VALUE {
CloseHandle(h_device);
h_device=INVALID_HANDLE_VALUE;
}
let mut device_interface_data:SP_DEVICE_INTERFACE_DATA=zeroed();
device_interface_data.cbSize=size_of::<SP_DEVICE_INTERFACE_DATA>() as _;
if SetupDiEnumDeviceInterfaces(h_dev_info, device_info_data, &GUID_DEVINTERFACE_DISK, index, &mut device_interface_data) == 0 {
let error_code=GetLastError();
if error_code!=ERROR_NO_MORE_ITEMS {
panic!("SetupDiEnumDeviceInterfaces: Error {}",error_code);
}
break;
} else {
let mut size={
let mut required_size=MaybeUninit::<u32>::uninit();
if SetupDiGetDeviceInterfaceDetailW(h_dev_info, &mut device_interface_data, null_mut(), 0, required_size.as_mut_ptr(), null_mut())==0 {
if GetLastError()==ERROR_INSUFFICIENT_BUFFER {
required_size.assume_init()
} else {
panic!("Error SetupDiGetDeviceInterfaceDetailW");
}
} else {
0
}
};
let mut buf:Vec<u8>=Vec::with_capacity(TryInto::<usize>::try_into(size).unwrap() + align_of::<SP_DEVICE_INTERFACE_DETAIL_DATA_W>()-1);
let align_offset=buf.as_mut_ptr().align_offset(align_of::<SP_DEVICE_INTERFACE_DETAIL_DATA_W>());
let device_iface_detail =&mut *(buf.as_mut_ptr().offset(align_offset.try_into().unwrap()) as *mut MaybeUninit<SP_DEVICE_INTERFACE_DETAIL_DATA_W>);
device_iface_detail.write(SP_DEVICE_INTERFACE_DETAIL_DATA_W {
cbSize: size_of::<SP_DEVICE_INTERFACE_DETAIL_DATA_W>().try_into().unwrap(),
DevicePath: [0],
});
if SetupDiGetDeviceInterfaceDetailW(h_dev_info, &mut device_interface_data, device_iface_detail.as_mut_ptr(), size, &mut size, null_mut())==0 {
println!("Error {}, Couldn't SetupDiGetDeviceInterfaceDetailW",GetLastError());
break;
}
let device_detail_data=device_iface_detail.assume_init_ref();
h_device=CreateFileW(device_detail_data.DevicePath.as_ptr(), 0, FILE_SHARE_READ, null_mut(), OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, null_mut());
if h_device==INVALID_HANDLE_VALUE {
println!("Couldn't open handle to device: Error {}",GetLastError());
break;
}
get_device_number(h_device);
}
index+=1;
}
if h_device!= INVALID_HANDLE_VALUE {
CloseHandle(h_device);
h_device=INVALID_HANDLE_VALUE;
}
}
}
fn get_device_number(h_device:*mut c_void)
{
unsafe {
let mut size=0_u32;
let mut disk_extents=MaybeUninit::<VOLUME_DISK_EXTENTS>::uninit();
disk_extents.write(zeroed());
let result=DeviceIoControl(h_device, IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS, null_mut(), 0, disk_extents.as_mut_ptr() as _, size_of::<VOLUME_DISK_EXTENTS>() as _, &mut size, null_mut());
/* It returns 0!! */
println!("Bytes returned: {}",size);
if result!=0 {
println!("Success");
} else {
/* This will be executed */
println!("get_device_number fail. Error {}",GetLastError());
}
}
}
fn get_enumerator_name(h_dev_info:HDEVINFO,device_info_data:PSP_DEVINFO_DATA) -> String
{
unsafe {
let mut buffer:[u8;MAX_PATH]=zeroed();
if SetupDiGetDeviceRegistryPropertyW(h_dev_info, device_info_data, SPDRP_ENUMERATOR_NAME, null_mut(), &mut buffer as _, (size_of::<u8>() * MAX_PATH) as _, null_mut()) != 0 {
ansi_to_string(&buffer)
} else {
"".to_string()
}
}
}
fn get_friendly_name(h_dev_info:HDEVINFO,device_info_data:PSP_DEVINFO_DATA) -> String
{
unsafe {
let mut buffer:[u8;MAX_PATH]=zeroed();
if SetupDiGetDeviceRegistryPropertyW(h_dev_info, device_info_data, SPDRP_FRIENDLYNAME, null_mut(), &mut buffer as _, (size_of::<u8>() * MAX_PATH) as _, null_mut()) != 0 {
ansi_to_string(&buffer)
} else {
"".to_string()
}
}
}
fn is_removable(h_dev_info:HDEVINFO,device_info_data:PSP_DEVINFO_DATA)->bool
{
unsafe {
let mut result=0_u8;
SetupDiGetDeviceRegistryPropertyW(h_dev_info, device_info_data, SPDRP_REMOVAL_POLICY, null_mut(), &mut result as _, size_of::<u32>() as _, null_mut());
match result as u32
{
CM_REMOVAL_POLICY_EXPECT_SURPRISE_REMOVAL|CM_REMOVAL_POLICY_EXPECT_ORDERLY_REMOVAL =>true,
_=>false
}
}
}
fn is_usb_drive(enumerator_name:&str) -> bool
{
["USBSTOR", "UASPSTOR", "VUSBSTOR","RTUSER", "CMIUCR", "EUCR","ETRONSTOR", "ASUSSTPT"].contains(&enumerator_name)
}
fn ansi_to_string(unsafe_utf8:&[u8])->String
{
match from_utf8(&unsafe_utf8.iter().filter(|c| **c != 0).map(|c| *c).collect::<Vec<u8>>() as _)
{
Err(err)=>{
println!("Error {}",err);
"".to_string()
},
Ok(res)=>res.trim().to_string()
}
}
fn main() {
unsafe {
let h_device_info=SetupDiGetClassDevsA(&GUID_DEVINTERFACE_DISK, null_mut(), null_mut(), DIGCF_PRESENT|DIGCF_DEVICEINTERFACE);
if h_device_info!=INVALID_HANDLE_VALUE {
let mut i=0;
let mut device_info_data:SP_DEVINFO_DATA=zeroed();
device_info_data.cbSize=size_of::<SP_DEVINFO_DATA>() as _;
while SetupDiEnumDeviceInfo(h_device_info, i, &mut device_info_data)!=0
{
let enumerator_name=get_enumerator_name(h_device_info, &mut device_info_data);
let friendly_name=get_friendly_name(h_device_info, &mut device_info_data);
if friendly_name.is_empty() {
continue;
}
println!("Name: {}",friendly_name);
println!("Is USB drive: {}",is_usb_drive(&enumerator_name));
println!("Is removable: {}",is_removable(h_device_info, &mut device_info_data));
get_detail_data(h_device_info, &mut device_info_data);
i+=1;
}
}
SetupDiDestroyDeviceInfoList(h_device_info);
}
}
Terminal output:
Name: SKHynix_HFM512GDHTNI-87A0B
Is USB drive: false
Is removable: false
DevicePath: \\?\scsi#disk&ven_nvme&prod_skhynix_hfm512gd#5&8980ef4&0&000000#{53f56307-b6bf-11d0-94f2-00a0c91efb8b}
Bytes returned: 0
get_device_number fail. Error 1
Everything went fine until DeviceIoControl always returns 0 and GetLastError() returns 1 (inside get_device_number() function). Can someone guide me figuring out what went wrong?
Thanks in advance
Related
I'm using binary trees to create a simple computation graph. I understand that linked lists are a pain in Rust, but it's a very convenient data structure for what I'm doing. I tried using Box and Rc<RefCell> for the children nodes, but it didn't work out how I wanted, so I used unsafe:
use std::ops::{Add, Mul};
#[derive(Debug, Copy, Clone)]
struct MyStruct {
value: i32,
lchild: Option<*mut MyStruct>,
rchild: Option<*mut MyStruct>,
}
impl MyStruct {
unsafe fn print_tree(&mut self, set_to_zero: bool) {
if set_to_zero {
self.value = 0;
}
println!("{:?}", self);
let mut nodes = vec![self.lchild, self.rchild];
while nodes.len() > 0 {
let child;
match nodes.pop() {
Some(popped_child) => child = popped_child.unwrap(),
None => continue,
}
if set_to_zero {
(*child).value = 0;
}
println!("{:?}", *child);
if !(*child).lchild.is_none() {
nodes.push((*child).lchild);
}
if !(*child).rchild.is_none() {
nodes.push((*child).rchild);
}
}
println!("");
}
}
impl Add for MyStruct {
type Output = Self;
fn add(self, other: Self) -> MyStruct {
MyStruct{
value: self.value + other.value,
lchild: Some(&self as *const _ as *mut _),
rchild: Some(&other as *const _ as *mut _),
}
}
}
impl Mul for MyStruct {
type Output = Self;
fn mul(self, other: Self) -> Self {
MyStruct{
value: self.value * other.value,
lchild: Some(&self as *const _ as *mut _),
rchild: Some(&other as *const _ as *mut _),
}
}
}
fn main() {
let mut tree: MyStruct;
{
let a = MyStruct{ value: 10, lchild: None, rchild: None };
let b = MyStruct{ value: 20, lchild: None, rchild: None };
let c = a + b;
println!("c.value: {}", c.value); // 30
let mut d = a + b;
println!("d.value: {}", d.value); // 30
d.value = 40;
println!("d.value: {}", d.value); // 40
let mut e = c * d;
println!("e.value: {}", e.value); // 1200
unsafe {
e.print_tree(false); // correct values
e.print_tree(true); // all zeros
e.print_tree(false); // all zeros, everything is set correctly
}
tree = e;
}
unsafe { tree.print_tree(false); } // same here, only zeros
}
Link to the playground
I honestly don't mind that much using unsafe, but is there a safe way doing it? How bad is the use of unsafe here?
You can just box both of the children, since you have a unidirectional tree:
use std::ops::{Add, Mul};
use std::fmt;
#[derive(Clone)]
struct MyStruct {
value: i32,
lchild: Option<Box<MyStruct>>,
rchild: Option<Box<MyStruct>>,
}
impl fmt::Debug for MyStruct {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
f.debug_struct("MyStruct")
.field("value", &self.value)
.field("lchild", &self.lchild.as_deref())
.field("rchild", &self.rchild.as_deref())
.finish()
}
}
impl MyStruct {
fn print_tree(&mut self, set_to_zero: bool) {
if set_to_zero {
self.value = 0;
}
println!("MyStruct {{ value: {:?}, lchild: {:?}, rchild: {:?} }}", self.value, &self.lchild as *const _, &self.rchild as *const _);
if let Some(child) = &mut self.lchild {
child.print_tree(set_to_zero);
}
if let Some(child) = &mut self.rchild {
child.print_tree(set_to_zero);
}
}
}
impl Add for MyStruct {
type Output = Self;
fn add(self, other: Self) -> MyStruct {
MyStruct {
value: self.value + other.value,
lchild: Some(Box::new(self)),
rchild: Some(Box::new(other)),
}
}
}
impl Mul for MyStruct {
type Output = Self;
fn mul(self, other: Self) -> Self {
MyStruct {
value: self.value * other.value,
lchild: Some(Box::new(self)),
rchild: Some(Box::new(other)),
}
}
}
fn main() {
let tree = {
let a = MyStruct {
value: 10,
lchild: None,
rchild: None,
};
let b = MyStruct {
value: 20,
lchild: None,
rchild: None,
};
let c = a.clone() + b.clone();
println!("c.value: {}", c.value); // 30
let mut d = a.clone() + b.clone();
println!("d.value: {}", d.value); // 30
d.value = 40;
println!("d.value: {}", d.value); // 40
let mut e = c * d;
println!("e.value: {}", e.value); // 1200
println!("");
e.print_tree(false); // correct values
println!("");
e.print_tree(true); // all zeros
println!("");
e.print_tree(false); // all zeros, everything is set correctly
println!("");
e
};
dbg!(tree);
}
I implemented Debug manually and reimplemented print_tree recursively. I don't know if there is a way to implement print_tree as mutable like that without recursion, but it's certainly possible if you take &self instead (removing the set_to_zero stuff).
playground
Edit: Turns out it is possible to mutably iterate over the tree values without recursion. The following code is derived from the playground in this comment by #Shepmaster.
impl MyStruct {
fn zero_tree(&mut self) {
let mut node_stack = vec![self];
let mut value_stack = vec![];
// collect mutable references to each value
while let Some(MyStruct { value, lchild, rchild }) = node_stack.pop() {
value_stack.push(value);
if let Some(child) = lchild {
node_stack.push(child);
}
if let Some(child) = rchild {
node_stack.push(child);
}
}
// iterate over mutable references to values
for value in value_stack {
*value = 0;
}
}
}
I have the following program that works fine on Linux:
Cargo.toml
[package]
name = "ansi-color-codec"
authors = ["Richard Neumann <mail#richard-neumann.de>"]
description = "Encode bytes as ANSI background colors"
license-file = "LICENSE"
homepage = "https://github.com/conqp/ansi-color-codec/"
repository = "https://github.com/conqp/ansi-color-codec/"
readme = "README.md"
documentation = "https://docs.rs/ansi-color-codec"
keywords = [ "ANSI", "color", "encoding"]
categories = ["command-line-utilities", "encoding"]
version = "0.3.8"
edition = "2021"
exclude = [
".gitignore",
"input.txt",
]
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
clap = { version = "4.0.23", features = ["derive"] }
ctrlc = "3.2.3"
[profile.release]
strip = true
lto = true
codegen-units = 1
panic = "abort"use std::iter::FlatMap;
src/lib.rs
const MASK_LOW: u8 = 0b00001111;
const MASK_HIGH: u8 = 0b11110000;
const MASK_BITS: u8 = 4;
const MASK_TRIPLET: u8 = MASK_LOW >> 1;
const COLOR_OFFSET_LOW: u8 = 40;
const COLOR_OFFSET_HIGH: u8 = 100;
const COLOR_CODE_LOW_MAX: u8 = MASK_TRIPLET;
const COLOR_CODE_MAX: u8 = MASK_LOW;
const COLOR_CODE_HIGH_BIT: u8 = 0b1000;
const MAX_DIGITS: u8 = 3;
const CODE_START: u8 = 0x1b;
const NUMBER_PREFIX: char = '[';
const NUMBER_SUFFIX: char = 'm';
const UNEXPECTED_TERMINATION_MSG: &str = "Byte stream terminated unexpectedly";
type ColorCodes<T> = FlatMap<T, [ColorCode; 2], fn(u8) -> [ColorCode; 2]>;
pub trait ColorCodec<T>
where
T: Iterator<Item = u8>,
{
fn ansi_color_encode(self) -> ColorCodes<T>;
fn ansi_color_decode(self) -> ColorCodesToBytes<ColorCodesFromBytes<T>>;
}
impl<T> ColorCodec<T> for T
where
T: Iterator<Item = u8>,
{
fn ansi_color_encode(self) -> ColorCodes<T> {
self.flat_map(|byte| byte.to_color_codes())
}
fn ansi_color_decode(self) -> ColorCodesToBytes<ColorCodesFromBytes<T>> {
ColorCodesToBytes::from(ColorCodesFromBytes::from(self))
}
}
#[derive(Debug, Eq, PartialEq)]
pub struct ColorCode {
number: u8,
}
impl ColorCode {
pub fn new(number: u8) -> Result<Self, String> {
if (0..=COLOR_OFFSET_LOW + COLOR_CODE_LOW_MAX).contains(&number)
|| (COLOR_OFFSET_HIGH..=COLOR_OFFSET_HIGH + COLOR_CODE_LOW_MAX).contains(&number)
{
Ok(Self { number })
} else {
Err(format!("Invalid color code: {}", number))
}
}
pub fn normalized(&self) -> u8 {
if self.number < COLOR_OFFSET_HIGH {
self.number - COLOR_OFFSET_LOW
} else {
self.number - COLOR_OFFSET_HIGH + COLOR_CODE_HIGH_BIT
}
}
}
impl TryFrom<u8> for ColorCode {
type Error = String;
fn try_from(value: u8) -> Result<Self, Self::Error> {
if value <= COLOR_CODE_LOW_MAX {
Self::new(value + COLOR_OFFSET_LOW)
} else if value <= COLOR_CODE_MAX {
Self::new((value & MASK_TRIPLET) + COLOR_OFFSET_HIGH)
} else {
Err(format!("Value out of bounds for color code: {}", value))
}
}
}
impl ToString for ColorCode {
fn to_string(&self) -> String {
format!("\x1b[{}m ", self.number)
}
}
trait ColorEncodable {
fn to_color_codes(&self) -> [ColorCode; 2];
fn from_color_codes(color_codes: [ColorCode; 2]) -> Self;
}
impl ColorEncodable for u8 {
fn to_color_codes(&self) -> [ColorCode; 2] {
[
ColorCode::try_from((self & MASK_HIGH) >> MASK_BITS).unwrap(),
ColorCode::try_from(self & MASK_LOW).unwrap(),
]
}
fn from_color_codes(color_codes: [ColorCode; 2]) -> Self {
(color_codes[0].normalized() << MASK_BITS) + color_codes[1].normalized()
}
}
#[derive(Debug, Eq, PartialEq)]
pub struct ColorCodesFromBytes<T>
where
T: Iterator<Item = u8>,
{
bytes: T,
}
impl<T> ColorCodesFromBytes<T>
where
T: Iterator<Item = u8>,
{
fn next_header(&mut self) -> Option<Result<(), String>> {
match self.bytes.next() {
Some(byte) => {
if byte == CODE_START {
match self.bytes.next() {
Some(byte) => {
if byte as char == NUMBER_PREFIX {
Some(Ok(()))
} else {
Some(Err(format!("Invalid number prefix: {}", byte)))
}
}
None => Some(Err(UNEXPECTED_TERMINATION_MSG.to_string())),
}
} else {
Some(Err(format!("Invalid start byte: {}", byte)))
}
}
None => None,
}
}
fn read_digits(&mut self) -> Result<String, String> {
let mut digits = String::new();
for count in 0..=MAX_DIGITS {
match self.bytes.next() {
Some(byte) => {
if byte.is_ascii_digit() {
if count < MAX_DIGITS {
digits.push(byte as char);
} else {
return Err(format!("Expected at most {} digits", MAX_DIGITS));
}
} else if byte as char == NUMBER_SUFFIX {
return if digits.is_empty() {
Err("Expected at least one digit".to_string())
} else {
Ok(digits)
};
} else {
return Err(format!("Encountered Unexpected byte \"{}\"", byte));
}
}
None => return Err(UNEXPECTED_TERMINATION_MSG.to_string()),
}
}
Ok(digits)
}
fn parse_color_code(&mut self) -> Result<u8, String> {
let digits = self.read_digits()?;
self.bytes.next(); // Discard bg-color encoded char
match digits.parse::<u8>() {
Ok(number) => Ok(number),
Err(_) => Err(format!("Could not parse u8 from {}", digits)),
}
}
}
impl<T> From<T> for ColorCodesFromBytes<T>
where
T: Iterator<Item = u8>,
{
fn from(bytes: T) -> Self {
Self { bytes }
}
}
impl<T> Iterator for ColorCodesFromBytes<T>
where
T: Iterator<Item = u8>,
{
type Item = Result<ColorCode, String>;
fn next(&mut self) -> Option<Self::Item> {
if let Err(msg) = self.next_header()? {
return Some(Err(msg));
}
match self.parse_color_code() {
Ok(sum) => {
if sum == 0 {
None
} else {
Some(ColorCode::new(sum))
}
}
Err(msg) => Some(Err(format!("{} while parsing color code", msg))),
}
}
}
#[derive(Debug, Eq, PartialEq)]
pub struct ColorCodesToBytes<T>
where
T: Iterator<Item = Result<ColorCode, String>>,
{
codes: T,
}
impl<T> From<T> for ColorCodesToBytes<T>
where
T: Iterator<Item = Result<ColorCode, String>>,
{
fn from(codes: T) -> Self {
Self { codes }
}
}
impl<T> Iterator for ColorCodesToBytes<T>
where
T: Iterator<Item = Result<ColorCode, String>>,
{
type Item = Result<u8, String>;
fn next(&mut self) -> Option<Self::Item> {
match self.codes.next() {
Some(high) => match high {
Ok(high) => match self.codes.next() {
Some(low) => match low {
Ok(low) => Some(Ok(u8::from_color_codes([high, low]))),
Err(msg) => Some(Err(msg)),
},
None => Some(Err("Missing second color code block".to_string())),
},
Err(msg) => Some(Err(msg)),
},
None => None,
}
}
}
src/main.rs
use ansi_color_codec::ColorCodec;
use clap::Parser;
use ctrlc::set_handler;
use std::io::{stdin, stdout, Read, Write};
use std::process::exit;
use std::sync::{
atomic::{AtomicBool, Ordering},
Arc,
};
const STDOUT_WRITE_ERR: &str = "Could not write bytes to STDOUT";
#[derive(Parser)]
#[clap(about, author, version)]
struct Args {
#[clap(short, long, name = "decode")]
pub decode: bool,
#[clap(short, long, name = "no-clear")]
pub no_clear: bool,
}
fn main() {
let args = Args::parse();
let running = Arc::new(AtomicBool::new(true));
let bytes = stream_stdin(running.clone());
set_handler(move || {
running.store(false, Ordering::SeqCst);
})
.expect("Error setting Ctrl-C handler");
if args.decode {
decode(bytes)
} else {
encode(bytes, !args.no_clear)
}
}
fn decode(bytes: impl Iterator<Item = u8>) {
for result in bytes.ansi_color_decode() {
match result {
Ok(byte) => {
stdout().write_all(&[byte]).expect(STDOUT_WRITE_ERR);
}
Err(msg) => {
eprintln!("{}", msg);
exit(1);
}
}
}
stdout().flush().expect("Could not flush STDOUT")
}
fn encode(bytes: impl Iterator<Item = u8>, clear: bool) {
for code in bytes.ansi_color_encode() {
stdout()
.write_all(code.to_string().as_bytes())
.expect(STDOUT_WRITE_ERR);
}
if clear {
println!("\x1b[0m ");
}
}
fn stream_stdin(running: Arc<AtomicBool>) -> impl Iterator<Item = u8> {
stdin()
.bytes()
.take_while(move |byte| byte.is_ok() && running.load(Ordering::SeqCst))
.map(|byte| byte.unwrap())
}
However, when I run
> echo "Windows doing Windows stuff" | ansi-color-codec | ansi-color-codec -d
on Windows, the program fails with
Invalid start byte: 239
When I inspect the first (three) bytes, I can see that ansi-color-codec -d receives the UTF-8 BOM from ansi-color-codec. But why? My program does not print it and only puts raw bytes onto STDOUT
Found the answer on Reddit:
EDIT : [SOLVED]
CAUSE : Windows powershell (at least) uses different code page from external programs ( rust program in this case ) which caused inconsistent inter process communication.
SOLUTION : Set the following environment variable for consistent communication with rust programs.
$OutputEncoding = [console]::InputEncoding = [console]::OutputEncoding = New-Object System.Text.UTF8Encoding
Hello my goal is to create a rust function which takes a process name as string and returns a PID.
I came up with this function:
pub unsafe fn get_proc_id(proc_name: String) -> u32 {
let mut proc_id: u32 = 0;
let mut h_snap = windows::Win32::System::Diagnostics::ToolHelp::CreateToolhelp32Snapshot(
TH32CS_SNAPPROCESS,
0,
);
let h_snap = match h_snap {
Ok(t) => t,
Err(e) => panic!("eror {}", e),
};
let mut proc_entry: PROCESSENTRY32 = PROCESSENTRY32 {
..PROCESSENTRY32::default()
};
proc_entry.dwSize = std::mem::size_of::<PROCESSENTRY32>() as u32;
let entry_ptr = &mut proc_entry as *mut PROCESSENTRY32;
if windows::Win32::System::Diagnostics::ToolHelp::Process32First(h_snap, entry_ptr).as_bool() {
loop {
let mut proc_exe_string: String = String::new();
if proc_exe_string.eq(&proc_name) {
proc_id = proc_entry.th32ProcessID;
break;
}
for e in proc_entry.szExeFile {
if e.0 != 0 {
proc_exe_string.push(e.0 as char)
}
}
println!("{}", proc_exe_string);
if !Process32Next(h_snap, entry_ptr).as_bool() {
break;
}
}
}
CloseHandle(h_snap);
return proc_id; }
The function prints some strange process names. for example I'm looking for "ac_client.exe" however the function shows this process as ac_client.exeexe.exeee and this is similar for most process names. some other examples:
chrome.exexeexe.exeee
Discord.exer.exee.exee
I've implemented the IWbemQuerySink in rust, in order to use ExecQueryAsync. Code is available in this pull request but I'll post the relevant bits here.
TLDR;
WMI is calling Release() 5 times on my QuerySink implementation, 1 which is understandable, and 4 which are not... There is only 1 call to AddRef by WMI, so the question is: where do this 4 other calls to Release comes from and is there a way to prevent them from happening?
With a lot of details
Here is a simplified implementation, it is working great:
use winapi::{
um::wbemcli::{
{IWbemClassObject,IWbemObjectSink, IWbemObjectSinkVtbl},
WBEM_S_NO_ERROR,
},
shared::{
ntdef::HRESULT,
wtypes::BSTR,
},
ctypes::{
c_long,
},
};
use com_impl::{ComImpl, VTable, Refcount};
use wio::com::ComPtr;
#[repr(C)]
#[derive(ComImpl)]
#[interfaces(IWbemObjectSink)]
pub struct QuerySink {
vtbl: VTable<IWbemObjectSinkVtbl>,
refcount: Refcount,
}
impl QuerySink {
pub fn new() -> ComPtr<IWbemObjectSink> {
let ptr = QuerySink::create_raw();
let ptr = ptr as *mut IWbemObjectSink;
unsafe { ComPtr::from_raw(ptr) }
}
}
// AddRef and Release methods are provided by com_impl
#[com_impl::com_impl]
unsafe impl IWbemObjectSink for QuerySink {
pub unsafe fn indicate(
&self,
_lObjectCount: c_long,
_apObjArray: *mut *mut IWbemClassObject
) -> HRESULT {
WBEM_S_NO_ERROR as i32
}
pub unsafe fn set_status(
&self,
_lFlags: c_long,
_hResult: HRESULT,
_strParam: BSTR,
_pObjParam: *mut IWbemClassObject
) -> HRESULT {
WBEM_S_NO_ERROR as i32
}
}
Let's try to use it in a simple call to ExecQueryAsync:
pub fn exec_async_query_native_wrapper(
&self,
query: impl AsRef<str>,
) -> Result<(), WMIError> {
let query_language = BStr::from_str("WQL")?;
let query = BStr::from_str(query.as_ref())?;
let p_sink: ComPtr<IWbemObjectSink> = QuerySink::new();
unsafe {
check_hres((*self.svc()).ExecQueryAsync(
query_language.as_bstr(),
query.as_bstr(),
WBEM_FLAG_BIDIRECTIONAL as i32,
ptr::null_mut(),
p_sink.as_raw(),
))?;
}
Ok(())
}
When executing it in a test, this code panics because the release function is called too many times, and the Refcount becomes negative:
attempt to subtract with overflow', C:\Users\apennamen\.cargo\registry\src\github.com-1ecc6299db9ec823\com-impl-0.1.1\src\lib.rs:139:9
stack backtrace:
0: std::panicking::begin_panic_handler
at /rustc/b32e6e6ac8921035177256ab6806e6ab0d4b9b94\/library\std\src\panicking.rs:493
1: core::panicking::panic_fmt
at /rustc/b32e6e6ac8921035177256ab6806e6ab0d4b9b94\/library\core\src\panicking.rs:92
2: core::panicking::panic
at /rustc/b32e6e6ac8921035177256ab6806e6ab0d4b9b94\/library\core\src\panicking.rs:50
3: com_impl::Refcount::release
at C:\Users\apennamen\.cargo\registry\src\github.com-1ecc6299db9ec823\com-impl-0.1.1\src\lib.rs:139
4: wmi::query_sink::QuerySink::__com_impl__IUnknown__Release
at .\src\query_sink.rs:31
5: CoMarshalInterface
6: CoMarshalInterface
.....
After some tests, I found that the Release method is called 4 times more than it should be by WMI. So when adding 4 calls to AddRef, the code works just fine. This is of course not satisfying as maybe tomorrow it will be 5 calls, or 10 or 2...
I don't know how to handle this problem, hopefully someone else ran into it
This is a hacky working code:
pub fn exec_async_query_native_wrapper(
&self,
query: impl AsRef<str>,
) -> Result<(), WMIError> {
let query_language = BStr::from_str("WQL")?;
let query = BStr::from_str(query.as_ref())?;
let p_sink: ComPtr<IWbemObjectSink> = QuerySink::new();
unsafe {
// FIXME hack the RefCount
for _ in 0..4 { p_sink.AddRef(); }
check_hres((*self.svc()).ExecQueryAsync(
query_language.as_bstr(),
query.as_bstr(),
WBEM_FLAG_BIDIRECTIONAL as i32,
ptr::null_mut(),
p_sink.as_raw(),
))?;
}
Ok(())
}
}
Thank you for reading this far :)
EDIT :
Here is the expanded macro code:
#[cfg(feature = "async-query")]
pub(crate) mod query_sink {
use winapi::{
um::wbemcli::{
{IWbemClassObject, IWbemObjectSink, IWbemObjectSinkVtbl}, WBEM_S_NO_ERROR,
WBEM_STATUS_COMPLETE,
},
shared::{
ntdef::HRESULT,
wtypes::BSTR,
winerror::{E_POINTER, E_FAIL},
},
ctypes::{c_long},
};
use com_impl::{ComImpl, VTable, Refcount};
use log::{trace, warn};
use std::ptr::NonNull;
use wio::com::ComPtr;
use crate::result_enumerator::IWbemClassWrapper;
use crate::WMIError;
/// Implementation for [IWbemObjectSink](https://learn.microsoft.com/en-us/windows/win32/api/wbemcli/nn-wbemcli-iwbemobjectsink).
/// This [Sink](https://en.wikipedia.org/wiki/Sink_(computing))
/// receives asynchronously the result of the query,
/// through Indicate calls. When finished,the SetStatus method
/// is called.
/// # <https://learn.microsoft.com/fr-fr/windows/win32/wmisdk/example--getting-wmi-data-from-the-local-computer-asynchronously>
#[repr(C)]
#[interfaces(IWbemObjectSink)]
pub struct QuerySink {
vtbl: VTable<IWbemObjectSinkVtbl>,
refcount: Refcount,
}
impl QuerySink {
fn create_raw() -> *mut Self {
Box::into_raw(Box::new(QuerySink {
vtbl: <Self as com_impl::BuildVTable<_>>::static_vtable(),
refcount: Default::default(),
}))
}
}
unsafe impl com_impl::BuildVTable<winapi::um::unknwnbase::IUnknownVtbl> for QuerySink {
const VTBL: winapi::um::unknwnbase::IUnknownVtbl = winapi::um::unknwnbase::IUnknownVtbl {
AddRef: Self::__com_impl__IUnknown__AddRef,
Release: Self::__com_impl__IUnknown__Release,
QueryInterface: Self::__com_impl__IUnknown__QueryInterface,
};
fn static_vtable() -> com_impl::VTable<winapi::um::unknwnbase::IUnknownVtbl> {
com_impl::VTable::new(&Self::VTBL)
}
}
#[allow(non_snake_case)]
impl QuerySink {
#[inline(never)]
unsafe extern "system" fn __com_impl__IUnknown__AddRef(
this: *mut winapi::um::unknwnbase::IUnknown,
) -> u32 {
let this = &*(this as *const Self);
this.refcount.add_ref()
}
#[inline(never)]
unsafe extern "system" fn __com_impl__IUnknown__Release(
this: *mut winapi::um::unknwnbase::IUnknown,
) -> u32 {
let ptr = this as *mut Self;
let count = (*ptr).refcount.release();
if count == 0 {
Box::from_raw(ptr);
}
count
}
#[inline(never)]
unsafe extern "system" fn __com_impl__IUnknown__QueryInterface(
this: *mut winapi::um::unknwnbase::IUnknown,
riid: *const winapi::shared::guiddef::IID,
ppv: *mut *mut winapi::ctypes::c_void,
) -> winapi::shared::winerror::HRESULT {
if ppv.is_null() {
return winapi::shared::winerror::E_POINTER;
}
if winapi::shared::guiddef::IsEqualIID(
&*riid,
&<winapi::um::unknwnbase::IUnknown as winapi::Interface>::uuidof(),
) || winapi::shared::guiddef::IsEqualIID(
&*riid,
&<IWbemObjectSink as winapi::Interface>::uuidof(),
) {
*ppv = this as *mut winapi::ctypes::c_void;
winapi::shared::winerror::S_OK
} else {
*ppv = std::ptr::null_mut();
winapi::shared::winerror::E_NOINTERFACE
}
}
}
impl QuerySink {
pub fn new() -> ComPtr<IWbemObjectSink> {
let ptr = QuerySink::create_raw();
let ptr = ptr as *mut IWbemObjectSink;
unsafe { ComPtr::from_raw(ptr) }
}
}
unsafe impl com_impl::BuildVTable<IWbemObjectSinkVtbl> for QuerySink {
const VTBL: IWbemObjectSinkVtbl = IWbemObjectSinkVtbl {
parent: <Self as com_impl::BuildVTable<_>>::VTBL,
Indicate: Self::__com_impl_stub__IWbemObjectSink__Indicate,
SetStatus: Self::__com_impl_stub__IWbemObjectSink__SetStatus,
};
fn static_vtable() -> com_impl::VTable<IWbemObjectSinkVtbl> {
com_impl::VTable::new(&Self::VTBL)
}
}
#[allow(non_snake_case)]
impl QuerySink {
#[inline(never)]
unsafe extern "system" fn __com_impl_stub__IWbemObjectSink__Indicate(
this: *mut IWbemObjectSink,
__com_arg_0: c_long,
__com_arg_1: *mut *mut IWbemClassObject,
) -> HRESULT {
let this = &*(this as *const Self);
Self::__com_impl_body__IWbemObjectSink__Indicate(this, __com_arg_0, __com_arg_1)
}
#[inline(never)]
unsafe extern "system" fn __com_impl_stub__IWbemObjectSink__SetStatus(
this: *mut IWbemObjectSink,
__com_arg_0: c_long,
__com_arg_1: HRESULT,
__com_arg_2: BSTR,
__com_arg_3: *mut IWbemClassObject,
) -> HRESULT {
let this = &*(this as *const Self);
Self::__com_impl_body__IWbemObjectSink__SetStatus(
this,
__com_arg_0,
__com_arg_1,
__com_arg_2,
__com_arg_3,
)
}
#[inline(always)]
unsafe extern "system" fn __com_impl_body__IWbemObjectSink__Indicate(
&self,
lObjectCount: c_long,
apObjArray: *mut *mut IWbemClassObject,
) -> HRESULT {
WBEM_S_NO_ERROR as i32
}
#[inline(always)]
unsafe extern "system" fn __com_impl_body__IWbemObjectSink__SetStatus(
&self,
lFlags: c_long,
_hResult: HRESULT,
_strParam: BSTR,
_pObjParam: *mut IWbemClassObject,
) -> HRESULT {
WBEM_S_NO_ERROR as i32
}
}
}
I have been trying to use the named pipes on Windows 10, but for some reason it doesn't seem to be either writing to it or reading from it correctly.
Running the below code will block on the spawned thread (output.read()) waiting for bytes to be written with input.write.
Let me know if I should add the imports.
fn main() {
let pipes = create_pipe(Path::new(r#"\\.\pipe\input"#));
let mut server: File = unsafe { File::from_raw_handle(pipes.0.as_raw_handle()) };
let client: File = unsafe { File::from_raw_handle(pipes.1.as_raw_handle()) };
let mut input = BufWriter::new(server);
let mut output = BufReader::new(client);
thread::spawn(move || {
let mut message = vec![];
loop {
output.read(&mut message).unwrap();
println!("m: {:#?}", message);
}
});
loop {
input.write("test".as_bytes()).unwrap();
thread::sleep(Duration::from_secs(1));
}
}
fn create_pipe(path: &Path) -> (Handle, Handle) {
let mut os_str: OsString = path.as_os_str().into();
os_str.push("\x00");
let u16_slice = os_str.encode_wide().collect::<Vec<u16>>();
let access_flags =
PIPE_ACCESS_DUPLEX | FILE_FLAG_OVERLAPPED | FILE_FLAG_FIRST_PIPE_INSTANCE | WRITE_DAC;
let server_handle: RawHandle = unsafe {
namedpipeapi::CreateNamedPipeW(
u16_slice.as_ptr(),
access_flags,
PIPE_TYPE_BYTE | PIPE_READMODE_BYTE | PIPE_WAIT,
1,
65536,
65536,
0,
ptr::null_mut(),
)
};
let mut attributes = SECURITY_ATTRIBUTES {
nLength: mem::size_of::<SECURITY_ATTRIBUTES>() as DWORD,
lpSecurityDescriptor: ptr::null_mut(),
bInheritHandle: true as BOOL,
};
let child_handle: RawHandle = unsafe {
fileapi::CreateFileW(
u16_slice.as_ptr(),
GENERIC_READ | GENERIC_WRITE | FILE_READ_ATTRIBUTES,
0,
&mut attributes as LPSECURITY_ATTRIBUTES,
OPEN_EXISTING,
0,
ptr::null_mut(),
)
};
if server_handle != INVALID_HANDLE_VALUE && child_handle != INVALID_HANDLE_VALUE {
let ret = unsafe { ConnectNamedPipe(server_handle, ptr::null_mut()) != 0 };
if !ret {
let err = unsafe { GetLastError() };
if err != 535 {
panic!("Pipe error");
}
}
(Handle(server_handle), Handle(child_handle))
} else {
panic!(io::Error::last_os_error())
}
}
#[derive(Debug)]
pub struct Handle(RawHandle);
unsafe impl Send for Handle {}
impl AsRawHandle for Handle {
fn as_raw_handle(&self) -> RawHandle {
self.0
}
}
impl FromRawHandle for Handle {
unsafe fn from_raw_handle(handle: RawHandle) -> Handle {
Handle(handle)
}
}