A service to run multiple instances of SETI on one machine

/*
*	$Header: $
*
*	$History: $
*/
// Note: Proxy/Stub Information
//      To build a separate proxy/stub DLL, 
//      run nmake -f setisrvps.mk in the project directory.
#include "stdafx.h"
#include <io.h>
#include <direct.h>
#include <process.h>
#include "resource.h"
#include <initguid.h>
#include "setisrv.h"

#include "setisrv_i.c"

#include <stdio.h>

void RunSETI(LPVOID);

CServiceModule _Module;

struct runInfo
{
	CString		   m_csWorkingFolder,
				   m_csExeName;
	HANDLE		   m_hArray[2];
	DWORD		   dwProcessor;
	CServiceModule *me;
};

enum arrayValues
{
	STOPEVENT,
	PROCESSHANDLE
};

BEGIN_OBJECT_MAP(ObjectMap)
END_OBJECT_MAP()

LPCTSTR FindOneOf(LPCTSTR p1, LPCTSTR p2)
{
    while (p1 != NULL && *p1 != NULL)
    {
        LPCTSTR p = p2;

        while (p != NULL && *p != NULL)
        {
            if (*p1 == *p)
                return CharNext(p1);
        
			p = CharNext(p);
        }

        p1 = CharNext(p1);
    }

    return NULL;
}

// Although some of these functions are big they are declared inline since they are only used once
inline HRESULT CServiceModule::RegisterServer(BOOL bRegTypeLib, BOOL bService)
{
    HRESULT hr = CoInitialize(NULL);

    if (FAILED(hr))
        return hr;

    // Remove any previous service since it may point to
    // the incorrect file
    Uninstall();

    // Add service entries
    UpdateRegistryFromResource(IDR_Setisrv, TRUE);

    // Adjust the AppID for Local Server or Service
    CRegKey keyAppID;
    
	LONG lRes = keyAppID.Open(HKEY_CLASSES_ROOT, _T("AppID"), KEY_WRITE);
    
	if (lRes != ERROR_SUCCESS)
        return lRes;

    CRegKey key;
    
	lRes = key.Open(keyAppID, _T("{F25915C1-A6EB-4BE6-BA9C-8F4054AF54C5}"), KEY_WRITE);
    
	if (lRes != ERROR_SUCCESS)
        return lRes;
    
	key.DeleteValue(_T("LocalService"));
    
    if (bService)
    {
        key.SetValue(_T("setisrv"), _T("LocalService"));
        key.SetValue(_T("-Service"), _T("ServiceParameters"));
        // Create service
        Install();
    }

    // Add object entries
    hr = CComModule::RegisterServer(bRegTypeLib);

    CoUninitialize();
    return hr;
}

inline HRESULT CServiceModule::UnregisterServer()
{
    HRESULT hr = CoInitialize(NULL);
    
	if (FAILED(hr))
        return hr;

    // Remove service entries
    UpdateRegistryFromResource(IDR_Setisrv, FALSE);
    
	// Remove service
    Uninstall();
    
	// Remove object entries
    CComModule::UnregisterServer(TRUE);
    CoUninitialize();
    return S_OK;
}

inline void CServiceModule::Init(_ATL_OBJMAP_ENTRY* p, HINSTANCE h, UINT nServiceNameID, const GUID* plibid)
{
    CComModule::Init(p, h, plibid);

    m_bService = TRUE;
    LoadString(h, nServiceNameID, m_szServiceName, sizeof(m_szServiceName) / sizeof(TCHAR));

    // set up the initial service status 
    m_hServiceStatus = NULL;
    m_status.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
    m_status.dwCurrentState = SERVICE_STOPPED;
    m_status.dwControlsAccepted = SERVICE_ACCEPT_STOP;
    m_status.dwWin32ExitCode = 0;
    m_status.dwServiceSpecificExitCode = 0;
    m_status.dwCheckPoint = 0;
    m_status.dwWaitHint = 0;
}

LONG CServiceModule::Unlock()
{
    LONG l = CComModule::Unlock();

    if (l == 0 && !m_bService)
        PostThreadMessage(dwThreadID, WM_QUIT, 0, 0);
    
	return l;
}

BOOL CServiceModule::IsInstalled()
{
    BOOL bResult = FALSE;

    SC_HANDLE hSCM = ::OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);

    if (hSCM != NULL)
    {
        SC_HANDLE hService = ::OpenService(hSCM, m_szServiceName, SERVICE_QUERY_CONFIG);
    
		if (hService != NULL)
        {
            bResult = TRUE;
            ::CloseServiceHandle(hService);
        }
        
		::CloseServiceHandle(hSCM);
    }

    return bResult;
}

