source-engine/utils/vmpi/vmpi_service_ui/vmpi_service_ui.cpp

//========= Copyright Valve Corporation, All rights reserved. ============//
//
// Purpose: 
//
// $NoKeywords: $
//
//=============================================================================//
// vmpi_service_ui.cpp : Defines the entry point for the application.
//

#include "stdafx.h"
#include "consolewnd.h"
#include "resource.h"
#include "tier0/dbg.h"
#include "tier1/strtools.h"
#include "shell_icon_mgr.h"
#include "vmpi.h"
#include "service_conn_mgr.h"
#include <io.h>
#include <time.h>



void UpdatePopupMenuState();



const char *g_pIconTooltip = VMPI_SERVICE_NAME;

IConsoleWnd *g_pConsoleWnd = NULL;
HINSTANCE g_hInstance = NULL;


#define MYWM_NOTIFYICON		(WM_APP+100)
CShellIconMgr g_ShellIconMgr;

bool g_bHighlightIconWhenBusy = false;


// STATE THE SERVICE OWNS.
int g_iCurState = 0; // One of the VMPI_SERVICE_STATE_ defines.
char *g_pPassword = NULL;
bool g_bScreensaverMode = false;


void LogString( const char *pStr, ... )
{
#ifdef VMPI_SERVICE_LOGS
	char str[4096];
	va_list marker;
	va_start( marker, pStr );
	_vsnprintf( str, sizeof( str ), pStr, marker );
	va_end( marker );
	
	static FILE *fp = fopen( "c:\\vmpi_service_ui.log", "wt" );
	if ( fp )
	{
		fprintf( fp, "%s", str );
		fflush( fp );
	}
#endif
}

char* GetLastErrorString()
{
	static char err[2048];
	
	LPVOID lpMsgBuf;
	FormatMessage( FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), (LPTSTR) &lpMsgBuf, 0, NULL );
	strncpy( err, (char*)lpMsgBuf, sizeof( err ) );
	LocalFree( lpMsgBuf );

	err[ sizeof( err ) - 1 ] = 0;
	return err;
}


// ------------------------------------------------------------------------------------------ //
// Persistent state in the registry.
// ------------------------------------------------------------------------------------------ //
void LoadStateFromRegistry()
{
	HKEY hKey = NULL;
	RegCreateKey( HKEY_LOCAL_MACHINE, VMPI_SERVICE_KEY, &hKey );
	if ( hKey )
	{
		DWORD val = 0;
		DWORD type = REG_DWORD;
		DWORD size = sizeof( val );
		if ( RegQueryValueEx( 
			hKey,
			"HighlightIconWhenBusy",
			0,
			&type,
			(unsigned char*)&val,
			&size ) == ERROR_SUCCESS && 
			type == REG_DWORD && 
			size == sizeof( val ) )
		{
			g_bHighlightIconWhenBusy = (val != 0);
		}
	}
}

void SaveStateToRegistry()
{
	HKEY hKey = NULL;
	RegCreateKey( HKEY_LOCAL_MACHINE, VMPI_SERVICE_KEY, &hKey );
	if ( hKey )
	{
		DWORD val = g_bHighlightIconWhenBusy;
		RegSetValueEx( 
			hKey,
			"HighlightIconWhenBusy",
			0,
			REG_DWORD,
			(unsigned char*)&val,
			sizeof( val ) );
	}
}


// ------------------------------------------------------------------------------------------ //
// Our CServiceConnMgr packet handler.
// ------------------------------------------------------------------------------------------ //

void UpdateAppIcon()
{
	if ( g_iCurState == VMPI_SERVICE_STATE_IDLE )
	{
		g_ShellIconMgr.ChangeIcon( IDI_WAITING_ICON );
	}
	else if ( g_iCurState == VMPI_SERVICE_STATE_BUSY )
	{
		if ( g_bHighlightIconWhenBusy )
			g_ShellIconMgr.ChangeIcon( IDI_BUSY_ICON );
		else
			g_ShellIconMgr.ChangeIcon( IDI_WAITING_ICON );
	}
	else
	{
		g_ShellIconMgr.ChangeIcon( IDI_DISABLED_ICON );
	}
}


class CUIConnMgr : public CServiceConnMgr
{
public:
	virtual void	HandlePacket( const char *pData, int len );
};


