yttrium-virtio-gpu/Balloon/app/utils.cpp

#include "stdafx.h"

extern LPWSTR ServiceName;
extern LPWSTR DisplayName;

extern CService srvc;

void ErrorHandler(char *s, int err)
{
    printf("Failed. Error %d ", err);

    LPTSTR lpMsgBuf;
    if (FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
                      NULL,
                      err,
                      MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
                      (LPTSTR)&lpMsgBuf,
                      0,
                      NULL) > 0)
    {
        _tprintf(L"%s\n", lpMsgBuf);
        LocalFree(lpMsgBuf);
    }
    else
    {
        printf("unknown error\n");
    }

    FILE *pLog;
    if (fopen_s(&pLog, "balloon.log", "a") == 0 && pLog)
    {
        fprintf(pLog, "%s failed, error code = %d\n", s, err);
        fclose(pLog);
    }

    ExitProcess(err);
}

void PrintMessage(char *s)
{
#ifdef DBG
    FILE *pLog;
    if (fopen_s(&pLog, "balloon.log", "a") == 0)
    {
        fprintf(pLog, "%s\n", s);
        fclose(pLog);
    }
#endif
}

void ShowUsage()
{
    printf("\n");
    printf("USAGE:\n");
    printf("blnsvr -i\tInstall service\n");
    printf("blnsvr -u\tUninstall service\n");
    printf("blnsvr -r\tRun service\n");
    printf("blnsvr -s\tStop service\n");
    printf("blnsvr status\tCurrent status\n");
    printf("\n");
}

BOOL InstallService()
{
    SC_HANDLE newService;
    SC_HANDLE scm;
    TCHAR szBuffer[255];
    TCHAR szPath[MAX_PATH];

    GetModuleFileName(GetModuleHandle(NULL), szPath, MAX_PATH);
    if (FAILED(StringCchCopy(szBuffer, 255, TEXT("\""))))
    {
        return FALSE;
    }
    if (FAILED(StringCchCat(szBuffer, 255, szPath)))
    {
        return FALSE;
    }
    if (FAILED(StringCchCat(szBuffer, 255, TEXT("\""))))
    {
        return FALSE;
    }

    scm = OpenSCManager(NULL, NULL, SC_MANAGER_CREATE_SERVICE);
    if (scm == NULL)
    {
        ErrorHandler("OpenSCManager", GetLastError());
    }
    newService = CreateService(scm,
                               ServiceName,
                               DisplayName,
                               SERVICE_ALL_ACCESS,
                               SERVICE_WIN32_OWN_PROCESS,
                               SERVICE_AUTO_START,
                               SERVICE_ERROR_NORMAL,
                               szBuffer,
                               NULL,
                               NULL,
                               NULL,
                               NULL,
                               NULL);
    if (!newService)
    {
        ErrorHandler("CreateService", GetLastError());
        return FALSE;
    }
    else
    {
        printf("Service Installed\n");
        ServiceRun();
    }

    CloseServiceHandle(newService);
    CloseServiceHandle(scm);

    return TRUE;
}

BOOL UninstallService()
{
    SC_HANDLE service;
    SC_HANDLE scm;
    BOOL res;
    SERVICE_STATUS status;

    scm = OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (!scm)
    {
        ErrorHandler("OpenSCManager", GetLastError());
    }

    service = OpenService(scm, ServiceName, SERVICE_ALL_ACCESS | DELETE);
    if (!service)
    {
        ErrorHandler("OpenService", GetLastError());
    }

    res = QueryServiceStatus(service, &status);
    if (!res)
    {
        ErrorHandler("QueryServiceStatus", GetLastError());
    }

    if (status.dwCurrentState != SERVICE_STOPPED)
    {
        printf("Stopping service...\n");
        res = ControlService(service, SERVICE_CONTROL_STOP, &status);
        if (!res)
        {
            ErrorHandler("ControlService", GetLastError());
        }
        Sleep(5000);
    }

    res = DeleteService(service);
    if (res)
    {
        printf("Service Uninstalled\n");
    }
    else
    {
        ErrorHandler("DeleteService", GetLastError());
    }

    CloseServiceHandle(service);
    CloseServiceHandle(scm);

    return TRUE;
}

