MemoryCheck/main.cpp

#include <iostream>
#include <stdio.h>
#include <fstream>
#include <Windows.h>
#include <string>
#include <tlhelp32.h>
#include <sstream>
#include <tchar.h>
#include "utils/StringHelper.h"
#include "src/MemoryCheck.h"
#include "easylogging++.h"

INITIALIZE_EASYLOGGINGPP

// Windows服务的服务名
char SERVICE_NAME[] = "MemoryCheck";
SERVICE_STATUS_HANDLE g_ServiceStatusHandle;
SERVICE_STATUS g_ServiceStatus;
DWORD g_ThreadID;
std::wofstream g_ofs;

// 服务主函数
VOID WINAPI ServiceMain(DWORD argc, LPWSTR* argv);

// 服务控制处理函数
VOID WINAPI ServiceCtrlHandler(DWORD CtrlCode);


void EasyLogConf()
{
    ////加载默认配置  init log
    //el::Configurations conf;
    //////设置为默认
    ////conf.setToDefault();
    ////设置日志输出格式
    //conf.setGlobally(el::ConfigurationType::Format, "%datetime %level [%fbase|%line] %msg");
    ////设置日志文件目录以及文件名
    //conf.setGlobally(el::ConfigurationType::Filename, "logs\\%datetime{%Y%M%d}.log");
    ////启用日志
    //conf.setGlobally(el::ConfigurationType::Enabled, "true");
    ////是否写入文件
    //conf.setGlobally(el::ConfigurationType::ToFile, "true");
    ////是否输出控制台
    //conf.setGlobally(el::ConfigurationType::ToStandardOutput, "false");
    //el::Loggers::reconfigureAllLoggers(conf);    //加载配置文件
    //el::Loggers::addFlag(el::LoggingFlag::HierarchicalLogging);
    //el::Loggers::setLoggingLevel(el::Level::Debug);   //设置日志等级

    string confPath = StringHelper::GetAppPath() + "\log.conf";
    el::Configurations conf(confPath);
    el::Loggers::reconfigureAllLoggers(conf);
}




int ConsoleMain()
{
    LOG(DEBUG) << "Service is running..." ;
    MemoryCheck::GetInstance().Init();
    MemoryCheck::GetInstance().Start();
 
    MemoryCheck::GetInstance().Stop();
    LOG(DEBUG) << "Service stoped..." << std::endl;
    return 0;
}

VOID WINAPI ServiceMain(DWORD argc, LPWSTR* argv)
{
    LOG(DEBUG) << "ServiceMain START...";
    g_ServiceStatus.dwCurrentState = SERVICE_START_PENDING;
    g_ServiceStatus.dwControlsAccepted = SERVICE_ACCEPT_STOP;
    // 注册服务控制处理函数
    g_ServiceStatusHandle = RegisterServiceCtrlHandler(SERVICE_NAME, ServiceCtrlHandler);
    if (g_ServiceStatusHandle == NULL)
    {
        LOG(ERROR) << "Handler not install";
        return;
    }

    SetServiceStatus(g_ServiceStatusHandle, &g_ServiceStatus);

    g_ServiceStatus.dwWin32ExitCode = S_OK;
    g_ServiceStatus.dwCheckPoint = 0;
    g_ServiceStatus.dwWaitHint = 0;
    g_ServiceStatus.dwCurrentState = SERVICE_RUNNING;
    SetServiceStatus(g_ServiceStatusHandle, &g_ServiceStatus);

    // 在 ServiceMain 中调用控制台程序的主函数
    ConsoleMain();

    // 设置服务状态为停止
    g_ServiceStatus.dwCurrentState = SERVICE_STOPPED;
    g_ServiceStatus.dwCheckPoint = 0;
    g_ServiceStatus.dwWaitHint = 0;

    if (!SetServiceStatus(g_ServiceStatusHandle, &g_ServiceStatus)){
        LOG(ERROR) << "SetServiceStatus failed, error code: " << GetLastError();
    }
    return;
}

void Init()
{
    g_ServiceStatusHandle = NULL;
    g_ServiceStatus.dwServiceType = SERVICE_WIN32_OWN_PROCESS | SERVICE_INTERACTIVE_PROCESS;
    g_ServiceStatus.dwCurrentState = SERVICE_START_PENDING;
    g_ServiceStatus.dwControlsAccepted = SERVICE_ACCEPT_STOP;
    g_ServiceStatus.dwWin32ExitCode = 0;
    g_ServiceStatus.dwServiceSpecificExitCode = 0;
    g_ServiceStatus.dwCheckPoint = 0;
    g_ServiceStatus.dwWaitHint = 0;
}


