minifilter 0.0.0

#![no_std]

mod sys;
use core::marker::*;
use core::mem::*;
use core::ptr::*;
use core::sync::atomic::*;

use sys::*;

static G_FILTER: AtomicPtr<core::ffi::c_void> = AtomicPtr::new(null_mut());

#[cfg(not(test))]
mod panic_handler {
    use super::*;
    #[panic_handler]
    fn panic(_info: &core::panic::PanicInfo) -> ! {
        unsafe { DbgPrint("PANIC!\n".as_ptr()) };
        loop {}
    }
}

struct ContextRc {
    context: POpaque,
}

impl ContextRc {
    pub fn wrap(context: POpaque) -> ContextRc {
        ContextRc { context }
    }
    pub fn raw(&self) -> POpaque {
        self.context
    }
}

impl Clone for ContextRc {
    fn clone(&self) -> Self {
        unsafe { FltReferenceContext(self.context) }
        ContextRc {
            context: self.context,
        }
    }
}

impl Drop for ContextRc {
    fn drop(&mut self) {
        unsafe { FltReleaseContext(self.context) }
    }
}

trait ContextTypeTag {
    fn sys_context_type() -> u16;
    fn sys_tag() -> u32;
}

#[derive(Clone)]
struct InstanceContextTag;
impl ContextTypeTag for InstanceContextTag {
    fn sys_context_type() -> u16 {
        2
    }
    fn sys_tag() -> u32 {
        0x40
    }
}

#[derive(Clone)]
struct Context<T, ContextType> {
    rc: ContextRc,
    phantom_t: core::marker::PhantomData<T>,
    phantom_context: core::marker::PhantomData<ContextType>,
}

extern "C" fn context_cleanup<T>(context: POpaque, _context_type: u16) {
    unsafe { drop_in_place(context as *mut T) }
}

impl<T: 'static + Send + Sync, ContextType: ContextTypeTag> Context<T, ContextType> {
    fn context_size() -> u64 {
        // TODO: align
        let size = size_of::<T>();
        size as u64
    }

    fn registration() -> FLT_CONTEXT_REGISTRATION {
        let cleanup = if needs_drop::<T>() {
            context_cleanup::<T> as *const _
        } else {
            null()
        };

        FLT_CONTEXT_REGISTRATION {
            ContextType: ContextType::sys_context_type(),
            Flags: 0,
            ContextCleanupCallback: cleanup,
            Size: Self::context_size(),
            PoolTag: ContextType::sys_tag(),
            ContextAllocateCallback: null(),
            ContextFreeCallback: null(),
            Reserved1: null_mut(),
        }
    }

    pub fn wrap(context: POpaque) -> Self {
        Context {
            rc: ContextRc { context },
            phantom_t: PhantomData,
            phantom_context: PhantomData,
        }
    }

    pub fn new(filter: POpaque, value: T, pool_type: i32) -> Result<Self, (T, NTSTATUS)> {
        unsafe {
            // TODO: align
            let mut context = null_mut();
            let status = FltAllocateContext(
                filter,
                ContextType::sys_context_type(),
                Self::context_size(),
                pool_type,
                addr_of_mut!(context),
            );
            if status < 0 {
                return Err((value, status));
            }
            write(context as *mut T, value);
            Ok(Context {
                rc: ContextRc { context },
                phantom_t: PhantomData,
                phantom_context: PhantomData,
            })
        }
    }

    pub fn get(&self) -> &T {
        unsafe { &*(self.rc.context as *const T) }
    }

    pub fn raw(&self) -> POpaque {
        self.rc.raw()
    }
}

struct InstanceContextInner {
    total_creates: AtomicU32,
}

type InstanceContext = Context<InstanceContextInner, InstanceContextTag>;

struct CallbackPackage {
    data: POpaque,
    flt_objects: *const FLT_RELATED_OBJECTS,
}

enum PreOpResult<T> {
    SuccessWithCallback(T),
    SuccessNoCallback,
}

enum PostOpResult {
    FinishedProcessing,
}

trait CompletionContext {
    fn into_ptr(self) -> POpaque;
    fn from_ptr(ptr: POpaque) -> Self;
}

struct PoolBox<T> {
    raw: *mut T,
    phantom: PhantomData<T>,
}

impl<T> CompletionContext for PoolBox<T> {
    fn into_ptr(self) -> POpaque {
        let raw = self.raw as POpaque;
        forget(self);
        raw
    }
    fn from_ptr(ptr: POpaque) -> Self {
        Self {
            raw: ptr as *mut T,
            phantom: PhantomData,
        }
    }
}

