Skip to content

Commit 92d82cc

Browse files
author
Requiem
committed
perf: improved VM::NVRAM execution speed
1 parent 9833b82 commit 92d82cc

1 file changed

Lines changed: 121 additions & 174 deletions

File tree

‎src/vmaware.hpp‎

Lines changed: 121 additions & 174 deletions
Original file line numberDiff line numberDiff line change
@@ -9456,176 +9456,155 @@ struct VM {
94569456
* @implements VM::NVRAM
94579457
*/
94589458
[[nodiscard]] static bool nvram_vars() {
9459-
struct VARIABLE_NAME {
9460-
ULONG NextEntryOffset;
9461-
GUID VendorGuid;
9462-
WCHAR Name[1];
9463-
};
9459+
struct VARIABLE_NAME { ULONG NextEntryOffset; GUID VendorGuid; WCHAR Name[1]; };
94649460
using PVARIABLE_NAME = VARIABLE_NAME*;
9461+
using NtEnumerateSystemEnvironmentValuesEx_t = NTSTATUS(__stdcall*)(ULONG, PVOID, PULONG);
94659462

9466-
using NtEnumerateSystemEnvironmentValuesEx_t = NTSTATUS(__stdcall*)(
9467-
ULONG InformationClass,
9468-
PVOID Buffer,
9469-
PULONG BufferLength);
9470-
9471-
// secure boot
9472-
bool found_dbDefault = false;
9473-
bool found_dbxDefault = false;
9474-
bool found_KEKDefault = false;
9475-
bool found_PKDefault = false;
9476-
9477-
// extra vars
9463+
bool found_dbDefault = false, found_dbxDefault = false, found_KEKDefault = false, found_PKDefault = false;
94789464
bool found_MORCL = false;
94799465

9480-
if (!util::is_admin()) {
9481-
return false;
9482-
}
9466+
if (!util::is_admin()) return false;
94839467

94849468
HANDLE hToken = nullptr;
9485-
if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken)) {
9486-
return false;
9487-
}
9469+
if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken)) return false;
94889470

