Service.h

/*
 * Copyright (c) 2025 Red Hat, Inc.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met :
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and / or other materials provided with the distribution.
 * 3. Neither the names of the copyright holders nor the names of their contributors
 *    may be used to endorse or promote products derived from this software
 *    without specific prior written permission.
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

#pragma once

template <typename From, typename To> void ConvertString(const From &from, To &to)
{
    to.Empty();
    for (INT i = 0; i < from.GetLength(); ++i)
    {
        to.AppendChar((char)from[i]);
    }
}

class CServiceManager
{
  public:
    CServiceManager(bool bQueryOnly = false)
    {
        m_Handle = OpenSCManager(NULL, NULL, bQueryOnly ? SC_MANAGER_ENUMERATE_SERVICE : SC_MANAGER_ALL_ACCESS);
        if (!m_Handle)
        {
            Log("%s - failed, error %d", __FUNCTION__, GetLastError());
        }
    }
    ~CServiceManager()
    {
        if (m_Handle)
        {
            CloseServiceHandle(m_Handle);
        }
    }
    bool Install(LPCTSTR name, LPCTSTR binary)
    {
        CString s;
        s.Format(_T("\"%s\" %s"), binary, name);
        SC_HANDLE h = CreateService(m_Handle,
                                    name,
                                    name,
                                    SERVICE_STOP | SERVICE_START | SERVICE_QUERY_STATUS,
                                    ServiceType(),
                                    SERVICE_AUTO_START,
                                    SERVICE_ERROR_NORMAL,
                                    s,
                                    NULL,
                                    0,
                                    NULL,
                                    NULL,
                                    NULL);
        if (h)
        {
            Log("%s: service installed", __FUNCTION__);
            CloseServiceHandle(h);
        }
        else
        {
            Log("%s: error %d", __FUNCTION__, GetLastError());
        }
        return h != NULL;
    }
    SC_HANDLE Handle() const
    {
        return m_Handle;
    }
    static ULONG ServiceType()
    {
        return SERVICE_WIN32_OWN_PROCESS;
    }

  protected:
    SC_HANDLE m_Handle;
};

class CService
{
  public:
    CService(LPCWSTR name, bool bQueryOnly = false) : m_Handle(NULL), m_ServiceManager(bQueryOnly)
    {
        ULONG flags = SERVICE_QUERY_STATUS;
        if (!bQueryOnly)
        {
            flags |= SERVICE_STOP | SERVICE_START | DELETE | SERVICE_USER_DEFINED_CONTROL;
        }
        if (m_ServiceManager.Handle())
        {
            m_Handle = OpenService(m_ServiceManager.Handle(), name, flags);
        }
        if (!m_Handle)
        {
            Log("%s %S - failed, error %d", __FUNCTION__, name, GetLastError());
        }
    }
    ~CService()
    {
        if (m_Handle)
        {
            CloseServiceHandle(m_Handle);
        }
    }
    ULONG Stop()
    {
        if (!m_Handle)
        {
            return ERROR_ACCESS_DENIED;
        }
        if (!ControlService(m_Handle, SERVICE_CONTROL_STOP, &m_Status))
        {
            ULONG error = GetLastError();
            Log("%s %u", __FUNCTION__, error);
            return error;
        }
        Log("%s done", __FUNCTION__);
        return ERROR_SUCCESS;
    }
    ULONG Start()
    {
        if (!m_Handle)
        {
            return ERROR_ACCESS_DENIED;
        }
        if (!StartService(m_Handle, 0, NULL))
        {
            ULONG error = GetLastError();
            Log("%s %u", __FUNCTION__, error);
            return error;
        }
        Log("%s done", __FUNCTION__);
        return ERROR_SUCCESS;
    }
    ULONG Query(ULONG &state)
    {
        if (!m_Handle)
        {
            return ERROR_ACCESS_DENIED;
        }
        if (!QueryServiceStatus(m_Handle, &m_Status))
        {
            return GetLastError();
        }
        state = m_Status.dwCurrentState;
        Log("%s = %s(0x%X)", __FUNCTION__, StateName(state), state);
        return ERROR_SUCCESS;
    }
    bool Delete()
    {
        bool b;
        if (!m_Handle)
        {
            return false;
        }
        b = !!DeleteService(m_Handle);
        if (!b)
        {
            ULONG error = GetLastError();
            Log("%s %u", __FUNCTION__, error);
        }
        else
        {
            Log("%s OK", __FUNCTION__);
        }
        return b;
    }
    bool NeedActionToStart()
    {
        ULONG state;
        if (Query(state))
        {
            return false;
        }
        return state == SERVICE_STOPPED || state == SERVICE_STOP_PENDING;
    }
    bool NeedActionToStop()
    {
        ULONG state;
        if (Query(state))
        {
            return false;
        }
        return state == SERVICE_RUNNING;
    }
    static LPCSTR StateName(ULONG state)
    {
        LPCSTR names[] = {"Unknown",
                          "Stopped",
                          "Start Pending",
                          "Stop Pending",
                          "Running",
                          "Continue Pending",
                          "Pause Pending",
                          "Paused"};
        if (state >= ELEMENTS_IN(names))
        {
            state = 0;
        }
        return names[state];
    }
    ULONG NotifyStatusChange(ULONG dwNotifyMask, PSERVICE_NOTIFY pNotify)
    {
        if (m_Handle)
        {
            ULONG state;
            if (!Query(state) && (m_Status.dwServiceType == SERVICE_KERNEL_DRIVER ||
                                  m_Status.dwServiceType == SERVICE_FILE_SYSTEM_DRIVER))
            {
                Log("%s: Not possible for kernel drivers", __FUNCTION__);
                return ERROR_INVALID_FLAGS;
            }
            return NotifyServiceStatusChange(m_Handle, dwNotifyMask, pNotify);
        }
        else
        {
            return ERROR_INVALID_HANDLE;
        }
    }
    void Control(UCHAR code)
    {
        if (m_Handle)
        {
            SERVICE_STATUS status;
            if (!ControlService(m_Handle, code, &status))
            {
                Log("%s: error %d", __FUNCTION__, GetLastError());
            }
        }
    }

  protected:
    CServiceManager m_ServiceManager;
    SC_HANDLE m_Handle;
    SERVICE_STATUS m_Status;
};

class CServiceRegKey : public CRegKey
{
  public:
    CServiceRegKey(CString Name)
    {
        CString s = TEXT("system\\currentcontrolset\\services\\");
        s += Name;
        Open(HKEY_LOCAL_MACHINE, s);
    }
};

class CServiceImplementation
{
  public:
    CServiceImplementation(LPCTSTR name) : m_Name(name)
    {
        Register(this);
    }
    LPCTSTR ServiceName() const
    {
        return m_Name;
    }
    bool IsEqual(LPCSTR AString, const CString &WString)
    {
        CString s;
        CStringA a = AString;
        ConvertString<CStringA, CString>(a, s);
        return !WString.CompareNoCase(s);
    }
    bool CheckInMain(LPCSTR FirstParameter)
    {
        char *data;
        ULONG done;
        if (WTSQuerySessionInformationA(WTS_CURRENT_SERVER_HANDLE, WTS_CURRENT_SESSION, WTSSessionId, &data, &done) &&
            done == 4)
        {
            ULONG session = *(ULONG *)data;
            Log("The application runs in session %d", session);
            WTSFreeMemory(data);
            INT_PTR count = m_Registered.GetCount();
            if (session == 0 && count && IsEqual(FirstParameter, m_Name))
            {
                SERVICE_TABLE_ENTRY *table = new SERVICE_TABLE_ENTRY[count + 1];
                table[count].lpServiceName = NULL;
                table[count].lpServiceProc = NULL;
                for (INT_PTR i = 0; i < count; ++i)
                {
                    CServiceImplementation *obj = m_Registered.GetAt(i);
                    table[i].lpServiceName = obj->m_Name.GetBuffer();
                    table[i].lpServiceProc = [](DWORD argc, LPTSTR *argv) {
                        CServiceImplementation *obj = NULL;
                        if (!argc || CServiceImplementation::Registered().GetCount() == 1)
                        {
                            obj = CServiceImplementation::Registered().GetAt(0);
                        }
                        else
                        {
                            obj = FindObject(argv[0]);
                        }
                        if (obj)
                        {
                            obj->ServiceMain(argc, argv);
                        }
                        else
                        {
                            Log("Nothing found to run: %d args", argc);
                        }
                    };
                }
                StartServiceCtrlDispatcher(table);
                delete[] table;
                return true;
            }
        }
        Log("This is not a service");
        return false;
    }
    bool Install(LPCWSTR BinPath)
    {
        CString s = BinPath ? BinPath : BinaryPath();
        CServiceManager mgr;
        bool b = mgr.Install(ServiceName(), s);
        return b;
    }
    bool Installed()
    {
        CService service(ServiceName(), true);
        ULONG state;
        return !service.Query(state);
    }
    bool Uninstall()
    {
        CService service(ServiceName());
        return service.Delete();
    }
    static const CAtlArray<CServiceImplementation *> &Registered()
    {
        return m_Registered;
    }
    static CString BinaryPath()
    {
        CString s;
        HMODULE hModule = GetModuleHandle(NULL);
        WCHAR path[MAX_PATH];
        GetModuleFileName(hModule, path, MAX_PATH);
        s = path;
        return s;
    }
    void Control(UCHAR code)
    {
        CService service(ServiceName());
        service.Control(code);
    }
    bool Start()
    {
        CService service(ServiceName());
        return service.Start() == ERROR_SUCCESS;
    }
    bool IsRunning()
    {
        CService service(ServiceName());
        ULONG state;
        return service.Query(state) == ERROR_SUCCESS && state == SERVICE_RUNNING;
    }
    void Stop()
    {
        CService service(ServiceName());
        service.Stop();
    }
    bool GetValue(LPCTSTR Name, ULONG &Value)
    {
        CServiceRegKey key(m_Name);
        return key.QueryDWORDValue(Name, Value) == ERROR_SUCCESS;
    }
    void SetValue(LPCTSTR Name, ULONG Value)
    {
        CServiceRegKey key(m_Name);
        key.SetDWORDValue(Name, Value);
    }
    bool GetString(LPCTSTR Name, CStringA &Value)
    {
        CServiceRegKey key(m_Name);
        TCHAR buffer[1024];
        ULONG len = ARRAYSIZE(buffer);
        if (key.QueryStringValue(Name, buffer, &len) == ERROR_SUCCESS)
        {
            CString s;
            buffer[len] = 0;
            s = buffer;
            ConvertString<CString, CStringA>(s, Value);
            return true;
        }
        return false;
    }
    void SetString(LPCTSTR Name, LPCSTR Value)
    {
        CServiceRegKey key(m_Name);
        CStringA s = Value;
        CString ws;
        ConvertString<CStringA, CString>(s, ws);
        key.SetStringValue(Name, ws);
    }

  protected:
    CString m_Name;
    void SetState(ULONG State)
    {
        m_State.Set(State);
    }
    class CServiceState
    {
      public:
        CServiceState() : m_StopEvent(true)
        {
            status.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
            status.dwCurrentState = SERVICE_STOPPED;
            status.dwControlsAccepted = SERVICE_ACCEPT_STOP | SERVICE_ACCEPT_SESSIONCHANGE | SERVICE_ACCEPT_PRESHUTDOWN;
            status.dwWaitHint = 30000;
            status.dwCheckPoint = 0;
            status.dwWin32ExitCode = NO_ERROR;
            status.dwServiceSpecificExitCode = 0;
        }
        SERVICE_STATUS_HANDLE hService = NULL;
        SERVICE_STATUS status;
        void Set(ULONG state)
        {
            status.dwCurrentState = state;
            Log("%s: %s", __FUNCTION__, CService::StateName(state));
            SetServiceStatus(hService, &status);
            if (state == SERVICE_STOPPED)
            {
                m_StopEvent.Set();
            }
        }
        void Stoppable(bool b)
        {
            if (b)
            {
                status.dwControlsAccepted |= SERVICE_ACCEPT_STOP;
            }
            else
            {
                status.dwControlsAccepted &= ~SERVICE_ACCEPT_STOP;
            }
            Log("%s = %d", __FUNCTION__, b);
            SetServiceStatus(hService, &status);
        }
        bool Stoppable()
        {
            bool b = (status.dwControlsAccepted & SERVICE_ACCEPT_STOP) != 0;
            Log("%s = %d", __FUNCTION__, b);
            return b;
        }
        CEvent m_StopEvent;
    };
    CServiceState m_State;
    void ServiceMain(_In_ ULONG dwArgc, _In_ LPTSTR *lpszArgv)
    {
        UNREFERENCED_PARAMETER(dwArgc);
        UNREFERENCED_PARAMETER(lpszArgv);
        m_State.hService = RegisterServiceCtrlHandlerEx(
                                                                                                            m_Name,
                                                                                                            [](DWORD dwControl,
                                                                                                               DWORD dwEventType,
                                                                                                               LPVOID lpEventData,
                                                                                                               LPVOID lpContext) -> DWORD {
                                                                                                                UNREFERENCED_PARAMETER(lpEventData);
                                                                                                                UNREFERENCED_PARAMETER(dwEventType);
                                                                                                                CServiceImplementation *obj = (CServiceImplementation *)lpContext;
                                                                                                                return obj->ControlHandler(dwControl,
                                                                                                                                           dwEventType,
                                                                                                                                           lpEventData);
                                                                                                            },
                                                                                                            this);
        if (m_State.hService)
        {
            m_State.Set(SERVICE_START_PENDING);
            if (OnStart())
            {
                m_State.Set(SERVICE_RUNNING);
            }
            m_State.m_StopEvent.Wait();
        }
        else
        {
            Log("Failed to register %S service, error %d", ServiceName(), GetLastError());
        }
    }
    virtual DWORD ControlHandler(ULONG dwControl, ULONG dwEventType, LPVOID lpEventData)
    {
        DWORD res = ERROR_CALL_NOT_IMPLEMENTED;
        UNREFERENCED_PARAMETER(dwEventType);
        UNREFERENCED_PARAMETER(lpEventData);
        Log("%s: %s", __FUNCTION__, Name<eServiceControl>(dwControl));
        switch (dwControl)
        {
            case SERVICE_CONTROL_PRESHUTDOWN:
            case SERVICE_CONTROL_SHUTDOWN:
            case SERVICE_CONTROL_STOP:
                if (!m_State.Stoppable())
                {
                    res = ERROR_CALL_NOT_IMPLEMENTED;
                }
                else if (m_State.status.dwCurrentState == SERVICE_RUNNING)
                {
                    res = NO_ERROR;
                    m_State.Set(SERVICE_STOP_PENDING);
                    if (OnStop())
                    {
                        m_State.Set(SERVICE_STOPPED);
                    }
                }
                else
                {
                    Log("%s: got STOP, current state %s(%d)",
                        __FUNCTION__,
                        CService::StateName(m_State.status.dwCurrentState),
                        m_State.status.dwCurrentState);
                }
                break;
            default:
                break;
        }
        return res;
    }
    virtual bool OnStop()
    {
        // if more processing required, return false
        // then later use m_State.Set(SERVICE_STOPPED)
        return true;
    }
    virtual bool OnStart()
    {
        // if more processing required, return false
        // then later use m_State.Set(SERVICE_RUNNING)
        // or m_State.Set(SERVICE_STOPPED)
        return true;
    }
    static void Register(CServiceImplementation *obj)
    {
        m_Registered.Add(obj);
    }
    static CServiceImplementation *FindObject(LPCTSTR name)
    {
        for (size_t i = 0; i < m_Registered.GetCount(); ++i)
        {
            CServiceImplementation *obj = m_Registered.GetAt(i);
            CString s = name;
            if (!s.CompareNoCase(obj->ServiceName()))
            {
                return obj;
            }
        }
        return NULL;
    }
    static CAtlArray<CServiceImplementation *> m_Registered;
};

class CThreadOwner
{
  public:
    CThreadOwner()
    {
    }
    ~CThreadOwner()
    {
        StopThread(true);
        while (IsThreadRunning())
        {
            Sleep(10);
        }
    }
    bool IsThreadRunning() const
    {
        return m_ThreadHandle != NULL;
    }
    enum tThreadState
    {
        tsNotRunning,
        tsRunning,
        tsAborted,
        tsDeleting
    };
    bool StartThread()
    {
        if (IsThreadRunning())
        {
            return false;
        }
        InterlockedCompareExchange(&m_State, tsRunning, tsNotRunning);
        m_ThreadHandle = (HANDLE)_beginthread(_ThreadProc, 0, this);
        if (IsThreadRunning())
        {
            return true;
        }
        else
        {
            InterlockedCompareExchange(&m_State, tsNotRunning, tsRunning);
            return false;
        }
    }
    virtual tThreadState StopThread(bool bDeleting = false)
    {
        LONG val;
        if (bDeleting)
        {
            val = InterlockedCompareExchange(&m_State, tsDeleting, tsRunning);
        }
        else
        {
            val = InterlockedCompareExchange(&m_State, tsAborted, tsRunning);
        }
        return (tThreadState)val;
    }

  private:
    LONG m_State = tsNotRunning;

  protected:
    tThreadState ThreadState()
    {
        return (tThreadState)m_State;
    }
    HANDLE m_ThreadHandle = NULL;
    virtual void ThreadProc() = 0;
    virtual void ThreadTerminated(tThreadState previous)
    {
        UNREFERENCED_PARAMETER(previous);
        m_ThreadHandle = NULL;
    }
    static void __cdecl _ThreadProc(PVOID param)
    {
        CThreadOwner *pOwner = (CThreadOwner *)param;
        pOwner->ThreadProc();
        LONG val = InterlockedExchange(&pOwner->m_State, tsNotRunning);
        pOwner->ThreadTerminated((tThreadState)val);
    }
};