inline BOOL CServiceModule::Install()
{
    if (IsInstalled())
        return TRUE;

    SC_HANDLE hSCM = ::OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);

    if (hSCM == NULL)
    {
        MessageBox(NULL, _T("Couldn't open service manager"), m_szServiceName, MB_OK);
        return FALSE;
    }

    // Get the executable file path
    TCHAR szFilePath[_MAX_PATH];
    
	::GetModuleFileName(NULL, szFilePath, _MAX_PATH);

    SC_HANDLE hService = ::CreateService(hSCM, m_szServiceName, m_szServiceName, SERVICE_ALL_ACCESS, SERVICE_WIN32_OWN_PROCESS, SERVICE_AUTO_START, SERVICE_ERROR_NORMAL, szFilePath, NULL, NULL, _T("RPCSS\0"), NULL, NULL);

    if (hService == NULL)
    {
        ::CloseServiceHandle(hSCM);
        MessageBox(NULL, _T("Couldn't create service"), m_szServiceName, MB_OK);
        return FALSE;
    }

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

inline BOOL CServiceModule::Uninstall()
{
    if (!IsInstalled())
        return TRUE;

    SC_HANDLE hSCM = ::OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);

    if (hSCM == NULL)
    {
        MessageBox(NULL, _T("Couldn't open service manager"), m_szServiceName, MB_OK);
        return FALSE;
    }

    SC_HANDLE hService = ::OpenService(hSCM, m_szServiceName, SERVICE_STOP | DELETE);

    if (hService == NULL)
    {
        ::CloseServiceHandle(hSCM);
        MessageBox(NULL, _T("Couldn't open service"), m_szServiceName, MB_OK);
        return FALSE;
    }

    SERVICE_STATUS status;
    
	::ControlService(hService, SERVICE_CONTROL_STOP, &status);

    BOOL bDelete = ::DeleteService(hService);
    
	::CloseServiceHandle(hService);
    ::CloseServiceHandle(hSCM);

    if (bDelete)
        return TRUE;

    MessageBox(NULL, _T("Service could not be deleted"), m_szServiceName, MB_OK);
    return FALSE;
}

///////////////////////////////////////////////////////////////////////////////////////
// Logging functions
void CServiceModule::LogEvent(LPCTSTR pFormat, ...)
{
    TCHAR    chMsg[256];
    HANDLE  hEventSource;
    LPTSTR  lpszStrings[1];
    va_list pArg;

    va_start(pArg, pFormat);
    _vstprintf(chMsg, pFormat, pArg);
    va_end(pArg);

    lpszStrings[0] = chMsg;

    if (m_bService)
    {
        /* Get a handle to use with ReportEvent(). */
        hEventSource = RegisterEventSource(NULL, m_szServiceName);
    
		if (hEventSource != NULL)
        {
            /* Write to event log. */
            ReportEvent(hEventSource, EVENTLOG_INFORMATION_TYPE, 0, 0, NULL, 1, 0, (LPCTSTR*) &lpszStrings[0], NULL);
            DeregisterEventSource(hEventSource);
        }
    }
    else
    {
        // As we are not running as a service, just write the error to the console.
        _putts(chMsg);
    }
}

//////////////////////////////////////////////////////////////////////////////////////////////
// Service startup and registration
inline void CServiceModule::Start()
{
    SERVICE_TABLE_ENTRY st[] =
    {
        { m_szServiceName, _ServiceMain },
        { NULL, NULL }
    };

    if (m_bService && !::StartServiceCtrlDispatcher(st))
    {
        m_bService = FALSE;
    }
    
	if (m_bService == FALSE)
        Run();
}

inline void CServiceModule::ServiceMain(DWORD /* dwArgc */, LPTSTR* /* lpszArgv */)
{
    // Register the control request handler
    m_status.dwCurrentState = SERVICE_START_PENDING;
    m_hServiceStatus = RegisterServiceCtrlHandler(m_szServiceName, _Handler);
    
	if (m_hServiceStatus == NULL)
    {
        LogEvent(_T("Handler not installed"));
        return;
    }
    
	SetServiceStatus(SERVICE_START_PENDING);

    m_status.dwWin32ExitCode = S_OK;
    m_status.dwCheckPoint = 0;
    m_status.dwWaitHint = 0;

    // When the Run function returns, the service has stopped.
    Run();
    SetServiceStatus(SERVICE_STOPPED);
    LogEvent(_T("Service stopped"));
}

inline void CServiceModule::Handler(DWORD dwOpcode)
{
    switch (dwOpcode)
    {
    case SERVICE_CONTROL_STOP:
        SetServiceStatus(SERVICE_STOP_PENDING);
        PostThreadMessage(dwThreadID, WM_QUIT, 0, 0);
        break;

    case SERVICE_CONTROL_PAUSE:
        break;
    
	case SERVICE_CONTROL_CONTINUE:
        break;
    
	case SERVICE_CONTROL_INTERROGATE:
        break;
    
	case SERVICE_CONTROL_SHUTDOWN:
        break;
    
	default:
        LogEvent(_T("Bad service request"));
    }
}

