Function Call Tracing in JScript

#include "stdafx.h"
#include "JSDUtils.h"
#include <atlbase.h>
#include "UtilityFunctions.h"

LPCWSTR szJSCRIPT_INPROC_KEY = L"CLSID\\" _T(szCLSID_JScript) L"\\InProcServer32";
LPCWSTR szDEBUG_APPS_KEY = L"JScriptDebug\\TraceApps";

BOOL JSDUtils::IsReplaceMode()
{
	CRegKey reg;
	BOOL bReplaceMode = TRUE;
	if (ERROR_SUCCESS == reg.Open(HKEY_CLASSES_ROOT,szJSCRIPT_INPROC_KEY))
	{
		WCHAR szBuffer[MAX_PATH+3];
		ULONG uSize = 256;
		if (ERROR_SUCCESS == reg.QueryValue(szBuffer,NULL,&uSize))
		{
			CComBSTR bstrServer = szBuffer;
			bstrServer.ToLower();

			// If its the original we are not in replace mode so leave debugging turned on
			if (wcsstr(bstrServer,L"jscript.dll"))
				bReplaceMode = FALSE;
		}
	}

	return bReplaceMode;
}

HRESULT JSDUtils::SetReplaceMode(BOOL bReplace)
{
	HRESULT hRes = E_FAIL;

	_bstr_t bstrModule;
	_bstr_t bstrProgId(L"JScriptDebug");
	if (bReplace)
		hRes = GetComponentPath(bstrProgId,bstrModule);
	else
	{
		WCHAR szBuffer[MAX_PATH];
		::GetWindowsDirectory(szBuffer,MAX_PATH);
		
		wcscat(szBuffer,L"\\System32\\JScript.dll");			// should be enough room !
		bstrModule = szBuffer;

		if (CheckFilePath(bstrModule))
			hRes = S_OK;
	}
	
	if (SUCCEEDED(hRes))
	{
		CRegKey reg;
		if (ERROR_SUCCESS == reg.Open(HKEY_CLASSES_ROOT,szJSCRIPT_INPROC_KEY))
		{
			if (ERROR_SUCCESS == reg.SetValue(bstrModule))
				hRes = S_OK;
		}
	}

	return hRes;
}

int JSDUtils::GetDebugApps(APPLIST& lstApps)
{
	lstApps.clear();

	CRegKey reg;
	if (ERROR_SUCCESS == reg.Open(HKEY_CLASSES_ROOT,szDEBUG_APPS_KEY))
	{
		DWORD dwIndex = 0;
		_bstr_t bstrName;
		_variant_t vntValue;
		while (ERROR_SUCCESS == EnumRegValues(reg,dwIndex,bstrName,vntValue))
		{
			if (bstrName.length())
				lstApps.push_back(new APPLICATION(bstrName,vntValue.lVal));
		}
	}

	return lstApps.size();
}

BOOL JSDUtils::SetDebugApp(LPCWSTR szModuleName,BOOL bEnabled)
{
	CRegKey reg;
	long lRet = E_FAIL;
	if (ERROR_SUCCESS == reg.Open(HKEY_CLASSES_ROOT,szDEBUG_APPS_KEY))
		lRet = reg.SetValue((DWORD)bEnabled,szModuleName);

	return lRet == ERROR_SUCCESS;
}

BOOL JSDUtils::RemoveDebugApp(LPCWSTR szModuleName)
{
	CRegKey reg;
	long lRet = E_FAIL;
	if (ERROR_SUCCESS == reg.Open(HKEY_CLASSES_ROOT,szDEBUG_APPS_KEY))
		lRet = reg.DeleteValue(szModuleName);

	return lRet == ERROR_SUCCESS;
}

HRESULT JSDUtils::IsDebugApp(LPCWSTR szModuleName)
{
	// Debugging is turned ON by default. If the registry has been modified so
	// that this engine is being used wherever jscript.dll would normally be used
	// then debugging is OFF.  To turn it on in this case you have to add the 
	// process name to TraceApps key of JScriptDebug.  This is also required even
	// if you have:
	//
	//					<SCRIPT language="JScriptDebug"> ... 
	//

	CRegKey reg;
		
	ATLASSERT(szModuleName && wcslen(szModuleName));
	if (NULL == szModuleName || wcslen(szModuleName) == 0)
		return E_INVALIDARG;

	_bstr_t bstrModule = szModuleName;
	
	APPLIST apps;
	if (GetDebugApps(apps))
	{
		APPLIST::iterator iter = apps.begin();
		while(iter != apps.end())
		{
			PAPPLICATION pApp = *iter;
			if (wcsicmp(pApp->m_sPath.c_str(),bstrModule) == 0 && pApp->m_bEnabled == TRUE)
				return S_OK;
			
			iter++;
		}
	}

	return S_FALSE;
}