void CUIConnMgr::HandlePacket( const char *pData, int len )
{
	if ( pData[0] != VMPI_SERVICE_UI_PROTOCOL_VERSION )
		return;
	
	int packetID = pData[1];
	
	int offset = 2;
	if ( packetID == VMPI_SERVICE_TO_UI_CONSOLE_TEXT )
	{
		Msg( &pData[offset] );
	}
	else if ( packetID == VMPI_SERVICE_TO_UI_STATE )
	{
		// Get the new state out..
		g_iCurState = *((int*)&pData[offset]);
		offset += 4;

		LogString( "New UI state: %d.\n", g_iCurState );
		
		// Update our icon.
		UpdateAppIcon();

		g_bScreensaverMode = (*((char*)&pData[offset]) != 0);
		++offset;

		// Store the current password.
		if ( g_pPassword )
			delete [] g_pPassword;

		const char *pStr = &pData[offset];
		g_pPassword = new char[ strlen( pStr ) + 1 ];
		strcpy( g_pPassword, pStr );
		offset += strlen( pStr ) + 1;

		UpdatePopupMenuState();
	}
	else if ( packetID == VMPI_SERVICE_TO_UI_PATCHING )
	{
		LogString( "Got a VMPI_SERVICE_TO_UI_PATCHING packet.\n" );
		
		int bExitAfter = pData[offset];
		++offset;
		
		char workingDir[MAX_PATH], commandLine[4096];
		V_strncpy( workingDir, &pData[offset], sizeof( workingDir ) );
		offset += V_strlen( workingDir ) + 1;
		
		V_strncpy( commandLine, &pData[offset], sizeof( commandLine ) );
		offset += V_strlen( commandLine ) + 1;
		
		// Run whatever they said to run.
		STARTUPINFO si;
		memset( &si, 0, sizeof( si ) );
		si.cb = sizeof( si );
		
		PROCESS_INFORMATION pi;
		memset( &pi, 0, sizeof( pi ) );
		if ( CreateProcess( NULL, commandLine, NULL, NULL, false, CREATE_NO_WINDOW, NULL, workingDir, &si, &pi ) )
		{
			LogString( "CreateProcess succeeded:\n%s\n", commandLine );

			CloseHandle( pi.hProcess );
			CloseHandle( pi.hThread );
			
			if ( bExitAfter )
				PostQuitMessage( 0 );
		}
		else
		{
			LogString( "CreateProcess failed: %s", GetLastErrorString() );
		}
	}
	else if ( packetID == VMPI_SERVICE_TO_UI_EXIT )
	{
		// Exit.
		PostQuitMessage( 0 );
	}
}


// ------------------------------------------------------------------------------------------ //
// Helpers.
// ------------------------------------------------------------------------------------------ //

void InternalSpew( const char *pMsg )
{
	if ( g_pConsoleWnd )
	{
		g_pConsoleWnd->PrintToConsole( pMsg );
	}

	OutputDebugString( pMsg );
}


SpewRetval_t MySpewOutputFunc( SpewType_t spewType, const char *pMsg )
{
	// Prepend the time.
	time_t aclock;
	time( &aclock );
	struct tm *newtime = localtime( &aclock );
	
	// Get rid of the \n.
	char timeString[512];
	Q_strncpy( timeString, asctime( newtime ), sizeof( timeString ) );
	char *pEnd = strstr( timeString, "\n" );
	if ( pEnd )
		*pEnd = 0;
	InternalSpew( timeString );
	
	InternalSpew( " - " );
	InternalSpew( pMsg );
	
	if ( spewType == SPEW_ASSERT )
		return SPEW_DEBUGGER;
	else if( spewType == SPEW_ERROR )
		return SPEW_ABORT;
	else
		return SPEW_CONTINUE;
}



CUIConnMgr g_ConnMgr;

void InitConsoleWindow()
{
	// Only initialize it once.
	if ( g_pConsoleWnd )
		return;

	g_pConsoleWnd = CreateConsoleWnd( 
		g_hInstance, 
		IDD_VMPI_SERVICE, 
		IDC_DEBUG_OUTPUT, 
		false );
}


// ------------------------------------------------------------------------------------------ //
// Implementation of IShellIconMgrHelper.
// ------------------------------------------------------------------------------------------ //

HMENU g_hMenu = NULL;
HMENU g_hPopupMenu = NULL;	// This is just a submenu of g_hMenu.


bool LoadPopupMenu()
{
	g_hMenu = LoadMenu( g_hInstance, MAKEINTRESOURCE( IDR_POPUP_MENU ) );
	if ( !g_hMenu )
	{
		Assert( false );
		Warning( "LoadMenu failed.\n" );
		return false;
	}

	g_hPopupMenu = GetSubMenu( g_hMenu, 0 );
	return true;
}