void WINAPI CServiceModule::_ServiceMain(DWORD dwArgc, LPTSTR* lpszArgv)
{
    _Module.ServiceMain(dwArgc, lpszArgv);
}

void WINAPI CServiceModule::_Handler(DWORD dwOpcode)
{
    _Module.Handler(dwOpcode); 
}

void CServiceModule::SetServiceStatus(DWORD dwState)
{
    m_status.dwCurrentState = dwState;
    ::SetServiceStatus(m_hServiceStatus, &m_status);
}

void CServiceModule::Run()
{
	HANDLE hStop;

    _Module.dwThreadID = GetCurrentThreadId();

    HRESULT hr = CoInitialize(NULL);
//  If you are running on NT 4.0 or higher you can use the following call
//  instead to make the EXE free threaded.
//  This means that calls come in on a random RPC thread
//  HRESULT hr = CoInitializeEx(NULL, COINIT_MULTITHREADED);

    _ASSERTE(SUCCEEDED(hr));

    // This provides a NULL DACL which will allow access to everyone.
    CSecurityDescriptor sd;

    sd.InitializeFromThreadToken();
    hr = CoInitializeSecurity(sd, -1, NULL, NULL, RPC_C_AUTHN_LEVEL_PKT, RPC_C_IMP_LEVEL_IMPERSONATE, NULL, EOAC_NONE, NULL);
    _ASSERTE(SUCCEEDED(hr));

    hr = _Module.RegisterClassObjects(CLSCTX_LOCAL_SERVER | CLSCTX_REMOTE_SERVER, REGCLS_MULTIPLEUSE);
    _ASSERTE(SUCCEEDED(hr));

    LogEvent(_T("Service started"));
	
    if (m_bService)
        SetServiceStatus(SERVICE_RUNNING);

	if ((hStop = CreateEvent(NULL, TRUE, FALSE, NULL)) != HANDLE(NULL))
	{
		TCHAR			tszTemp[_MAX_PATH],
					    tszDrive[_MAX_DRIVE],
						tszPath[_MAX_PATH],
						tszFileName[_MAX_FNAME],
						tszExtension[_MAX_EXT];
		CString			csPath,
						csSetiPath;
		WIN32_FIND_DATA FindFileData;
		HANDLE			hFind;

		//	Get our module name and search for seti subdirectories
		GetModuleFileName(NULL, tszTemp, sizeof(tszTemp));
		_splitpath(tszTemp, tszDrive, tszPath, tszFileName, tszExtension);
		csPath = tszDrive;
		csPath += tszPath;
		csSetiPath = csPath;
		csPath += _T("*.*");

		if ((hFind = FindFirstFile(csPath, &FindFileData)) != INVALID_HANDLE_VALUE)
		{
			do
			{
				if (FindFileData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)
				{
					if (_tcscmp(FindFileData.cFileName, _T(".")) != 0 && _tcscmp(FindFileData.cFileName, _T("..")) != 0)
					{
						//	Found a candidate directory, does it contain seti.exe?
						runInfo *sRun = new runInfo;
						
						sRun->m_csWorkingFolder = csSetiPath;
						sRun->m_csWorkingFolder += FindFileData.cFileName;
						sRun->m_csExeName = sRun->m_csWorkingFolder;
						sRun->m_csExeName += _T("\\seti.exe");
						sRun->me = this;
						sRun->dwProcessor = 0;
						sRun->m_hArray[STOPEVENT] = hStop;
						sRun->m_hArray[PROCESSHANDLE] = HANDLE(NULL);

						if (_access(sRun->m_csExeName, 0) == 0)
						{
							// Now check if there's a process affinity file...
							CString csTemp;

							for (int i = 0; i < 32; i++)
							{
								csTemp.Format(_T("%s\\%d"), sRun->m_csWorkingFolder, i);

								if (_access(csTemp, 0) == 0)
									sRun->dwProcessor |= 0x1 << i;
							}

							_beginthread(RunSETI, 0, LPVOID(sRun));
						}
						else
							delete sRun;
					}
				}
			} while (FindNextFile(hFind, &FindFileData));

			FindClose(hFind);
			LogEvent(_T("Finished directory search loop"));
		}

		MSG msg;
    
		while (GetMessage(&msg, 0, 0, 0))
			DispatchMessage(&msg);

		SetEvent(hStop);
	}
	else
		LogEvent(_T("Unable to create Stop Event, exiting"));

	CloseHandle(hStop);
    _Module.RevokeClassObjects();
    CoUninitialize();
}