impl<T> Drop for PoolBox<T> {
    fn drop(&mut self) {
        unsafe { ExFreePoolWithTag(self.raw as POpaque, 0) }
    }
}

impl<T> PoolBox<T> {
    fn from_raw(raw: *mut T) -> Self {
        Self {
            raw,
            phantom: PhantomData,
        }
    }

    fn get_raw(&self) -> *mut T {
        self.raw
    }
}

trait FilterInterface {
    type CreateCompletionContext: CompletionContext;
    fn pre_create(callback_package: &CallbackPackage)
        -> PreOpResult<Self::CreateCompletionContext>;
    fn post_create(
        callback_package: &CallbackPackage,
        completion_context: Self::CreateCompletionContext,
    ) -> PostOpResult;
}

struct FilterImpl;

impl FilterInterface for FilterImpl {
    type CreateCompletionContext = PoolBox<UNICODE_STRING>;
    fn pre_create(
        callback_package: &CallbackPackage,
    ) -> PreOpResult<Self::CreateCompletionContext> {
        unsafe {
            if (*(*callback_package.flt_objects).FileObject).Flags & 0x00400000 != 0 {
                // volume open
                return PreOpResult::SuccessNoCallback;
            }

            let mut raw_context: POpaque = null_mut();
            let status = FltGetInstanceContext(
                (*callback_package.flt_objects).Instance,
                addr_of_mut!(raw_context),
            );
            if status < 0 {
                DbgPrint(
                    "pre_create: FltGetInstanceContext failed (0x%x)\n\0".as_ptr(),
                    status,
                );
                return PreOpResult::SuccessNoCallback;
            }
            let instance_context = InstanceContext::wrap(raw_context);
            let total_creates = instance_context
                .get()
                .total_creates
                .fetch_add(1, Ordering::SeqCst)
                + 1;
            DbgPrint("pre_create: %u created so far\n\0".as_ptr(), total_creates);

            let process = FltGetRequestorProcess(callback_package.data);
            let mut image_file_name: *mut UNICODE_STRING = null_mut();
            let status = SeLocateProcessImageName(process, addr_of_mut!(image_file_name));
            if status < 0 {
                DbgPrint(
                    "pre_create: SeLocateProcessImageName  failed (0x%x)\n\0".as_ptr(),
                    status,
                );
                return PreOpResult::SuccessNoCallback;
            }
            let image_file_name = PoolBox::from_raw(image_file_name);

            DbgPrint(
                "pre_create: process %wZ\n\0".as_ptr(),
                image_file_name.get_raw(),
            );

            PreOpResult::SuccessWithCallback(image_file_name)
        }
    }

    fn post_create(
        callback_package: &CallbackPackage,
        image_file_name: Self::CreateCompletionContext,
    ) -> PostOpResult {
        unsafe {
            DbgPrint(
                "post_create: process %wZ\n\0".as_ptr(),
                image_file_name.get_raw(),
            );

            PostOpResult::FinishedProcessing
        }
    }
}

struct Filter<Interface> {
    filter: POpaque,
    marker: core::marker::PhantomData<Interface>,
}

impl<Interface> Filter<Interface> {
    unsafe fn start(self) -> Result<(), NTSTATUS> {
        let status = FltStartFiltering(self.filter);
        if status < 0 {
            FltUnregisterFilter(self.filter);
            return Err(status);
        }

        Ok(())
    }
}

unsafe extern "C" fn flt_unload(_flags: u32) -> NTSTATUS {
    DbgPrint("flt_unload\n\0".as_ptr());
    FltUnregisterFilter(G_FILTER.load(core::sync::atomic::Ordering::Acquire));
    0
}

unsafe extern "C" fn flt_instance_setup(
    flt_objects: *const FLT_RELATED_OBJECTS,
    _flags: u32,
    _volume_device_type: u32,
    volume_filesystem_type: i32,
) -> NTSTATUS {
    if volume_filesystem_type != 3 {
        DbgPrint(
            "flt_instance_setup: File System is not FAT (0x%x), ignoring\n\0".as_ptr(),
            volume_filesystem_type,
        );
        return 0xC01C000Fu32 as NTSTATUS;
    }

    DbgPrint("flt_instance_setup: Attaching to FAT volume\n\0".as_ptr());

    let context = match InstanceContext::new(
        (*flt_objects).Filter,
        InstanceContextInner {
            total_creates: AtomicU32::new(0),
        },
        1,
    ) {
        Ok(context) => context,
        Err((_, status)) => {
            DbgPrint(
                "flt_instance_setup: context allocation failed (0x%x)\n\0".as_ptr(),
                status,
            );
            return status;
        }
    };

    let status = FltSetInstanceContext((*flt_objects).Instance, 0, context.raw(), null_mut());

    if status < 0 {
        DbgPrint(
            "flt_instance_setup: FltSetInstanceContext failed (0x%x)\n\0".as_ptr(),
            status,
        );
        return status;
    }

    0
}