bool IsInstalled()
{
    SC_HANDLE hSCM = ::OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (hSCM == NULL) {
        LOG(ERROR) << "OpenSCManager error..."<<GetLastError();
        return false;
    }

    SC_HANDLE hService = ::OpenService(hSCM, SERVICE_NAME, SERVICE_QUERY_CONFIG);   
    if (hService == NULL) {
        ::CloseServiceHandle(hSCM);
        DWORD err = GetLastError();
        if (err == ERROR_SERVICE_DOES_NOT_EXIST) {
            LOG(INFO) << "Service is not installed";
            return false;
        }
        else {
            LOG(ERROR) << "OpenService error..." << GetLastError();
            return false;
        }
    }

    ::CloseServiceHandle(hService);
    ::CloseServiceHandle(hSCM);
    return true;
}


bool Install()
{
    if (IsInstalled())return true;
    SC_HANDLE hSCM = ::OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (hSCM == NULL)return false;
    TCHAR szFilePath[MAX_PATH];
    ::GetModuleFileName(NULL, szFilePath, MAX_PATH);

    SC_HANDLE hService = ::CreateService(hSCM, SERVICE_NAME, SERVICE_NAME, SERVICE_ALL_ACCESS,
        SERVICE_WIN32_OWN_PROCESS | SERVICE_INTERACTIVE_PROCESS, SERVICE_AUTO_START, SERVICE_ERROR_NORMAL,
        szFilePath, NULL, NULL, _T(""), NULL, NULL);
    if (hService == NULL) {
        ::CloseServiceHandle(hSCM);
        LOG(ERROR) << "CreateService false...error:"<<GetLastError();
        return false;
    }

    // 创建服务后立即启动
    if (!StartService(hService, 0, NULL)){
        LOG(ERROR)  << "StartService false...error:"<<GetLastError();
    }
    else {
        LOG(DEBUG) << "服务启动成功...";
    }

    ::CloseServiceHandle(hService);
    ::CloseServiceHandle(hSCM);
    return true;
}


bool Uninstall()
{
    SC_HANDLE schSCManager;
    SC_HANDLE schService;

    schSCManager = OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (schSCManager == NULL){
        LOG(ERROR) << " 打开SCManager失败:" << GetLastError();
        return false;
    }

    schService = OpenService(schSCManager, SERVICE_NAME, SERVICE_ALL_ACCESS);
    if (schService == NULL) {
        LOG(ERROR) << "打开服务失败：" << GetLastError();
        CloseServiceHandle(schSCManager);
        return false;
    }

    bool result = true;
    SERVICE_STATUS status;
    ::ControlService(schService, SERVICE_CONTROL_STOP, &status);
    if (!DeleteService(schService)){
        LOG(ERROR) << "删除服务失败：" << GetLastError();
        result = false;
    }
    else{
        LOG(ERROR) << "服务删除成功" ;
    }

    CloseServiceHandle(schService);
    CloseServiceHandle(schSCManager);
    return result;
}



VOID WINAPI ServiceCtrlHandler(DWORD CtrlCode)
{
    switch (CtrlCode)
    {
    case SERVICE_CONTROL_STOP:
        //设置服务状态为停止
        g_ServiceStatus.dwCurrentState = SERVICE_STOP_PENDING;
        g_ServiceStatus.dwCheckPoint = 1;

        if (!SetServiceStatus(g_ServiceStatusHandle, &g_ServiceStatus))
        {
            std::cerr << "SetServiceStatus failed, error code: " << GetLastError() << std::endl;
        }
        Sleep(500);
        // 设置服务状态为停止
        g_ServiceStatus.dwCurrentState = SERVICE_STOPPED;
        g_ServiceStatus.dwCheckPoint = 0;
        g_ServiceStatus.dwWaitHint = 0;

        if (!SetServiceStatus(g_ServiceStatusHandle, &g_ServiceStatus))
        {
            std::cerr << "SetServiceStatus failed, error code: " << GetLastError() << std::endl;
        }

        PostThreadMessage(g_ThreadID, WM_CLOSE, 0, 0);
        break;
    case SERVICE_CONTROL_SHUTDOWN:
        exit(0);
        break;
    default:
        break;
    }

    return;
}



int APIENTRY _tWinMain(HINSTANCE hinstance, HINSTANCE hPrevInstance, LPTSTR lpCmdLine, int nCmdShow){
    EasyLogConf();    //easylog配置

    Init();
    g_ThreadID = ::GetCurrentThreadId();
    SERVICE_TABLE_ENTRY ServiceTable[] =
    {
        { SERVICE_NAME, (LPSERVICE_MAIN_FUNCTION)ServiceMain },
        { NULL, NULL }
    };

    if (strcmp(lpCmdLine, "/install") == 0) {
        LOG(DEBUG) << "install start...";
        Install();
    }
    else if (strcmp(lpCmdLine, "/uninstall") == 0)
    {
        LOG(DEBUG) << "uninstall start...";
        if (!Uninstall()) {
            LOG(DEBUG) << "uninstall error(" << GetLastError() << ")...";
        }
    }
    else {
        if (!::StartServiceCtrlDispatcher(ServiceTable))
        {
            LOG(DEBUG) << "start error(" << GetLastError() << ")...";
        }
    }

    return 0;

}