void TermPopupMenu()
{
	if ( g_hMenu )
	{
		DestroyMenu( g_hMenu );
		g_hMenu = NULL;
	}
	g_hPopupMenu = NULL;
}


void UpdatePopupMenuState()
{
	bool bEnabled = (g_iCurState == VMPI_SERVICE_STATE_IDLE || g_iCurState == VMPI_SERVICE_STATE_BUSY);
	EnableMenuItem( g_hPopupMenu, ID_ENABLE_WORKER,  bEnabled  ? MF_GRAYED : MF_ENABLED );
	EnableMenuItem( g_hPopupMenu, ID_DISABLE_WORKER, !bEnabled ? MF_GRAYED : MF_ENABLED );

	// Enable or disable console items.
	EnableMenuItem( g_hPopupMenu, ID_SHOW_CONSOLE_WINDOW,  g_pConsoleWnd->IsVisible() ? MF_GRAYED : MF_ENABLED );
	EnableMenuItem( g_hPopupMenu, ID_HIDE_CONSOLE_WINDOW, !g_pConsoleWnd->IsVisible() ? MF_GRAYED : MF_ENABLED );

	CheckMenuItem( g_hPopupMenu, ID_SCREENSAVER_MODE, g_bScreensaverMode ? MF_CHECKED : MF_UNCHECKED );
	CheckMenuItem( g_hPopupMenu, ID_HIGHLIGHT_ICON_WHEN_BUSY, g_bHighlightIconWhenBusy ? MF_CHECKED : MF_UNCHECKED );
}


int CALLBACK SetPasswordDlgProc(
  HWND hwndDlg,  // handle to dialog box
  UINT uMsg,     // message
  WPARAM wParam, // first message parameter
  LPARAM lParam  // second message parameter
)
{
	switch( uMsg )
	{
		case WM_INITDIALOG:
		{
			if ( g_pPassword )
			{
				HWND hWnd = GetDlgItem( hwndDlg, IDC_PASSWORD );
				SetWindowText( hWnd, g_pPassword );
			}
		}
		break;

		case WM_COMMAND:
		{
			switch( wParam )
			{
				case IDOK:
				{
					// Set our new password.
					HWND hWnd = GetDlgItem( hwndDlg, IDC_PASSWORD );
					if ( hWnd )
					{
						char tempBuf[512];
						GetWindowText( hWnd, tempBuf, sizeof( tempBuf ) );

						// Send it to the service.
						CUtlVector<char> data;
						data.AddToTail( VMPI_SERVICE_UPDATE_PASSWORD );
						data.AddMultipleToTail( strlen( tempBuf ) + 1, tempBuf );
						g_ConnMgr.SendPacket( -1, data.Base(), data.Count() );
					}
					EndDialog( hwndDlg, 0 );
				}
				break;
			
				case IDCANCEL:
				{
					EndDialog( hwndDlg, 0 );
				}
				break;
			}
		}
		break;
	}

	return FALSE;
}


class CShellIconMgrHelper : public IShellIconMgrHelper
{
public:
	virtual HINSTANCE	GetHInstance()
	{
		return g_hInstance;
	}