// When in "Replace Mode" (i.e. JScript/InProcServer32 points to this dll) an applications
// must be listed under JScriptDebug in HKCR in order to use this debugging version.  This
// is to avoid potential problems with system processes or other applications which use jscript.
// If the calling app is not listed in the TraceApps key then the following function
// returns a reference to an instance of the real JScript class factory

HRESULT JSDUtils::GetJScriptClassObject(LPVOID* ppv)
{
	// Note - this has to use the hard-coded string for jscript - because one of the 
	// simplest ways to do debugging is to temporarily replace the jscript CLSID with
	// ours - we need the original one here.  We also have to load the class factory 
	// and create the instance from the reall jscript dll - just calling CoCreateInstance
	// will cause an infinite recursion when COM calls this code again.

	CLSID clsidJScript;
	HRESULT hRes = CLSIDFromString(_T(szCLSID_JScript),&clsidJScript);

	HINSTANCE hJScriptDLL = ::CoLoadLibrary(L"jscript.dll",FALSE);
	if (NULL != hJScriptDLL)
	{
		fnDLLGETCLASSOBJECT pfnGCO = (fnDLLGETCLASSOBJECT)GetProcAddress(hJScriptDLL,"DllGetClassObject");
		if (pfnGCO)
			hRes = pfnGCO(clsidJScript,IID_IClassFactory,ppv);
	}

	return hRes;
} 

BOOL JSDUtils::readSettings(
	BOOL* pbTrace,
	BOOL* pbInstrument,
	BOOL* pbBlankLines,
	BOOL* pbRaiseErrors,
	int*  piIndent,
	LPWSTR szScriptDump)
{
	BOOL bOK = FALSE;

	CRegKey reg;
	if (ERROR_SUCCESS == reg.Open(HKEY_CLASSES_ROOT,L"JScriptDebug\\Settings"))
	{
		// Should tracing begin straight away or only when turned on
		DWORD dw = 0;
		if (pbTrace && ERROR_SUCCESS == reg.QueryValue(dw,L"Trace"))
			*pbTrace = dw ? true : false;

		// Is fn entry/exit instrumentation used ?
		dw = 0;
		if (pbInstrument && ERROR_SUCCESS == reg.QueryValue(dw,L"Instrument"))
			*pbInstrument = dw ? true : false;

		// Blank line between trace statements ?
		dw = 0;
		if (pbBlankLines && ERROR_SUCCESS == reg.QueryValue(dw,L"BlankLines"))
			*pbBlankLines = dw ? true : false;

		// Set indent size - default is 4
		dw = 0;
		if (piIndent && ERROR_SUCCESS == reg.QueryValue(dw,L"Indent"))
			*piIndent = dw;

		// Raise errors
		dw = 0;
		if (pbRaiseErrors && ERROR_SUCCESS == reg.QueryValue(dw,L"_RaiseErrors"))
			*pbRaiseErrors = dw ? true : false;

		if (szScriptDump)
		{
			dw = MAX_PATH;
			reg.QueryValue(szScriptDump,L"_ScriptDump",&dw);
		}

		bOK = TRUE;

	}
	
	return bOK;
}

BOOL JSDUtils::writeSettings(
	BOOL* pbTrace,
	BOOL* pbInstrument,
	BOOL* pbBlankLines,
	BOOL* pbRaiseErrors,
	int*  piIndent,
	LPWSTR szScriptDump)
{
	BOOL bOK = FALSE;

	CRegKey reg;
	if (ERROR_SUCCESS == reg.Open(HKEY_CLASSES_ROOT,L"JScriptDebug\\Settings"))
	{
		// Should tracing begin straight away or only when turned on
		DWORD dw = 0;
		if (pbTrace)
			reg.SetValue(*pbTrace ? 1 : 0,L"Trace");

		// Is fn entry/exit instrumentation used ?
		if (pbInstrument)
			reg.SetValue(*pbInstrument ? 1 : 0,L"Instrument");

		// Blank line between trace statements ?
		if (pbBlankLines)
			reg.SetValue(*pbBlankLines ? 1 : 0,L"BlankLines");

		// Set indent size - default is 4
		if (piIndent)
			reg.SetValue((DWORD)*piIndent,L"Indent");

		// Raise errors
		if (pbRaiseErrors)
			reg.SetValue(*pbRaiseErrors ? 1 : 0,L"_RaiseErrors");

		if (szScriptDump)
			reg.SetValue(szScriptDump,L"_ScriptDump");

		bOK = TRUE;

	}
	
	return bOK;
}