9489-
auto token_closer = [&](HANDLE token) {
9490-
if (token) {
9491-
TOKEN_PRIVILEGES tp{};
9492-
tp.PrivilegeCount = 1;
9493-
LookupPrivilegeValue(nullptr, SE_SYSTEM_ENVIRONMENT_NAME, &tp.Privileges[0].Luid);
9494-
tp.Privileges[0].Attributes = 0;
9495-
AdjustTokenPrivileges(token, FALSE, &tp, sizeof(TOKEN_PRIVILEGES), nullptr, nullptr);
9496-
CloseHandle(token);
9471+
LUID luid{};
9472+
bool priv_enabled = false;
9473+
auto cleanup = [&]() {
9474+
if (priv_enabled && hToken) {
9475+
TOKEN_PRIVILEGES tpDisable{};
9476+
tpDisable.PrivilegeCount = 1;
9477+
tpDisable.Privileges[0].Luid = luid;
9478+
tpDisable.Privileges[0].Attributes = 0;
9479+
AdjustTokenPrivileges(hToken, FALSE, &tpDisable, sizeof(TOKEN_PRIVILEGES), nullptr, nullptr);
9480+
}
9481+
if (hToken) {
9482+
CloseHandle(hToken);
9483+
hToken = nullptr;
94979484
}
94989485
};
9499-
std::unique_ptr<void, decltype(token_closer)> token_guard(hToken, token_closer);
95009486

9501-
LUID luid{};
9502-
if (!LookupPrivilegeValue(nullptr, SE_SYSTEM_ENVIRONMENT_NAME, &luid)) {
9503-
return false;
9504-
}
9505-
9506-
TOKEN_PRIVILEGES tp{};
9507-
tp.PrivilegeCount = 1;
9508-
tp.Privileges[0].Luid = luid;
9509-
tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
9510-
9511-
if (!AdjustTokenPrivileges(hToken, FALSE, &tp, sizeof(TOKEN_PRIVILEGES), nullptr, nullptr) || GetLastError() != ERROR_SUCCESS) {
9512-
return false;
9487+
if (!LookupPrivilegeValue(nullptr, SE_SYSTEM_ENVIRONMENT_NAME, &luid)) { cleanup(); return false; }
9488+
9489+
TOKEN_PRIVILEGES tpEnable{};
9490+
tpEnable.PrivilegeCount = 1;
9491+
tpEnable.Privileges[0].Luid = luid;
9492+
tpEnable.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
9493+
AdjustTokenPrivileges(hToken, FALSE, &tpEnable, sizeof(TOKEN_PRIVILEGES), nullptr, nullptr);
9494+
if (GetLastError() != ERROR_SUCCESS) { cleanup(); return false; }
9495+
priv_enabled = true;
9496+
9497+
bool hasFunction = false;
9498+
bool success = false;
9499+
std::vector<BYTE> resBuffer;
9500+
ULONG bufferLength = 0;
9501+
HMODULE ntdll = util::get_ntdll();
9502+
if (ntdll) {
9503+
const char* names[] = { "NtEnumerateSystemEnvironmentValuesEx" };
9504+
void* functions[1] = { nullptr };
9505+
util::get_function_address(ntdll, names, functions, 1);
9506+
auto NtEnum = reinterpret_cast<NtEnumerateSystemEnvironmentValuesEx_t>(functions[0]);
9507+
if (NtEnum) {
9508+
hasFunction = true;
9509+
NTSTATUS status = NtEnum(1, nullptr, &bufferLength);
9510+
if (bufferLength != 0) {
9511+
try { resBuffer.resize(bufferLength); }
9512+
catch (...) { resBuffer.clear(); bufferLength = 0; }
9513+
if (!resBuffer.empty()) {
9514+
status = NtEnum(1, resBuffer.data(), &bufferLength);
9515+
if (status == 0) { success = true; resBuffer.resize(bufferLength); }
9516+
else resBuffer.clear();
9517+
}
9518+
}
9519+
}
95139520
}
95149521

9515-
util::EnumerateFirmwareResult res = util::enumerate_firmware_variables();
9516-
9517-
if (!res.hasFunction) {
9522+
if (!hasFunction) {
95189523
debug("NVRAM: Handle to ntdll.dll could not be obtained, possibly tampered");
9519-
return true; // returning true on purpose
9524+
cleanup();
9525+
return true;
95209526
}
9521-
9522-
if (!res.success) {
9527+
if (!success) {
95239528
debug("NVRAM: System is not UEFI");
9524-
return false; // NOT UEFI
9529+
cleanup();
9530+
return false;
95259531
}
95269532

9527-
auto contains_redhat_ascii_ci = [](const std::vector<BYTE>& buf) -> bool {
9533+
auto contains_redhat_ascii_ci = [](const std::vector<BYTE>& buf)->bool {
95289534
if (buf.empty()) return false;
95299535
std::string s(reinterpret_cast<const char*>(buf.data()), buf.size());
9530-
for (auto& c : s) c = static_cast<char>(::tolower(static_cast<u8>(c)));
9536+
for (auto& c : s) c = static_cast<char>(::tolower(static_cast<unsigned char>(c)));
95319537
return s.find("red hat secure boot") != std::string::npos;
95329538
};
9533-
auto contains_redhat_utf16le_ci = [](const std::vector<BYTE>& buf) -> bool {
9534-
if (buf.size() < 2) return false;
9535-
if (buf.size() % 2 != 0) return false;
9539+
auto contains_redhat_utf16le_ci = [](const std::vector<BYTE>& buf)->bool {
9540+
if (buf.size() < 2 || (buf.size() % 2) != 0) return false;
95369541
const WCHAR* wptr = reinterpret_cast<const WCHAR*>(buf.data());
95379542
size_t wlen = buf.size() / sizeof(WCHAR);
95389543
try {
95399544
std::wstring ws(wptr, wlen);
95409545
for (auto& wc : ws) wc = static_cast<wchar_t>(::towlower(wc));
9541-
std::wstring needle = L"red hat secure boot";
9542-
return ws.find(needle) != std::wstring::npos;
9543-
}
9544-
catch (...) {
9545-
return false;
9546+
return ws.find(L"red hat secure boot") != std::wstring::npos;
95469547
}
9548+
catch (...) { return false; }
95479549
};
95489550

9549-
PVARIABLE_NAME varName = reinterpret_cast<PVARIABLE_NAME>(res.buffer.data());
9550-
const size_t bufSize = res.buffer.size();
9551-
9551+
PVARIABLE_NAME varName = reinterpret_cast<PVARIABLE_NAME>(resBuffer.data());
9552+
const size_t bufSize = resBuffer.size();
95529553
constexpr size_t MAX_NAME_BYTES = 4096;
95539554

95549555
while (true) {
9555-
const uintptr_t basePtr = reinterpret_cast<uintptr_t>(res.buffer.data());
9556+
const uintptr_t basePtr = reinterpret_cast<uintptr_t>(resBuffer.data());
95569557
const uintptr_t curPtr = reinterpret_cast<uintptr_t>(varName);
9557-
95589558
if (curPtr < basePtr) break;
95599559
const size_t offset = static_cast<size_t>(curPtr - basePtr);
95609560
if (offset >= bufSize) break;
95619561

95629562
const size_t nameOffset = offsetof(VARIABLE_NAME, Name);
95639563
if (bufSize - offset < nameOffset) break;
95649564

9565+
// read name
95659566
std::wstring nameStr;
9566-
{
9567-
size_t nameMaxBytes = 0;
9568-
if (varName->NextEntryOffset != 0) {
9569-
const SIZE_T ne = static_cast<SIZE_T>(varName->NextEntryOffset);
9570-
if (ne <= nameOffset) {
9571-
return false;
9572-
}
9573-
if (ne > bufSize - offset) {
9574-
break;
9575-
}
9576-
nameMaxBytes = ne - nameOffset;
9577-
}
9578-
else {
9579-
if (offset + nameOffset >= bufSize) {
9580-
return false;
9581-
}
9582-
nameMaxBytes = bufSize - (offset + nameOffset);
9583-
}
9584-
9585-
if (nameMaxBytes > MAX_NAME_BYTES) nameMaxBytes = MAX_NAME_BYTES;
9586-
9587-
if (nameMaxBytes >= sizeof(WCHAR)) {
9588-
const WCHAR* namePtr = reinterpret_cast<const WCHAR*>(reinterpret_cast<const BYTE*>(varName) + nameOffset);
9589-
const size_t maxChars = nameMaxBytes / sizeof(WCHAR);
9590-
size_t realChars = 0;
9591-
while (realChars < maxChars && namePtr[realChars] != L'\0') ++realChars;
9592-
if (realChars == maxChars) return false;
9593-
nameStr.assign(namePtr, realChars);
9594-
}
9595-
else {
9596-
nameStr.clear();
9597-
}
9567+
size_t nameMaxBytes = 0;
9568+
if (varName->NextEntryOffset != 0) {
9569+
const SIZE_T ne = static_cast<SIZE_T>(varName->NextEntryOffset);
9570+
if (ne <= nameOffset) { cleanup(); return false; }
9571+
if (ne > bufSize - offset) break;
9572+
nameMaxBytes = ne - nameOffset;
95989573
}
9599-
9600-
std::wstring guidStr;
9601-
{
9602-
auto guid_to_wstring = [](const GUID& g) -> std::wstring {
9603-
wchar_t buf[40] = {};
9604-
int written = _snwprintf_s(
9605-
buf, _countof(buf), _TRUNCATE,
9606-
L"{%08lX-%04hX-%04hX-%02X%02X-%02X%02X%02X%02X%02X%02X}",
9607-
static_cast<unsigned long>(g.Data1),
9608-
static_cast<u16>(g.Data2),
9609-
static_cast<u16>(g.Data3),
9610-
static_cast<u32>(g.Data4[0]),
9611-
static_cast<u32>(g.Data4[1]),
9612-
static_cast<u32>(g.Data4[2]),
9613-
static_cast<u32>(g.Data4[3]),
9614-
static_cast<u32>(g.Data4[4]),
9615-
static_cast<u32>(g.Data4[5]),
9616-
static_cast<u32>(g.Data4[6]),
9617-
static_cast<u32>(g.Data4[7])
9618-
);
9619-
if (written <= 0) return std::wstring();
9620-
return std::wstring(buf);
9621-
};
9622-
9623-
guidStr = guid_to_wstring(varName->VendorGuid);
9624-
if (guidStr.empty()) return true;
9574+
else {
9575+
if (offset + nameOffset >= bufSize) { cleanup(); return false; }
9576+
nameMaxBytes = bufSize - (offset + nameOffset);
96259577
}
9578+
if (nameMaxBytes > MAX_NAME_BYTES) nameMaxBytes = MAX_NAME_BYTES;
9579+
if (nameMaxBytes >= sizeof(WCHAR)) {
9580+
const WCHAR* namePtr = reinterpret_cast<const WCHAR*>(reinterpret_cast<const BYTE*>(varName) + nameOffset);
9581+
const size_t maxChars = nameMaxBytes / sizeof(WCHAR);
9582+
size_t realChars = 0;
9583+
while (realChars < maxChars && namePtr[realChars] != L'\0') ++realChars;
9584+
if (realChars == maxChars) { cleanup(); return false; }
9585+
nameStr.assign(namePtr, realChars);
9586+
}
9587+
9588+
auto guid_to_wstring = [](const GUID& g)->std::wstring {
9589+
wchar_t buf[40] = {};
9590+
int written = _snwprintf_s(buf, _countof(buf), _TRUNCATE,
9591+
L"{%08lX-%04hX-%04hX-%02X%02X-%02X%02X%02X%02X%02X%02X}",
9592+
static_cast<unsigned long>(g.Data1),
9593+
static_cast<u16>(g.Data2),
9594+
static_cast<u16>(g.Data3),
9595+
static_cast<u32>(g.Data4[0]), static_cast<u32>(g.Data4[1]),
9596+
static_cast<u32>(g.Data4[2]), static_cast<u32>(g.Data4[3]),
9597+
static_cast<u32>(g.Data4[4]), static_cast<u32>(g.Data4[5]),
9598+
static_cast<u32>(g.Data4[6]), static_cast<u32>(g.Data4[7]));
9599+
if (written <= 0) return std::wstring();
9600+
return std::wstring(buf);
9601+
};
9602+
std::wstring guidStr = guid_to_wstring(varName->VendorGuid);
9603+
if (guidStr.empty()) { cleanup(); return true; }
96269604

96279605
if (!nameStr.empty() && nameStr.rfind(L"VMM", 0) == 0) {
96289606
debug("NVRAM: Detected hypervisor signature");
9607+
cleanup();
96299608
return true;
96309609
}
96319610

@@ -9635,79 +9614,47 @@ struct VM {
96359614
else if (nameStr == L"dbxDefault") found_dbxDefault = true;
96369615

96379616
std::vector<BYTE> valueBuf;
9638-
DWORD readLen = 0;
9639-
DWORD required = 0;
9640-
required = GetFirmwareEnvironmentVariableW(nameStr.c_str(), guidStr.c_str(), nullptr, 0);
9617+
DWORD required = GetFirmwareEnvironmentVariableW(nameStr.c_str(), guidStr.c_str(), nullptr, 0);
96419618
if (required > 0) {
96429619
valueBuf.resize(required);
9643-
readLen = GetFirmwareEnvironmentVariableW(nameStr.c_str(), guidStr.c_str(), valueBuf.data(), required);
9644-
if (readLen == 0) {
9645-
valueBuf.clear();
9646-
}
9647-
else {
9648-
valueBuf.resize(readLen);
9649-
}
9620+
DWORD readLen = GetFirmwareEnvironmentVariableW(nameStr.c_str(), guidStr.c_str(), valueBuf.data(), required);
9621+
if (readLen == 0) valueBuf.clear(); else valueBuf.resize(readLen);
96509622
}
96519623
else {
96529624
const DWORD fallbackSize = 8192;
96539625
valueBuf.resize(fallbackSize);
9654-
readLen = GetFirmwareEnvironmentVariableW(nameStr.c_str(), guidStr.c_str(), valueBuf.data(), fallbackSize);
9655-
if (readLen > 0) {
9656-
valueBuf.resize(readLen);
9657-
}
9658-
else {
9659-
valueBuf.clear();
9660-
}
9626+
DWORD readLen = GetFirmwareEnvironmentVariableW(nameStr.c_str(), guidStr.c_str(), valueBuf.data(), fallbackSize);
9627+
if (readLen > 0) valueBuf.resize(readLen); else valueBuf.clear();
96619628
}
96629629

9663-
if (nameStr == L"MemoryOverwriteRequestControlLock") {
9664-
found_MORCL = true;
9665-
}
9630+
if (nameStr == L"MemoryOverwriteRequestControlLock") found_MORCL = true;
96669631

96679632
if (nameStr == L"PKDefault") {
96689633
bool pk_has_redhat = false;
9669-
if (!valueBuf.empty()) {
9670-
if (contains_redhat_utf16le_ci(valueBuf) || contains_redhat_ascii_ci(valueBuf)) {
9671-
pk_has_redhat = true;
9672-
}
9673-
}
9634+
if (!valueBuf.empty() && (contains_redhat_utf16le_ci(valueBuf) || contains_redhat_ascii_ci(valueBuf)))
9635+
pk_has_redhat = true;
96749636
if (pk_has_redhat) {
96759637
debug("NVRAM: QEMU detected");
9676-
return core::add(brands::QEMU);
9638+
bool added = core::add(brands::QEMU);
9639+
cleanup();
9640+
return added;
96779641
}
96789642
}
96799643

96809644
if (varName->NextEntryOffset == 0) break;
9681-
96829645
const SIZE_T ne = static_cast<SIZE_T>(varName->NextEntryOffset);
96839646
const size_t nextOffset = offset + ne;
96849647
if (nextOffset <= offset || nextOffset > bufSize) break;
9685-
9686-
varName = reinterpret_cast<PVARIABLE_NAME>(reinterpret_cast<PBYTE>(res.buffer.data()) + nextOffset);
9648+
varName = reinterpret_cast<PVARIABLE_NAME>(reinterpret_cast<PBYTE>(resBuffer.data()) + nextOffset);
96879649
}
96889650

9689-
if (!found_MORCL) {
9690-
debug("NVRAM: Missing MemoryOverwriteRequestControlLock");
9691-
return true;
9692-
}
9693-
9694-
if (!found_dbDefault) {
9695-
debug("NVRAM: Missing dbDefault");
9696-
return true;
9697-
}
9698-
if (!found_dbxDefault) {
9699-
debug("NVRAM: Missing dbxDefault");
9700-
return true;
9701-
}
9702-
if (!found_KEKDefault) {
9703-
debug("NVRAM: Missing KEKDefault");
9704-
return true;
9705-
}
9706-
if (!found_PKDefault) {
9707-
debug("NVRAM: Missing PKDefault");
9708-
return true;
9709-
}
9651+
if (!found_MORCL) { debug("NVRAM: Missing MemoryOverwriteRequestControlLock"); cleanup(); return true; }
9652+
if (!found_dbDefault) { debug("NVRAM: Missing dbDefault"); cleanup(); return true; }
9653+
if (!found_dbxDefault) { debug("NVRAM: Missing dbxDefault"); cleanup(); return true; }
9654+
if (!found_KEKDefault) { debug("NVRAM: Missing KEKDefault"); cleanup(); return true; }
9655+
if (!found_PKDefault) { debug("NVRAM: Missing PKDefault"); cleanup(); return true; }
97109656

9657+
cleanup();
97119658
return false;
97129659
}
97139660

0 commit comments

Comments
 (0)