Skip to content

Commit

Permalink
feat(agent): initial DVC server implementation (#1053)
Browse files Browse the repository at this point in the history
  • Loading branch information
pacmancoder authored Nov 8, 2024
1 parent 5fa998e commit caa5ffa
Show file tree
Hide file tree
Showing 32 changed files with 2,410 additions and 288 deletions.
268 changes: 236 additions & 32 deletions Cargo.lock

Large diffs are not rendered by default.

Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,7 @@ use std::rc::Rc;
use futures::Future;
use hyper;

use super::request as __internal_request;
use super::{configuration, Error};
use super::{configuration, request as __internal_request, Error};
use crate::models;

pub struct DefaultApiClient<C: hyper::client::connect::Connect>
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,4 @@
use http;
use hyper;
use serde_json;
use {http, hyper, serde_json};

#[derive(Debug)]
pub enum Error {
Expand Down
Original file line number Diff line number Diff line change
@@ -1,13 +1,10 @@
use std::collections::HashMap;
use std::pin::Pin;

use futures;
use futures::future::*;
use futures::Future;
use hyper;
use hyper::header::{HeaderValue, AUTHORIZATION, CONTENT_LENGTH, CONTENT_TYPE, USER_AGENT};
use serde;
use serde_json;
use {futures, hyper, serde, serde_json};

use super::{configuration, Error};

Expand Down
1 change: 1 addition & 0 deletions crates/win-api-wrappers/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ base16ct = { version = "0.2", features = ["alloc"] }
anyhow = "1.0"
thiserror = "1.0"
tracing = "0.1"
uuid = { version = "1.10", features = ["v4"] }

[dependencies.windows]
version = "0.58.0"
Expand Down
38 changes: 38 additions & 0 deletions crates/win-api-wrappers/src/event.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
use std::sync::Arc;

use windows::Win32::Foundation::HANDLE;
use windows::Win32::System::Threading::{CreateEventW, SetEvent};

use crate::handle::Handle;

/// RAII wrapper for WinAPI event handle.
#[derive(Debug, Clone)]
pub struct Event {
handle: Arc<Handle>,
}

impl Event {
pub fn new_unnamed() -> anyhow::Result<Self> {
// SAFETY: No preconditions.
let raw_handle = unsafe { CreateEventW(None, false, false, None) }?;

// SAFETY: `CreateEventW` always returns a valid handle on success.
let handle = unsafe { Handle::new_owned(raw_handle) }?;

Ok(Self {
handle: Arc::new(handle),
})
}

pub fn raw(&self) -> HANDLE {
self.handle.raw()
}

pub fn set(&self) -> anyhow::Result<()> {
// SAFETY: No preconditions.
unsafe {
SetEvent(self.handle.raw())?;
}
Ok(())
}
}
2 changes: 2 additions & 0 deletions crates/win-api-wrappers/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ mod lib_win {
mod error;
pub use error::Error;

pub mod event;
pub mod handle;
pub mod identity;
pub mod process;
Expand All @@ -15,6 +16,7 @@ mod lib_win {
pub mod thread;
pub mod token;
pub mod utils;
pub mod wts;

// Allowed since the goal is to replicate the Windows API crate so that it's familiar, which itself uses the raw names from the API.
#[allow(
Expand Down
42 changes: 23 additions & 19 deletions crates/win-api-wrappers/src/process.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
use std::collections::HashMap;
use std::ffi::{c_void, OsString};
use std::fmt::Debug;
use std::mem::{self};
use std::os::windows::ffi::OsStringExt;
use std::path::{Path, PathBuf};
use std::{ptr, slice};
Expand Down Expand Up @@ -30,8 +29,8 @@ use windows::Win32::System::LibraryLoader::{
GetModuleFileNameW, GetModuleHandleExW, GetProcAddress, GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS,
};
use windows::Win32::System::Threading::{
CreateProcessAsUserW, CreateRemoteThread, GetCurrentProcess, GetExitCodeProcess, OpenProcess, OpenProcessToken,
QueryFullProcessImageNameW, TerminateProcess, WaitForSingleObject, CREATE_UNICODE_ENVIRONMENT,
CreateProcessAsUserW, CreateRemoteThread, GetCurrentProcess, GetExitCodeProcess, GetProcessId, OpenProcess,
OpenProcessToken, QueryFullProcessImageNameW, TerminateProcess, WaitForSingleObject, CREATE_UNICODE_ENVIRONMENT,
EXTENDED_STARTUPINFO_PRESENT, INFINITE, LPPROC_THREAD_ATTRIBUTE_LIST, LPTHREAD_START_ROUTINE, PEB,
PROCESS_ACCESS_RIGHTS, PROCESS_BASIC_INFORMATION, PROCESS_CREATION_FLAGS, PROCESS_INFORMATION, PROCESS_NAME_WIN32,
PROCESS_TERMINATE, STARTUPINFOEXW, STARTUPINFOW, STARTUPINFOW_FLAGS,
Expand Down Expand Up @@ -122,7 +121,7 @@ impl Process {
let thread = self.create_thread(
// SAFETY: `LoadLibraryW` fits the type. It takes one argument that is the name of the library.
Some(unsafe {
mem::transmute::<*const c_void, unsafe extern "system" fn(*mut c_void) -> u32>(load_library)
core::mem::transmute::<*const c_void, unsafe extern "system" fn(*mut c_void) -> u32>(load_library)
}),
Some(allocation.address),
)?;
Expand Down Expand Up @@ -152,7 +151,7 @@ impl Process {
)?
};

// SAFETY: The handle is owned by us, we opened the ressource above.
// SAFETY: The handle is owned by us, we opened the resource above.
let handle = unsafe { Handle::new_owned(handle) }?;

Ok(Thread::from(handle))
Expand Down Expand Up @@ -242,6 +241,11 @@ impl Process {
})
}

pub fn pid(&self) -> u32 {
// SAFETY: Safe to call with a valid process handle.
unsafe { GetProcessId(self.handle.raw()) }
}

/// Reads process memory at a specified address into a buffer.
/// The buffer is not read.
/// Returns the number of bytes read.
Expand All @@ -266,13 +270,13 @@ impl Process {
Ok(bytes_read)
}

/// Reads a stucture from process memory at a specified address.
/// Reads a structure from process memory at a specified address.
///
/// # Safety
///
/// - `address` must point to a valid and correctly sized instance of the structure.
pub unsafe fn read_struct<T: Sized>(&self, address: *const c_void) -> Result<T> {
let mut buf = vec![0; mem::size_of::<T>()];
let mut buf = vec![0; size_of::<T>()];

// SAFETY: Based on the security requirements of the function, the `address` should
// point to a valid and correctly sized instance of `T`.
Expand All @@ -286,7 +290,7 @@ impl Process {
}
}

/// Reads a continous array of a structure from process memory at a specified address.
/// Reads a continuous array of a structure from process memory at a specified address.
///
/// # Safety
///
Expand All @@ -301,10 +305,10 @@ impl Process {
// However, we assume that the data will be alined as `Vec` wants.
let data = unsafe { data.align_to_mut::<u8>().1 };

// SAFETY: `read_memory` does not read `data`, so we can safely pass an unitialized buffer.
// SAFETY: `read_memory` does not read `data`, so we can safely pass an uninitialized buffer.
let read_bytes = unsafe { self.read_memory(address.cast(), data) }?;

if count * mem::size_of::<T>() == read_bytes {
if count * size_of::<T>() == read_bytes {
// SAFETY: Buffer can hold `count` items and was filled up to that point.
unsafe { buf.set_len(count) };

Expand Down Expand Up @@ -380,31 +384,31 @@ impl Peb<'_> {
let image_path_name = unsafe {
self.process.read_array(
raw_params.ImagePathName.Buffer.as_ptr(),
raw_params.ImagePathName.Length as usize / mem::size_of::<u16>(),
raw_params.ImagePathName.Length as usize / size_of::<u16>(),
)?
};

// SAFETY: We assume `raw_params.CommandLine` is truthful and valid.
let command_line = unsafe {
self.process.read_array(
raw_params.CommandLine.Buffer.as_ptr(),
raw_params.CommandLine.Length as usize / mem::size_of::<u16>(),
raw_params.CommandLine.Length as usize / size_of::<u16>(),
)?
};

// SAFETY: We assume `raw_params.DesktopInfo` is truthful and valid.
let desktop = unsafe {
self.process.read_array(
raw_params.DesktopInfo.Buffer.as_ptr(),
raw_params.DesktopInfo.Length as usize / mem::size_of::<u16>(),
raw_params.DesktopInfo.Length as usize / size_of::<u16>(),
)?
};

// SAFETY: We assume `raw_params.CurrentDirectory` is truthful and valid.
let working_directory = unsafe {
self.process.read_array(
raw_params.CurrentDirectory.DosPath.Buffer.as_ptr(),
raw_params.CurrentDirectory.DosPath.Length as usize / mem::size_of::<u16>(),
raw_params.CurrentDirectory.DosPath.Length as usize / size_of::<u16>(),
)?
};

Expand Down Expand Up @@ -579,8 +583,8 @@ impl ProcessEntry32Iterator {
};

// SAFETY: It is safe to zero out the structure as it is a simple POD type.
let mut process_entry: PROCESSENTRY32W = unsafe { mem::zeroed() };
process_entry.dwSize = mem::size_of::<PROCESSENTRY32W>()
let mut process_entry: PROCESSENTRY32W = unsafe { core::mem::zeroed() };
process_entry.dwSize = size_of::<PROCESSENTRY32W>()
.try_into()
.expect("BUG: PROCESSENTRY32W size always fits in u32");

Expand Down Expand Up @@ -700,7 +704,7 @@ pub fn create_process_as_user(
// SAFETY: As per `CreateEnvironmentBlock` documentation: We must specify
// `CREATE_UNICODE_ENVIRONMENT` and call `DestroyEnvironmentBlock` after
// `CreateProcessAsUser` call.
// - `CREATE_UNICODE_ENVIRONMENT` is always set unconditionaly.
// - `CREATE_UNICODE_ENVIRONMENT` is always set unconditionally.
// - `DestroyEnvironmentBlock` is called in the `ProcessEnvironment` destructor.
//
// Therefore, all preconditions are met to safely call `CreateEnvironmentBlock`.
Expand Down Expand Up @@ -742,10 +746,10 @@ pub fn create_process_as_user(
)
}?;

// SAFETY: The handle is owned by us, we opened the ressource above.
// SAFETY: The handle is owned by us, we opened the resource above.
let process = unsafe { Handle::new_owned(raw_process_information.hProcess).map(Process::from)? };

// SAFETY: The handle is owned by us, we opened the ressource above.
// SAFETY: The handle is owned by us, we opened the resource above.
let thread = unsafe { Handle::new_owned(raw_process_information.hThread).map(Thread::from)? };

Ok(ProcessInformation {
Expand Down
10 changes: 5 additions & 5 deletions crates/win-api-wrappers/src/rpc.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use std::ffi::c_void;
use std::ptr::{self, NonNull};
use std::{mem, slice};
use std::slice;

use windows::core::GUID;
use windows::Win32::Foundation::{ERROR_MORE_DATA, E_INVALIDARG};
Expand Down Expand Up @@ -51,9 +51,9 @@ impl RpcBindingHandle {
server_principal_name.resize((attribs.ServerPrincipalNameBufferLength / 2) as usize, 0);

attribs.ClientPrincipalName = client_principal_name.as_mut_ptr();
attribs.ClientPrincipalNameBufferLength = (client_principal_name.len() * mem::size_of::<u16>()).try_into()?;
attribs.ClientPrincipalNameBufferLength = (client_principal_name.len() * size_of::<u16>()).try_into()?;
attribs.ServerPrincipalName = server_principal_name.as_mut_ptr();
attribs.ServerPrincipalNameBufferLength = (server_principal_name.len() * mem::size_of::<u16>()).try_into()?;
attribs.ServerPrincipalNameBufferLength = (server_principal_name.len() * size_of::<u16>()).try_into()?;

// SAFETY: No preconditions.
let status = unsafe { RpcServerInqCallAttributesW(Some(self.0), &mut attribs as *mut _ as *mut c_void) };
Expand Down Expand Up @@ -161,14 +161,14 @@ impl RpcServerInterfacePointer {
let addr = unsafe { raw_dispatch_table.add(i) }.cast_mut();

// SAFETY: Assume the address points to a valid handler which is not currently in use.
let old_prot = unsafe { set_memory_protection(addr.cast(), mem::size_of::<*const ()>(), PAGE_READWRITE) }?;
let old_prot = unsafe { set_memory_protection(addr.cast(), size_of::<*const ()>(), PAGE_READWRITE) }?;

// TODO: See if it could be possible to freeze other threads during switch or to do an atomic switch.
// SAFETY: Because of previous assumption and memory protection, this should succeed.
unsafe { *addr = *new_handler };

// SAFETY: Address is already assumed to be valid.
let _ = unsafe { set_memory_protection(addr.cast(), mem::size_of::<*const ()>(), old_prot) }?;
let _ = unsafe { set_memory_protection(addr.cast(), size_of::<*const ()>(), old_prot) }?;
}

Ok(())
Expand Down
17 changes: 8 additions & 9 deletions crates/win-api-wrappers/src/security/acl.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
use std::alloc::Layout;
use std::mem::{self};
use std::{ptr, slice};

use anyhow::{bail, Result};
Expand Down Expand Up @@ -83,7 +82,7 @@ impl Ace {
};

// SAFETY: We are adding to the pointer in byte aligned mode to access next field.
ptr = unsafe { ptr.byte_add(mem::size_of::<ACE_HEADER>()) };
ptr = unsafe { ptr.byte_add(size_of::<ACE_HEADER>()) };

#[allow(clippy::cast_ptr_alignment)] // FIXME(DGW-221): Raw* hack is flawed.
// SAFETY: Buffer is at least `size_of::<ACE_HEADER> + size_of::<u32>` big.
Expand All @@ -92,7 +91,7 @@ impl Ace {
};

// SAFETY: We are adding to the pointer in byte aligned mode to access next field.
ptr = unsafe { ptr.byte_add(mem::size_of::<u32>()) };
ptr = unsafe { ptr.byte_add(size_of::<u32>()) };

// SAFETY: Buffer is at least `size_of::<ACE_HEADER> + size_of::<u32> + body.len()` big.
unsafe { ptr.copy_from(body.as_ptr(), body.len()) };
Expand All @@ -109,21 +108,21 @@ impl Ace {
// SAFETY: Assume that the pointer points to a valid ACE_HEADER if not null.
let header = unsafe { ptr.as_ref() }.ok_or_else(|| Error::NullPointer("ACE header"))?;

if (header.AceSize as usize) < mem::size_of::<ACE_HEADER>() + mem::size_of::<u32>() {
if (header.AceSize as usize) < size_of::<ACE_HEADER>() + size_of::<u32>() {
bail!(Error::from_win32(ERROR_INVALID_DATA));
}

// SAFETY: Assume that the header is followed by a 4 byte access mask.
ptr = unsafe { ptr.byte_add(mem::size_of::<ACE_HEADER>()) };
ptr = unsafe { ptr.byte_add(size_of::<ACE_HEADER>()) };

// SAFETY: Assume that the header is followed by a 4 byte access mask.
#[allow(clippy::cast_ptr_alignment)] // FIXME(DGW-221): Raw* hack is flawed.
let access_mask = unsafe { ptr.cast::<u32>().read() };

// SAFETY: Assume buffer is big enough to fit Ace data.
ptr = unsafe { ptr.byte_add(mem::size_of::<u32>()) };
ptr = unsafe { ptr.byte_add(size_of::<u32>()) };

let body_size = header.AceSize as usize - mem::size_of::<ACE_HEADER>() - mem::size_of::<u32>();
let body_size = header.AceSize as usize - size_of::<ACE_HEADER>() - size_of::<u32>();

// SAFETY: `body_size` must be >= 0 because of previous check. Pointer is valid.
let body = unsafe { slice::from_raw_parts(ptr.cast::<u8>(), body_size) };
Expand Down Expand Up @@ -158,10 +157,10 @@ impl Acl {

pub fn to_raw(&self) -> Result<Vec<u8>> {
let raw_aces = self.aces.iter().map(Ace::to_raw).collect::<Result<Vec<_>>>()?;
let size = mem::size_of::<ACL>() + raw_aces.iter().map(Vec::len).sum::<usize>();
let size = size_of::<ACL>() + raw_aces.iter().map(Vec::len).sum::<usize>();

// Align on u32 boundary
let size = (size + mem::size_of::<u32>() - 1) & !3;
let size = (size + size_of::<u32>() - 1) & !3;

let mut buf = vec![0; size];

Expand Down
9 changes: 4 additions & 5 deletions crates/win-api-wrappers/src/thread.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
use std::ffi::c_void;
use std::fmt::Debug;
use std::mem::{self};

use anyhow::{bail, Result};

Expand Down Expand Up @@ -32,7 +31,7 @@ impl Thread {
pub fn get_by_id(id: u32, desired_access: THREAD_ACCESS_RIGHTS) -> Result<Self> {
// SAFETY: No preconditions.
let handle = unsafe { OpenThread(desired_access, false, id)? };
// SAFETY: The handle is owned by us, we opened the ressource above.
// SAFETY: The handle is owned by us, we opened the resource above.
let handle = unsafe { Handle::new_owned(handle)? };

Ok(Self::from(handle))
Expand Down Expand Up @@ -172,9 +171,9 @@ impl ThreadAttributeType<'_> {

pub fn size(&self) -> usize {
match self {
ThreadAttributeType::ParentProcess(_) => mem::size_of::<HANDLE>(),
ThreadAttributeType::ExtendedFlags(_) => mem::size_of::<u32>(),
ThreadAttributeType::HandleList(h) => mem::size_of::<HANDLE>() * h.len(),
ThreadAttributeType::ParentProcess(_) => size_of::<HANDLE>(),
ThreadAttributeType::ExtendedFlags(_) => size_of::<u32>(),
ThreadAttributeType::HandleList(h) => size_of::<HANDLE>() * h.len(),
}
}
}
Loading

0 comments on commit caa5ffa

Please sign in to comment.