Rust中WinAPI GetTokenInformation包装器的优化咨询

问题描述 #

我在为WinAPI的GetTokenInformation函数创建Rust辅助包装器时遇到了困难，初步实现了如下代码，但对返回类型存疑，希望得到优化建议。

当前实现代码 #

/// The wrapper over GetTokenInformation
pub fn get_token_information(
    token_handle: winapi::um::winnt::HANDLE,
    token_information_class: u32
) -> Result<Box<[u8]>, _> {
    let mut required_len = 0;
    unsafe {
        winapi::um::securitybaseapi::GetTokenInformation(
            token_handle,
            token_information_class,
            std::ptr::null_mut(),
            0,
            &mut required_len,
        )

        let mut mem = vec![0u8; required_len as usize].into_boxed_slice();
        winapi::um::securitybaseapi::GetTokenInformation(
            token_handle,
            token_information_class,
            mem.as_mut_ptr() as _,
            required_len,
            &mut required_len,
        );

        Ok(mem)
    }
}

/// Sample function using get_token_information wrapper.
fn get_user_sid() {
    ...  
    let token_user = get_token_information(process_token_handle.0, winapi::um::winnt::TokenUser)?;
    let sid = unsafe {
        let ptr = token_user.as_ptr() as *const TOKEN_USER;
        convert_sid_to_string_sid((*ptr).User.Sid)?
    };
}

已获反馈与优化需求 #

更新：根据@cafce25和@IInspectable的回复，上述代码用于内部使用基本可行，但仍希望获取其他优化意见。

优化建议 #

• 完善错误处理逻辑
  当前代码未检查GetTokenInformation的返回值，存在潜在风险：

  1. 第一次调用需确认失败原因是ERROR_INSUFFICIENT_BUFFER（通过GetLastError()获取），避免忽略其他未知错误。
  2. 第二次调用必须检查返回值，调用失败时应获取对应错误码并返回Err，而非直接返回内存块。
     修正片段示例：

  use winapi::um::errhandlingapi::GetLastError;
  use winapi::shared::minwindef::ERROR_INSUFFICIENT_BUFFER;
  
  // 第一次调用获取所需长度
  let ret = winapi::um::securitybaseapi::GetTokenInformation(
      token_handle,
      token_information_class,
      std::ptr::null_mut(),
      0,
      &mut required_len,
  );
  if ret == 0 {
      let err = unsafe { GetLastError() };
      if err != ERROR_INSUFFICIENT_BUFFER {
          return Err(err);
      }
  }
  
  // 第二次调用填充数据
  let ret = winapi::um::securitybaseapi::GetTokenInformation(
      token_handle,
      token_information_class,
      mem.as_mut_ptr() as _,
      required_len,
      &mut required_len,
  );
  if ret == 0 {
      return Err(unsafe { GetLastError() });
  }
  

• 引入强类型枚举提升参数安全性
  将token_information_class从裸u32替换为对应WinAPI枚举的强类型封装，避免传入无效值：

  #[repr(u32)]
  pub enum TokenInformationClass {
      TokenUser = winapi::um::winnt::TokenUser,
      TokenGroups = winapi::um::winnt::TokenGroups,
      // 按需添加其他需要的枚举值
  }
  
  // 函数参数修改为
  token_information_class: TokenInformationClass
  

• 提供泛型版本简化类型转换
  当前返回Box<[u8]>需要用户手动转换为目标结构体，可实现泛型版本直接返回封装后的结构体，减少重复unsafe代码：

  pub fn get_token_information<T>(
      token_handle: winapi::um::winnt::HANDLE,
      token_info_class: TokenInformationClass
  ) -> Result<Box<T>, u32>
  where
      T: Sized,
  {
      let mut required_len = 0;
      unsafe {
          // 第一次调用获取长度
          let ret = winapi::um::securitybaseapi::GetTokenInformation(
              token_handle,
              token_info_class as u32,
              std::ptr::null_mut(),
              0,
              &mut required_len,
          );
          if ret == 0 {
              let err = GetLastError();
              if err != ERROR_INSUFFICIENT_BUFFER {
                  return Err(err);
              }
          }
  
          // 检查内存长度是否足够容纳目标结构体
          if required_len < std::mem::size_of::<T>() as u32 {
              return Err(winapi::shared::minwindef::ERROR_BAD_LENGTH);
          }
  
          // 分配内存并填充数据
          let mut mem = vec![0u8; required_len as usize].into_boxed_slice();
          let ret = winapi::um::securitybaseapi::GetTokenInformation(
              token_handle,
              token_info_class as u32,
              mem.as_mut_ptr() as _,
              required_len,
              &mut required_len,
          );
          if ret == 0 {
              return Err(GetLastError());
          }
  
          // 转换为目标结构体类型
          Ok(Box::from_raw(mem.as_mut_ptr() as *mut T))
      }
  }
  

  使用示例：

  let token_user: Box<winapi::um::winnt::TOKEN_USER> = get_token_information(
      process_token_handle.0,
      TokenInformationClass::TokenUser
  )?;
  let sid = convert_sid_to_string_sid(token_user.User.Sid)?;
  

• 优化内存分配灵活性
  若不需要固定长度的内存块，可直接返回Vec<u8>而非Box<[u8]>，Vec支持更多操作（如切片、转换），且可轻松转换为Box<[u8]>，提升代码灵活性。

内容的提问来源于stack exchange，提问作者Rob D