unsafe extern "C" fn flt_query_teardown(flt_objects: POpaque, _flags: u32) -> NTSTATUS {
    DbgPrint("flt_query_teardown\n\0".as_ptr());

    0
}

unsafe extern "C" fn pre_create<Interface: FilterInterface>(
    data: POpaque,
    flt_objects: *const FLT_RELATED_OBJECTS,
    completion_context: *mut POpaque,
) -> i32 {
    let callback_package = CallbackPackage { data, flt_objects };

    match Interface::pre_create(&callback_package) {
        PreOpResult::SuccessWithCallback(completion_context_value) => {
            *completion_context = completion_context_value.into_ptr();
            0
        }
        PreOpResult::SuccessNoCallback => 1,
    }
}

unsafe extern "C" fn post_create<Interface: FilterInterface>(
    data: POpaque,
    flt_objects: *const FLT_RELATED_OBJECTS,
    completion_context: POpaque,
    flags: u32,
) -> i32 {
    let callback_package = CallbackPackage { data, flt_objects };
    let completion_context = Interface::CreateCompletionContext::from_ptr(completion_context);

    match Interface::post_create(&callback_package, completion_context) {
        PostOpResult::FinishedProcessing => 0,
    }
}

unsafe fn register_filter<Interface: FilterInterface>(
    driver_object: POpaque,
) -> Result<Filter<Interface>, NTSTATUS> {
    let context_registration = [
        <InstanceContext>::registration(),
        FLT_CONTEXT_REGISTRATION {
            ContextType: 0xFFFF,
            Flags: 0,
            ContextCleanupCallback: core::ptr::null(),
            Size: 0,
            PoolTag: 0,
            ContextAllocateCallback: core::ptr::null(),
            ContextFreeCallback: core::ptr::null(),
            Reserved1: core::ptr::null_mut(),
        },
    ];
    let operation_registration = [
        FLT_OPERATION_REGISTRATION {
            MajorFunction: 0, //create
            Flags: 0,
            PreOperation: pre_create::<Interface> as *const _,
            PostOperation: post_create::<Interface> as *const _,
            Reserved1: core::ptr::null_mut(),
        },
        FLT_OPERATION_REGISTRATION {
            MajorFunction: 0x80,
            Flags: 0,
            PreOperation: core::ptr::null(),
            PostOperation: core::ptr::null(),
            Reserved1: core::ptr::null_mut(),
        },
    ];

    let registration = FLT_REGISTRATION {
        Size: core::mem::size_of::<FLT_REGISTRATION>() as u16,
        Version: 0x0202, // ehh
        Flags: 0,
        ContextRegistration: context_registration.as_ptr(),
        OperationRegistration: operation_registration.as_ptr(),
        FilterUnloadCallback: flt_unload as *const _,
        InstanceSetupCallback: flt_instance_setup as *const _,
        InstanceQueryTeardownCallback: flt_query_teardown as *const _,
        InstanceTeardownStartCallback: core::ptr::null(),
        InstanceTeardownCompleteCallback: core::ptr::null(),
        GenerateFileNameCallback: core::ptr::null(),
        NormalizeNameComponentCallback: core::ptr::null(),
        NormalizeContextCleanupCallback: core::ptr::null(),
        TransactionNotificationCallback: core::ptr::null(),
        NormalizeNameComponentExCallback: core::ptr::null(),
        // SectionNotificationCallback: core::ptr::null(),
    };

    let mut filter: POpaque = core::ptr::null_mut();
    let status = FltRegisterFilter(driver_object, addr_of!(registration), addr_of_mut!(filter));
    if status < 0 {
        return Err(status);
    }
    G_FILTER.store(filter, core::sync::atomic::Ordering::Release);
    Ok(Filter {
        filter,
        marker: core::marker::PhantomData,
    })
}

#[export_name = "entry"]
pub unsafe extern "C" fn entry(driver_object: POpaque) -> NTSTATUS {
    DbgPrint("Silverwind - Loaded\n\0".as_ptr());

    let filter = match register_filter::<FilterImpl>(driver_object) {
        Ok(filter) => filter,
        Err(status) => {
            DbgPrint("register failed! 0x%x\n\0".as_ptr(), status);
            return status;
        }
    };

    if let Err(status) = filter.start() {
        DbgPrint("start failed! 0x%x\n\0".as_ptr(), status);
        return status;
    }

    0
}