/////////////////////////////////////////////////////////////////////////////
//
extern "C" int WINAPI _tWinMain(HINSTANCE hInstance, HINSTANCE /*hPrevInstance*/, LPTSTR lpCmdLine, int /*nShowCmd*/)
{
    lpCmdLine = GetCommandLine(); //this line necessary for _ATL_MIN_CRT
    _Module.Init(ObjectMap, hInstance, IDS_SERVICENAME, &LIBID_SETISRVLib);
    _Module.m_bService = TRUE;

    TCHAR szTokens[] = _T("-/");

    LPCTSTR lpszToken = FindOneOf(lpCmdLine, szTokens);

    while (lpszToken != NULL)
    {
        if (lstrcmpi(lpszToken, _T("UnregServer"))==0)
            return _Module.UnregisterServer();

        // Register as Local Server
        if (lstrcmpi(lpszToken, _T("RegServer"))==0)
            return _Module.RegisterServer(TRUE, FALSE);
        
        // Register as Service
        if (lstrcmpi(lpszToken, _T("Service"))==0)
            return _Module.RegisterServer(TRUE, TRUE);
        
        lpszToken = FindOneOf(lpszToken, szTokens);
    }

    // Are we Service or Local Server
    CRegKey keyAppID;
    LONG	lRes = keyAppID.Open(HKEY_CLASSES_ROOT, _T("AppID"), KEY_READ);
    
	if (lRes != ERROR_SUCCESS)
        return lRes;

    CRegKey key;
    
	lRes = key.Open(keyAppID, _T("{F25915C1-A6EB-4BE6-BA9C-8F4054AF54C5}"), KEY_READ);
    
	if (lRes != ERROR_SUCCESS)
        return lRes;

    TCHAR szValue[_MAX_PATH];
    DWORD dwLen = _MAX_PATH;
    
	lRes = key.QueryValue(szValue, _T("LocalService"), &dwLen);
    _Module.m_bService = FALSE;

    if (lRes == ERROR_SUCCESS)
        _Module.m_bService = TRUE;

    _Module.Start();

    // When we get here, the service has been stopped
    return _Module.m_status.dwWin32ExitCode;
}

void RunSETI(LPVOID data)
{
	{
		runInfo *me = (runInfo *) data;
		BOOL    bStop = FALSE;

		_ASSERTE(me);

		STARTUPINFO			si;
		PROCESS_INFORMATION pi;
		CString				csControlFile;

		//	Create our control file name
		csControlFile.Format(_T("%s\\disable"), me->m_csWorkingFolder);

restart:

		memset(&pi, 0, sizeof(pi));
		memset(&si, 0, sizeof(si));
		si.cb = sizeof(si);
		si.dwFlags = STARTF_USESHOWWINDOW;
		si.wShowWindow = SW_HIDE;

		me->me->LogEvent(_T("About to start %s"), me->m_csExeName);
		
		if (_access(csControlFile, 0) == 0 || CreateProcess(me->m_csExeName, NULL, NULL, NULL, FALSE, CREATE_NO_WINDOW, NULL, me->m_csWorkingFolder, &si, &pi))
		{
			me->m_hArray[PROCESSHANDLE] = pi.hProcess;

			if (me->dwProcessor)
			{
				me->me->LogEvent("Attempting to set %s to processor %lx", me->m_csExeName, me->dwProcessor);
				SetProcessAffinityMask(pi.hProcess, me->dwProcessor);
			}

			while (bStop == FALSE)
			{
				switch (WaitForMultipleObjects(me->m_hArray[PROCESSHANDLE] != HANDLE(NULL) ? 2 : 1, me->m_hArray, FALSE, 60000))
				{
				case WAIT_TIMEOUT:
					//  We timed out, check if we should change status...
					if (me->m_hArray[PROCESSHANDLE] != (HANDLE) NULL)
					{
						//	SETI's running but we have a disable file, so kill SETI
						if (_access(csControlFile, 0) == 0)
						{
							me->me->LogEvent("Disabling %s on request", me->m_csExeName);
							TerminateProcess(me->m_hArray[PROCESSHANDLE], 1);
							me->m_hArray[PROCESSHANDLE] = HANDLE(NULL);
							memset(&pi, 0, sizeof(pi));
						}
					}
					else
					{
						if (_access(csControlFile, 0) != 0)
						{
							//	SETI's not running but we have no disable file so
							//	start SETI
							me->me->LogEvent("Restarting %s on request", me->m_csExeName);
							goto restart;
						}
					}

					continue;

				case WAIT_OBJECT_0:
					bStop = TRUE;
					break;

				case WAIT_OBJECT_0 + 1:
					goto restart;
				}
			}

			TerminateProcess(me->m_hArray[PROCESSHANDLE], 1);
			delete me;
		}
		else
			me->me->LogEvent(_T("Attempting to start %s returned errcode %d"), me->m_csExeName, GetLastError());
	}

	_endthread();
}