使用服务做守护进程

转自：http://www.godebug.org/index.php/archives/5/#comment-36

建完blog后第一次正式写文章，有点小激动

在工作时上头要求写一个服务程序，当机器启动时这个服务会自动启动某个程序，并且在程序挂掉的时候自动把他开起来。说白了就是个守护进程，当时我觉得写成服务比较麻烦，普通方法创建的子进程跟服务进程一样也是system权限显示界面很烦人，不如在注册表的run下面加点东西直接完事，可惜咱是一届码农头头不听咱的，最终还是写成了服务。

在写的时候发现有不少人也在找这方面的代码，所以写完了就索性发到刚建成的blog上方便后来人，也方便自己查找。

就一个cpp文件就不放github了。

代码：

#define _WIN32_WINNT 0x502
#define _CRT_SECURE_NO_WARNINGS
#include <Windows.h>
#include <Shlwapi.h>
#include <WtsApi32.h>
#include <UserEnv.h>
#include <iostream>
#include <fstream>
#include <atlbase.h>

static const char* const lpServiceName = "ProtectService";
static SERVICE_STATUS_HANDLE hServiceStatus = NULL;
static SERVICE_STATUS   ServiceStatus = {0};
static char szCurDir[MAX_PATH+1] = {0};
static bool bRun = false;
static HANDLE hProcess = NULL;
static FILE* fLog = NULL;
static std::ofstream ofs_log;

bool InstallService();
VOID WINAPI ServiceMain(DWORD dwArgc,LPTSTR *lpszArgv);
VOID WINAPI HandlerFunc(DWORD dwControl);
HANDLE RunAsLoggedUser(const char* lpPath,char* lpCmdLine);
void WorkFunc();

int main(int argc, char **argv)
{
    GetModuleFileName(NULL,szCurDir,MAX_PATH);
    *strrchr(szCurDir,'\\') = '\0';

    char szLogPath[MAX_PATH+1];
    sprintf(szLogPath,"%s\\Service.log",szCurDir);
    fLog = fopen(szLogPath,"a+");
    setvbuf(fLog,NULL,_IONBF,1024);

    ofs_log.open(szLogPath,std::ios::app);

    SERVICE_TABLE_ENTRY ServiceTable[2];
    char szBuffer[30];
    strcpy(szBuffer,lpServiceName);
    ServiceTable[0].lpServiceName = szBuffer;
    ServiceTable[0].lpServiceProc = (LPSERVICE_MAIN_FUNCTION)ServiceMain;

    ServiceTable[1].lpServiceName = NULL;
    ServiceTable[1].lpServiceProc = NULL;

    // 启动服务的控制分派机线程
    if (!StartServiceCtrlDispatcher(ServiceTable))
    {
        ofs_log<<"程序不是以服务方式启动"<<std::endl;
        std::cout<<"程序不是以服务方式启动，要创建服务并启动吗？y/n:";
        if (getchar() == 'y')
        {
            // 安装服务
            if (!InstallService())
            {
                ofs_log<<"安装服务失败"<<std::endl;
                std::cout<<"安装服务失败，具体原因请查看日志"<<std::endl;
                return 1;
            }
            ofs_log<<"安装服务成功"<<std::endl;
        }

    }
    return 0;
}

bool InstallService()
{
    SC_HANDLE hSCManager = OpenSCManager(NULL,NULL,SC_MANAGER_ALL_ACCESS);
    if (!hSCManager)
    {
        ofs_log<<"OpenSCManager失败，错误码为："<<GetLastError()<<std::endl;
        return false;
    }

    SC_HANDLE hService = OpenService(hSCManager,lpServiceName,SERVICE_QUERY_CONFIG);
    if (hService)
    {
        ofs_log<<"服务已经存在"<<std::endl;
        CloseServiceHandle(hSCManager);
        CloseServiceHandle(hService);
        return false;
    }

    char    szPath[MAX_PATH+1];
    GetModuleFileName(NULL,szPath,MAX_PATH);
    hService = CreateService(hSCManager,
        lpServiceName,
        lpServiceName,
        SERVICE_ALL_ACCESS,
        SERVICE_WIN32_OWN_PROCESS | SERVICE_INTERACTIVE_PROCESS,
        SERVICE_AUTO_START,
        SERVICE_ERROR_NORMAL,
        szPath,
        NULL,
        NULL,
        NULL,
        NULL,
        NULL);
    if (!hService)
    {
        ofs_log<<"CreateService失败，错误码为"<<GetLastError()<<std::endl;
        CloseServiceHandle(hSCManager);
        return false;
    }

    if (!StartService(hService,0,NULL))
    {
        ofs_log<<"服务安装成功，但是启动服务失败，错误码为："<<GetLastError()<<std::endl;
        CloseServiceHandle(hService);
        CloseServiceHandle(hSCManager);
        return false;
    }

    CloseServiceHandle(hService);
    CloseServiceHandle(hSCManager);

    return true;
}

