driver.cpp

#include "driver.h"
#include <filesystem>
#include <bcrypt.h>
#include <fstream>
#include <cwctype>
#include <sstream>
#include <iomanip>
#include <chrono>
#include <random>

#define PAGE_SIZE 0x1000 

#pragma comment(lib, "advapi32.lib")
#pragma comment(lib, "bcrypt.lib")

std::wstring generate_hex_wstr()
{
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<unsigned int> dis(0x10000, 0xFFFFF); // 限制在 5 位十六进制范围

    std::wstringstream wss;
    wss << std::hex << std::setw(5) << std::setfill(L'0') << dis(gen);

    return wss.str();
}
//必须要随机化，不然有bug
std::wstring svc_name = L"WpnUserService_" + generate_hex_wstr();

BOOL ReadPhysicalMemory(HANDLE hDevice, LARGE_INTEGER physicalAddress, DWORD sizeToRead, std::vector<BYTE>& outBuffer) {
    if (hDevice == INVALID_HANDLE_VALUE || sizeToRead == 0) return FALSE;

    DWORD pageOffset = physicalAddress.QuadPart % PAGE_SIZE;
    LARGE_INTEGER alignedAddress; alignedAddress.QuadPart = physicalAddress.QuadPart - pageOffset;
    DWORD totalRequiredLength = pageOffset + sizeToRead;
    DWORD alignedSize = ((totalRequiredLength + PAGE_SIZE - 1) / PAGE_SIZE) * PAGE_SIZE;

    BYTE* tempBuffer = (BYTE*)VirtualAlloc(NULL, alignedSize, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
    if (!tempBuffer) return FALSE;

    OVERLAPPED ol = {0};
    ol.Offset = alignedAddress.LowPart;
    ol.OffsetHigh = alignedAddress.HighPart;

    DWORD bytesRead = 0;
    BOOL success = ReadFile(hDevice, tempBuffer, alignedSize, &bytesRead, &ol);

    if (success && bytesRead >= totalRequiredLength) {
        outBuffer.resize(sizeToRead);
        std::memcpy(outBuffer.data(), tempBuffer + pageOffset, sizeToRead);
    }else{
        success = FALSE;
    }

    VirtualFree(tempBuffer, 0, MEM_RELEASE);
    return success;
}

Bytes phys_read(HANDLE hDevice, uint64_t addr, uint64_t size) {

    LARGE_INTEGER testAddress; 
    testAddress.QuadPart = addr;
    std::vector<BYTE> resultData;
    if (!ReadPhysicalMemory(hDevice, testAddress, size, resultData)) {
        printf("Read_PhysicalMemory fail \n");
    }
    return (resultData.size() >= size) ? resultData : Bytes{};
}
//this is have a bug，if run too many time if wail fail, better it to random to gen svc_name
bool load_driver(std::string filepath = "")
{
    auto path =
        std::filesystem::absolute(
            filepath.empty() ? "Magnet.sys" : filepath
        ).wstring();

    //constexpr wchar_t svc_name[] =
    //    L"AMD Crash DriverServie";


    SC_HANDLE h_scm =
        OpenSCManagerW(
            nullptr,
            nullptr,
            SC_MANAGER_ALL_ACCESS);

    if (!h_scm)
    {
        printf("OpenSCManager failed: %lu\n",
            GetLastError());
        return false;
    }
    
    SC_HANDLE h_svc =
        CreateServiceW(
            h_scm,
            svc_name.c_str(),
            svc_name.c_str(),
            SERVICE_ALL_ACCESS,
            SERVICE_KERNEL_DRIVER,
            SERVICE_DEMAND_START,
            SERVICE_ERROR_NORMAL,
            path.c_str(),
            nullptr,
            nullptr,
            nullptr,
            nullptr,
            nullptr);

    if (!h_svc)
    {
        DWORD err = GetLastError();

        printf("CreateServiceW failed: %lu\n", err);

        if (err == ERROR_SERVICE_EXISTS ||
            err == ERROR_SERVICE_MARKED_FOR_DELETE)
        {
            h_svc =
                OpenServiceW(
                    h_scm,
                    svc_name.c_str(),
                    SERVICE_ALL_ACCESS);

            if (!h_svc)
            {
                printf("OpenServiceW failed: %lu\n",
                    GetLastError());

                CloseServiceHandle(h_scm);
                return false;
            }
        }
        else
        {
            CloseServiceHandle(h_scm);
            return false;
        }
    }

    BOOL ok =
        StartServiceW(
            h_svc,
            0,
            nullptr);

    if (!ok)
    {
        DWORD err = GetLastError();

        printf("StartServiceW failed: %lu\n", err);

        if (err == ERROR_SERVICE_ALREADY_RUNNING)
        {
            ok = TRUE;
        }
    }

    CloseServiceHandle(h_svc);
    CloseServiceHandle(h_scm);
    return ok == TRUE;
}


//this is have a bug，i donot know why it sometime will fail
bool stop_driver()
{
    SC_HANDLE h_scm =
        OpenSCManager(nullptr, nullptr, SC_MANAGER_ALL_ACCESS);

    if (!h_scm)
        return false;
    

    SC_HANDLE h_svc =
        OpenServiceW(
            h_scm,
            svc_name.c_str(),
            SERVICE_ALL_ACCESS);

    if (!h_svc)
    {
        CloseServiceHandle(h_scm);
        return false;
    }

    SERVICE_STATUS status;
    ControlService(
        h_svc,
        SERVICE_CONTROL_STOP,
        &status);

    DeleteService(h_svc);

    CloseServiceHandle(h_svc);
    CloseServiceHandle(h_scm);

    return true;
}