	virtual int		WindowProc( void *pWnd, int uMsg, long wParam, long lParam )
	{
		HWND hWnd = (HWND)pWnd;

		switch( uMsg )
		{
			// Right button brings up the popup menu.
			case MYWM_NOTIFYICON:
			{
				if ( lParam == WM_RBUTTONDOWN )
				{
					POINT cursorPos;
					GetCursorPos( &cursorPos );

					UpdatePopupMenuState();

					// Make a popup menu.
					SetForegroundWindow( hWnd );
					TrackPopupMenu( g_hPopupMenu, TPM_RIGHTALIGN | TPM_BOTTOMALIGN, cursorPos.x, cursorPos.y, 0, hWnd, NULL );
					return 0;
				}
				else if ( lParam == WM_LBUTTONDOWN )
				{
					// Left button brings up the console.
					g_pConsoleWnd->SetVisible( true );
					UpdatePopupMenuState();
					return 0;
				}
			}
			break;
			
			case WM_COMMAND:
			{
				switch( wParam )
				{
					case ID_ENABLE_WORKER:
					{
						char cPacket = VMPI_SERVICE_ENABLE;
						g_ConnMgr.SendPacket( -1, &cPacket, 1 );
					}
					break;

					case ID_DISABLE_WORKER:
					{
						char cPacket = VMPI_SERVICE_DISABLE;
						g_ConnMgr.SendPacket( -1, &cPacket, 1 );
					}
					break;
					
					case ID_KILLCURRENTJOB:
					{
						char cPacket = VMPI_KILL_PROCESS;
						g_ConnMgr.SendPacket( -1, &cPacket, 1 );
					}
					break;

					case ID_HIGHLIGHT_ICON_WHEN_BUSY:
					{
						g_bHighlightIconWhenBusy = !g_bHighlightIconWhenBusy;
						SaveStateToRegistry();
						UpdateAppIcon();
						UpdatePopupMenuState();
					}
					break;

					case ID_SCREENSAVER_MODE:
					{
						g_bScreensaverMode = !g_bScreensaverMode;
						char cPacket[2] = { VMPI_SERVICE_SCREENSAVER_MODE, g_bScreensaverMode };
						g_ConnMgr.SendPacket( -1, cPacket, sizeof( cPacket ) );
					}
					break;


					case ID_SHOW_CONSOLE_WINDOW:
					{
						g_pConsoleWnd->SetVisible( true );
						UpdatePopupMenuState();
					}
					break;

					case ID_HIDE_CONSOLE_WINDOW:
					{
						g_pConsoleWnd->SetVisible( false );
						UpdatePopupMenuState();
					}
					break;
				
					case ID_SET_PASSWORD:
					{
						DialogBox( g_hInstance, MAKEINTRESOURCE( IDD_SET_PASSWORD ), NULL, SetPasswordDlgProc );
					}
					break;

					case ID_EXIT_SERVICE:
					{
						// Quit the service app..
						char cPacket = VMPI_SERVICE_EXIT;
						g_ConnMgr.SendPacket( -1, &cPacket, 1 );

						// Stop showing the icon.
						g_ShellIconMgr.Term();

						// Wait for a bit for the connection to go away.
						DWORD startTime = GetTickCount();
						while ( GetTickCount()-startTime < 2000 )
						{
							g_ConnMgr.Update();
							if ( !g_ConnMgr.IsConnected() )
								break;
							else
								Sleep( 10 );
						}							
						
						// Quit the UI app.
						PostQuitMessage( 0 );
						return 1;
					}
					break;
				}
			}
			break;
		}	

		return DefWindowProc( (HWND)hWnd, uMsg, wParam, lParam );
	}
};

CShellIconMgrHelper g_ShellIconMgrHelper;



int APIENTRY WinMain(HINSTANCE hInstance,
                     HINSTANCE hPrevInstance,
                     LPSTR     lpCmdLine,
                     int       nCmdShow)
{
 	g_hInstance = hInstance;


	LogString( "vmpi_service_ui startup.\n" );

	// Don't run multiple instances.	
	HANDLE hMutex = CreateMutex( NULL, FALSE, "vmpi_service_ui_mutex" );
	if ( hMutex && GetLastError() == ERROR_ALREADY_EXISTS )
		return 1;


	// Hook spew output.
	SpewOutputFunc( MySpewOutputFunc );
	InitConsoleWindow();
	LogString( "Setup console window.\n" );

	LoadStateFromRegistry();

	// Setup the popup menu.
	if( !LoadPopupMenu() )
	{
		return false;
	}
	UpdatePopupMenuState();


	// Setup the tray icon.
	Msg( "Waiting for jobs...\n" );
	if ( !g_ShellIconMgr.Init( &g_ShellIconMgrHelper, g_pIconTooltip, MYWM_NOTIFYICON, IDI_WAITING_ICON ) )
	{
		return false;
	}


	// Connect to the VMPI service.
	g_ConnMgr.InitClient();

	LogString( "Entering main loop.\n" );
	
	while ( 1 )
	{
		MSG msg;
		msg.message = !WM_QUIT;	// So it doesn't accidentally exit.
		while ( PeekMessage( &msg, NULL, 0, 0, PM_REMOVE ) )
		{
			if ( msg.message == WM_QUIT )
				break;
			
			TranslateMessage( &msg );
			DispatchMessage( &msg );
		}
		if ( msg.message == WM_QUIT )
			break;
	
		g_ConnMgr.Update();
		if ( !g_ConnMgr.IsConnected() )
		{
			g_ShellIconMgr.ChangeIcon( IDI_UNCONNECTED );
		}

		Sleep( 30 );
	}

	// Important that we call this instead of letting the destructor do it because it deletes its 
	// socket and it needs to cleanup some threads.
	g_ConnMgr.Term();
	
	g_ShellIconMgr.Term();
	
	return 0;
}