VOID WINAPI ServiceMain( DWORD dwArgc,LPTSTR *lpszArgv )
{
    ServiceStatus.dwServiceType = SERVICE_WIN32;
    ServiceStatus.dwCurrentState = SERVICE_RUNNING;
    ServiceStatus.dwControlsAccepted = SERVICE_ACCEPT_SHUTDOWN | SERVICE_ACCEPT_STOP;
    hServiceStatus = RegisterServiceCtrlHandler(lpServiceName,HandlerFunc);
    if (!hServiceStatus)
    {
        ofs_log<<"RegisterServiceCtrlHandler失败，错误码为："<<GetLastError()<<std::endl;
        return;
    }
    if (!SetServiceStatus(hServiceStatus,&ServiceStatus))
    {
        ofs_log<<"SetServiceStatus失败，错误码为："<<GetLastError()<<std::endl;
        return;
    }

    bRun = true;
    WorkFunc();
}

VOID WINAPI HandlerFunc( DWORD dwControl )
{
    ofs_log<<"收到服务消息："<<dwControl<<std::endl;

    if (dwControl == SERVICE_CONTROL_STOP || dwControl == SERVICE_CONTROL_SHUTDOWN)
    {
        // 停止服务
        ofs_log<<"服务停止"<<std::endl;
        ServiceStatus.dwCurrentState = SERVICE_STOPPED;
        bRun = false;
        if (hProcess)
        {
            TerminateProcess(hProcess,0);
        }
    }

    SetServiceStatus(hServiceStatus,&ServiceStatus);
}

HANDLE RunAsCreator(const char* lpPath, char* lpCmdLine)
{
    if (!PathFileExists(lpPath))
    {
        ofs_log<<"指定的要启动的程序不存在"<<std::endl;
        return NULL;
    }

    char szSubProcCurDir[MAX_PATH];
    strcpy(szSubProcCurDir,lpPath);
    *strrchr(szSubProcCurDir,'\\') = '\0';

    STARTUPINFO si = {0};
    PROCESS_INFORMATION pi = {0};
    si.cb = sizeof(si);

    if (!CreateProcess(lpPath,lpCmdLine,NULL,NULL,FALSE,NORMAL_PRIORITY_CLASS | CREATE_NEW_CONSOLE,NULL,szSubProcCurDir,&si,&pi))
    {
        ofs_log<<"CreateProcess失败，错误码为："<<GetLastError()<<std::endl;
        return NULL;
    }

    if (pi.hThread && pi.hThread != INVALID_HANDLE_VALUE)
    {
        CloseHandle(pi.hThread);
    }

    return pi.hProcess;
}

HANDLE RunAsLoggedUser( const char* lpPath,char* lpCmdLine )
{
    if (!PathFileExists(lpPath))
    {
        ofs_log<<"指定的要启动的程序不存在"<<std::endl;
        return NULL;
    }

    DWORD dwSid = WTSGetActiveConsoleSessionId();
    ofs_log<<"当前登录用户ID为："<<dwSid<<std::endl;

    HANDLE hExistingToken = NULL;
    if (!WTSQueryUserToken(dwSid,&hExistingToken))
    {
        ofs_log<<"WTSQueryUserToken失败，错误码为："<<GetLastError()<<std::endl;
        return NULL;
    }

    HANDLE hNewToken = NULL;
    if (!DuplicateTokenEx(hExistingToken,MAXIMUM_ALLOWED,NULL,SecurityIdentification,TokenPrimary,&hNewToken))
    {
        ofs_log<<"DuplicateTokenEx失败，错误码为："<<GetLastError()<<std::endl;
        CloseHandle(hExistingToken);
        return NULL;
    }
    CloseHandle(hExistingToken);

    LPVOID pEnv = NULL;
    DWORD dwCreationFlags = NORMAL_PRIORITY_CLASS | CREATE_NEW_CONSOLE;
    if (CreateEnvironmentBlock(&pEnv,hNewToken,FALSE))
    {
        dwCreationFlags |= CREATE_UNICODE_ENVIRONMENT;
    }

    STARTUPINFO si = {0};
    PROCESS_INFORMATION pi = {0};
    si.cb = sizeof(si);
    si.lpDesktop = "WinSta0\\Default";

    char szSubProcCurDir[MAX_PATH+1];
    strcpy(szSubProcCurDir,lpPath);
    *strrchr(szSubProcCurDir,'\\') = '\0';

    if (!CreateProcessAsUser(hNewToken,lpPath,lpCmdLine,NULL,NULL,FALSE,dwCreationFlags,pEnv,szSubProcCurDir,&si,&pi))
    {
        ofs_log<<"CreateProcessAsUser失败，错误码为："<<GetLastError()<<std::endl;
        if (pEnv)
        {
            DestroyEnvironmentBlock(pEnv);
        }
        return NULL;
    }

    if (pi.hThread && pi.hThread != INVALID_HANDLE_VALUE)
    {
        CloseHandle(pi.hThread);
    }

    if (pEnv)
    {
        DestroyEnvironmentBlock(pEnv);
    }

    return pi.hProcess;
}