BOOL ServiceRun()
{
    SC_HANDLE scm, Service;
    SERVICE_STATUS ssStatus;
    DWORD dwOldCheckPoint;
    DWORD dwStartTickCount;
    DWORD dwWaitTime;
    DWORD dwStatus;

    scm = OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (!scm)
    {
        ErrorHandler("OpenSCManager", GetLastError());
    }

    Service = OpenService(scm, ServiceName, SERVICE_ALL_ACCESS);
    if (!Service)
    {
        ErrorHandler("OpenService", GetLastError());
        return FALSE;
    }
    else
    {
        StartService(Service, 0, NULL);
        srvc.GetStatus(Service);

        if (!QueryServiceStatus(Service, &ssStatus))
        {
            ErrorHandler("QueryServiceStatus", GetLastError());
        }
        dwStartTickCount = GetTickCount();
        dwOldCheckPoint = ssStatus.dwCheckPoint;

        while (ssStatus.dwCurrentState == SERVICE_START_PENDING)
        {
            dwWaitTime = ssStatus.dwWaitHint / 10;

            if (dwWaitTime < 1000)
            {
                dwWaitTime = 1000;
            }
            else if (dwWaitTime > 10000)
            {
                dwWaitTime = 10000;
            }

            Sleep(dwWaitTime);

            if (!QueryServiceStatus(Service, &ssStatus))
            {
                break;
            }

            if (ssStatus.dwCheckPoint > dwOldCheckPoint)
            {
                dwStartTickCount = GetTickCount();
                dwOldCheckPoint = ssStatus.dwCheckPoint;
            }
            else
            {
                if (GetTickCount() - dwStartTickCount > ssStatus.dwWaitHint)
                {
                    break;
                }
            }
        }

        if (ssStatus.dwCurrentState == SERVICE_RUNNING)
        {
            srvc.GetStatus(Service);
            dwStatus = NO_ERROR;
        }
        else
        {

            printf("\nService not started.\n");
            printf("  Current State: %d\n", ssStatus.dwCurrentState);
            printf("  Exit Code: %d\n", ssStatus.dwWin32ExitCode);
            printf("  Service Specific Exit Code: %d\n", ssStatus.dwServiceSpecificExitCode);
            printf("  Check Point: %d\n", ssStatus.dwCheckPoint);
            printf("  Wait Hint: %d\n", ssStatus.dwWaitHint);
            dwStatus = GetLastError();
        }
    }

    CloseServiceHandle(scm);
    CloseServiceHandle(Service);
    return TRUE;
}

BOOL ServiceControl(int ctrl)
{
    SC_HANDLE service;
    SC_HANDLE scm;
    BOOL res = TRUE;
    SERVICE_STATUS status;

    scm = OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (!scm)
    {
        ErrorHandler("OpenSCManager", GetLastError());
    }

    service = OpenService(scm, ServiceName, SERVICE_ALL_ACCESS);
    if (!service)
    {
        ErrorHandler("OpenService", GetLastError());
    }

    if (ctrl == SERVICE_CONTROL_STOP)
    {
        printf("Service is stopping...\n");
        res = ControlService(service, SERVICE_CONTROL_STOP, &status);
    }

    if (!res)
    {
        ErrorHandler("ControlService", GetLastError());
    }
    else
    {
        srvc.GetStatus(service);
    }

    CloseServiceHandle(service);
    CloseServiceHandle(scm);

    return TRUE;
}

BOOL GetConfiguration()
{
    SC_HANDLE service;
    SC_HANDLE scm;
    BOOL res;
    LPQUERY_SERVICE_CONFIG buffer;
    DWORD sizeNeeded;

    scm = OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (!scm)
    {
        ErrorHandler("OpenSCManager", GetLastError());
    }

    service = OpenService(scm, ServiceName, SERVICE_QUERY_CONFIG);
    if (!service)
    {
        ErrorHandler("OpenService", GetLastError());
    }

    buffer = (LPQUERY_SERVICE_CONFIG)LocalAlloc(LPTR, 4096);
    if (!buffer)
    {
        ErrorHandler("LocalAlloc", GetLastError());
    }
    res = QueryServiceConfig(service, buffer, 4096, &sizeNeeded);
    if (!res)
    {
        ErrorHandler("QueryServiceConfig", GetLastError());
    }

    printf("Service name:\t%S\n", buffer->lpDisplayName);
    printf("Service type:\t%d\n", buffer->dwServiceType);
    printf("Start type:\t%d\n", buffer->dwStartType);
    printf("Start name:\t%S\n", buffer->lpServiceStartName);
    printf("Path:\t\t%S\n", buffer->lpBinaryPathName);

    LocalFree(buffer);

    CloseServiceHandle(service);
    CloseServiceHandle(scm);
    return TRUE;
}

BOOL ChangeConfig()
{
    SC_HANDLE service;
    SC_HANDLE scm;
    BOOL res;
    SC_LOCK lock;

    scm = OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS | GENERIC_WRITE);
    if (!scm)
    {
        ErrorHandler("OpenSCManager", GetLastError());
    }
    lock = LockServiceDatabase(scm);
    if (lock == 0)
    {
        ErrorHandler("LockServiceDatabase", GetLastError());
    }
    service = OpenService(scm, ServiceName, SERVICE_ALL_ACCESS);
    if (!service)
    {
        ErrorHandler("OpenService", GetLastError());
    }
    res = ChangeServiceConfig(service,
                              SERVICE_NO_CHANGE,
                              SERVICE_NO_CHANGE,
                              SERVICE_NO_CHANGE,
                              NULL,
                              NULL,
                              NULL,
                              NULL,
                              NULL,
                              NULL,
                              NULL);
    if (!res)
    {
        UnlockServiceDatabase(lock);
        ErrorHandler("ChangeServiceConfig", GetLastError());
    }

    res = UnlockServiceDatabase(lock);
    if (!res)
    {
        ErrorHandler("UnlockServiceDatabase", GetLastError());
    }
    CloseServiceHandle(service);
    CloseServiceHandle(scm);
    return TRUE;
}