HANDLE RunAsSpecifiedUser(const char* lpUserName,const char* lpPassword,const char* lpPath,char* lpCmdLine)
{
    if (!PathFileExists(lpPath))
    {
        ofs_log<<"指定的要启动的程序不存在"<<std::endl;
        return NULL;
    }

    HANDLE hExistingToken = NULL;
    if (!LogonUser(lpUserName,NULL,lpPassword,LOGON32_LOGON_INTERACTIVE,LOGON32_PROVIDER_DEFAULT,&hExistingToken))
    {
        ofs_log<<"LogonUser失败，错误码为："<<GetLastError()<<std::endl;
        return NULL;
    }

    HANDLE hNewToken = NULL;
    if (!DuplicateTokenEx(hExistingToken,MAXIMUM_ALLOWED,NULL,SecurityIdentification,TokenPrimary,&hNewToken))
    {
        ofs_log<<"DuplicateTokenEx失败，错误码为："<<GetLastError()<<std::endl;
        CloseHandle(hExistingToken);
        return NULL;
    }
    CloseHandle(hExistingToken);

    LPVOID pEnv = NULL;
    DWORD dwCreationFlags = NORMAL_PRIORITY_CLASS | CREATE_NEW_CONSOLE;
    if (CreateEnvironmentBlock(&pEnv,hNewToken,FALSE))
    {
        dwCreationFlags |= CREATE_UNICODE_ENVIRONMENT;
    }

    STARTUPINFO si = {0};
    PROCESS_INFORMATION pi = {0};
    si.cb = sizeof(si);
    //si.lpDesktop = "WinSta0\\Default";

    char szSubProcCurDir[MAX_PATH+1];
    strcpy(szSubProcCurDir,lpPath);
    *strrchr(szSubProcCurDir,'\\') = '\0';

    if (!CreateProcessAsUser(hNewToken,lpPath,lpCmdLine,NULL,NULL,FALSE,dwCreationFlags,pEnv,szSubProcCurDir,&si,&pi))
    {
        ofs_log<<"CreateProcessAsUser失败，错误码为："<<GetLastError()<<std::endl;
        if (pEnv)
        {
            DestroyEnvironmentBlock(pEnv);
        }
        return NULL;
    }

    if (pi.hThread && pi.hThread != INVALID_HANDLE_VALUE)
    {
        CloseHandle(pi.hThread);
    }

    if (pEnv)
    {
        DestroyEnvironmentBlock(pEnv);
    }

    return pi.hProcess;
}

void WorkFunc()
{
    char szIniPath[MAX_PATH+1];
    sprintf(szIniPath,"%s\\config.ini",szCurDir);
    if (!PathFileExists(szCurDir))
    {
        ofs_log<<"配置文件不存在"<<std::endl;
    }

    char szProgPath[MAX_PATH+1];
    GetPrivateProfileString("PROGRAM","PATH","",szProgPath,MAX_PATH,szIniPath);
    char szCmdLine[200];
    GetPrivateProfileString("PROGRAM","CMD","",szCmdLine,200,szIniPath);
    int method = GetPrivateProfileInt("USERINFO","METHOD",1,szIniPath);
    char szUserName[100];
    char szPassword[100];
    GetPrivateProfileString("USERINFO","UID","",szUserName,99,szIniPath);
    GetPrivateProfileString("USERINFO","PWD","",szPassword,99,szIniPath);

    while (bRun)
    {
        switch (method)
        {
        case 1:
            hProcess = RunAsCreator(szProgPath,szCmdLine);
            break;
        case 2:
            hProcess = RunAsLoggedUser(szProgPath,szCmdLine);
            break;
        case 3:
            hProcess = RunAsSpecifiedUser(szUserName,szPassword,szProgPath,szCmdLine);
            break;
        default:
            ofs_log<<"未知的启动方式："<<method<<std::endl;
            return;
        }

        if (!hProcess)
        {
            ofs_log<<"创建进程失败失败"<<std::endl;
        }
        WaitForSingleObject(hProcess,INFINITE);
        CloseHandle(hProcess);
        hProcess = NULL;
        ofs_log<<"被守护进程退出"<<std::endl;
        Sleep(2000);
    }
}