Shell Extensions for .NET Assemblies

// AssemblyInfo.cpp : Implementation of CAssemblyInfo

#include "stdafx.h"
#include "resource.h"
#include "AssemblyInfo.h"
#include <oleauto.h>

#define MAKEPTR(cast, ptr, addValue) (cast)((DWORD_PTR)(ptr) + (DWORD_PTR)(addValue))

// Define the char array for hex characters.
const WCHAR c_szHex[] = L"0123456789abcdefg";

// CAssemblyInfo

// Returns success codes of S_OK and S_FALSE in place of the optional boolean return
// value, or an error code and IErrorInfo otherwise.
STDMETHODIMP CAssemblyInfo::IsAssembly(BSTR path, VARIANT_BOOL* retVal)
{
	HRESULT hr;
	FILEINFO FileInfo = { FILEINFOTYPE_CORHEADER };

	hr = GetFileInfo(path, &FileInfo);

	// Don't err when retVal is NULL because
	// internal implementation may pass NULL.
	if (SUCCEEDED(hr) && retVal)
	{
		*retVal = FileInfo.varfCorHeader;
	}

	return hr;
}

STDMETHODIMP CAssemblyInfo::GetFileType(BSTR path, FileType* retVal)
{
	CW2TEX<MAX_PATH> szPath(path);
	FILEINFO FileInfo = { FILEINFOTYPE(FILEINFOTYPE_CORHEADER | FILEINFOTYPE_MACHINEARCH) };
	FileType ft = Other; // Default to Other.
	HRESULT hr = NOERROR;
	LPTSTR pszExtension;

	if (path && retVal)
	{
		// Check the extension first to boost performance.
		pszExtension = PathFindExtension(szPath);
		if (_tcsicmp(pszExtension, TEXT(".dll")) == 0)
		{
			ft = Win32Library;
		}
		else if (_tcsicmp(pszExtension, TEXT(".exe")) == 0)
		{
			ft = Win32Executable;
		}
		else if (_tcsicmp(pszExtension, TEXT(".netmodule")) == 0)
		{
			ft = DotNETModule;
		}

		// For EXE, DLL, and MOD get the file information.
		if (Other != ft)
		{
			hr = GetFileInfo(path, &FileInfo);
			if (SUCCEEDED(hr))
			{
				// If this is a managed assembly and not a .netmodule increment the FileType
				// by the relative offset for managed assemblies.
				if (DotNETModule != ft && VARIANT_TRUE == FileInfo.varfCorHeader)
					ft = FileType(ft + OFFSET_DOTNET);

				// If the file targets IA64 increment relative offset for IA64.
				if (IMAGE_FILE_MACHINE_IA64 == FileInfo.wMachineArch)
					ft = FileType(ft + OFFSET_IA64);
				// Else if the file targets X64 increment by relative offset for X64.
				else if (IMAGE_FILE_MACHINE_AMD64 == FileInfo.wMachineArch)
					ft = FileType(ft + OFFSET_X64);
			}
		}

		*retVal = ft;
	}
	else
	{
		hr = E_POINTER;
	}

	return hr;
}

STDMETHODIMP CAssemblyInfo::GetPublicKeyToken(BSTR path, BSTR* retVal)
{
	BYTE b;
	CComBSTR* bstrToken = NULL;
	HRESULT hr;
	LPBYTE pToken = NULL;
	ULONG cbToken = 0;
	ULONG i, j;

	static BOOL fError = FALSE;

	if (path && retVal)
	{
		// Default the return value to NULL.
		*retVal = NULL;

		// First make sure the path points to an assembly.
		hr = IsAssembly(path, NULL);
		if (S_OK == hr)
		{
			try
			{
				if (!fError && g_mscoree.StrongNameTokenFromAssembly(path, &pToken, &cbToken))
				{
					bstrToken = new CComBSTR(cbToken * 2); // 2 chars / byte for hex.

					// Encode the BYTE[] to hexadecimal.
					for (i = 0, j = 0; i < cbToken; i++)
					{
						b = pToken[i];
						bstrToken->m_str[j++] = c_szHex[b >> 4];
						bstrToken->m_str[j++] = c_szHex[b & 0xf];
					}

					// Free the buffer allocated by StrongNameTokenFromAssembly.
					g_mscoree.StrongNameFreeBuffer(pToken);

					bstrToken->CopyTo(retVal);
				}
			}
			catch (...)
			{
				fError = TRUE;
			}
		}
	}
	else
	{
		hr = E_POINTER;
	}

	if (bstrToken)
	{
		delete bstrToken;
		bstrToken = NULL;
	}

	return hr;
}

STDMETHODIMP CAssemblyInfo::GetPublicKey(BSTR path, SAFEARRAY** retVal)
{
	HRESULT hr;
	LPBYTE pKeyBlob = NULL;
	LPBYTE pToken = NULL;
	SAFEARRAY* psa;
	ULONG cbKeyBlob = 0;
	ULONG cbToken = 0;
	LPBYTE pbData = NULL;
	
	static BOOL fError = FALSE;

	if (path && retVal)
	{
		// First make sure the path points to an assembly.
		hr = IsAssembly(path, NULL);
		if (S_OK == hr)
		{
			try
			{
				if (!fError && g_mscoree.StrongNameTokenFromAssemblyEx(path, &pToken, &cbToken, &pKeyBlob, &cbKeyBlob) && cbKeyBlob > 0)
				{
					psa = SafeArrayCreateVector(VT_UI1, 0, cbKeyBlob);
					if (psa)
					{
						// Lock the SA and copy the element data.
						hr = SafeArrayAccessData(psa, (LPVOID*)&pbData);
						if (SUCCEEDED(hr))
						{
							memcpy(pbData, pKeyBlob, (size_t)cbKeyBlob);
							hr = SafeArrayUnaccessData(psa);
							if (SUCCEEDED(hr))
								*retVal = psa;
						}
					}
					else
						hr = E_OUTOFMEMORY;

					// Free the buffers allocated by StrongNameTokenFromAssembly.
					g_mscoree.StrongNameFreeBuffer(pToken);
					g_mscoree.StrongNameFreeBuffer(pKeyBlob);

				}
			}
			catch (...)
			{
				fError = TRUE;
			}
		}
	}
	else
	{
		hr = E_POINTER;
	}

	return hr;
}

HRESULT CAssemblyInfo::GetFileInfo(BSTR path, LPFILEINFO pFileInfo)
{
	CAtlFile file;
	CAtlFileMappingBase map;
	HRESULT hr;
	LPVOID pFile;

	if (path && pFileInfo)
	{
		// Open the file image referenced by the path parameter.
		hr = file.Create(CW2TEX<MAX_PATH>(path), GENERIC_READ, FILE_SHARE_READ, OPEN_EXISTING,
			FILE_ATTRIBUTE_READONLY | FILE_FLAG_RANDOM_ACCESS);
		if (FAILED(hr))
			return AtlReportError(GetObjectCLSID(), ERR_FILE_OPEN, IID_IAssemblyInfo, hr);

		// Map the file into memory. The destructor will unmap the file, which will
		// release the file handle when unmapped.
		hr = map.MapFile(file);
		if (FAILED(hr))
		{
			file.Close();
			return AtlReportError(GetObjectCLSID(), ERR_FILE_MAP, IID_IAssemblyInfo, hr);
		}

		pFile = map.GetData();

		// Determine if the COM+ header is found in the file.
		if (0 != (pFileInfo->fit & FILEINFOTYPE_CORHEADER))
		{
			hr = HasCorHeader(pFile);
			if (SUCCEEDED(hr))
			{
				pFileInfo->varfCorHeader = (S_OK == hr) ? VARIANT_TRUE : VARIANT_FALSE;
			}
		}

		// Get the machine architecture from the executable.
		if (SUCCEEDED(hr) && 0 != (pFileInfo->fit & FILEINFOTYPE_MACHINEARCH))
		{
			hr = GetMachineArch(pFile, &pFileInfo->wMachineArch);
		}

		// Close the file and return.
		file.Close();
	}
	else
	{
		hr = E_POINTER;
	}

	return hr;
}

HRESULT CAssemblyInfo::GetNtHeaders(LPVOID pFile, PIMAGE_NT_HEADERS* ppNtHeaders)
{
	HRESULT hr;
	PIMAGE_DOS_HEADER pDosHeader;

	if (pFile && ppNtHeaders)
	{
		// Get the PE header from the file.
		pDosHeader = reinterpret_cast<PIMAGE_DOS_HEADER>(pFile);
		if (IMAGE_DOS_SIGNATURE == pDosHeader->e_magic)
		{
			*ppNtHeaders = MAKEPTR(PIMAGE_NT_HEADERS, pDosHeader, pDosHeader->e_lfanew);
			hr = S_OK;
		}
		else
		{
			// The PE/COFF executable is not supported.
			hr = CO_E_NOT_SUPPORTED;
		}
	}
	else
	{
		hr = E_POINTER;
	}

	return hr;
}

HRESULT CAssemblyInfo::GetComDirectory(PIMAGE_NT_HEADERS pNtHeaders, PIMAGE_DATA_DIRECTORY pComDirectory)
{
	HRESULT hr;

	if (pNtHeaders && pComDirectory)
	{
		// Get the NT headers.
		if (IMAGE_NT_SIGNATURE == pNtHeaders->Signature)
		{
			// Get the data directories from the appropriate optional NT header.
			switch (pNtHeaders->OptionalHeader.Magic)
			{
			case IMAGE_NT_OPTIONAL_HDR32_MAGIC:
				*pComDirectory = reinterpret_cast<PIMAGE_NT_HEADERS32>
					(pNtHeaders)->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR];
				hr = S_OK;
				break;

			case IMAGE_NT_OPTIONAL_HDR64_MAGIC:
				*pComDirectory = reinterpret_cast<PIMAGE_NT_HEADERS64>
					(pNtHeaders)->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR];
				hr = S_OK;
				break;

			default:
				hr = CO_E_NOT_SUPPORTED;
			}
		}
		else
		{
			// The header signature is invalid or the header format is not supported.
			hr = CO_E_NOT_SUPPORTED;
		}
	}
	else
	{
		hr = E_POINTER;
	}

	return hr;
}

HRESULT CAssemblyInfo::HasCorHeader(LPVOID pFile)
{
	HRESULT hr;
	PIMAGE_NT_HEADERS pNtHeaders = NULL;
	IMAGE_DATA_DIRECTORY ComDirectory;

	hr = GetNtHeaders(pFile, &pNtHeaders);
	if (SUCCEEDED(hr))
	{
		hr = GetComDirectory(pNtHeaders, &ComDirectory);
		if (SUCCEEDED(hr))
		{
			// If the VirtualAddress is set, a valid COM+ header is present.
			hr = (ComDirectory.VirtualAddress && ComDirectory.Size) ? S_OK : S_FALSE;
		}
	}

	return hr;
}

HRESULT CAssemblyInfo::GetMachineArch(LPVOID pFile, LPWORD pwMachine)
{
	HRESULT hr;
	PIMAGE_NT_HEADERS pNtHeaders = NULL;

	if (pwMachine)
	{
		hr = GetNtHeaders(pFile, &pNtHeaders);
		if (SUCCEEDED(hr) && pNtHeaders)
		{
			*pwMachine = pNtHeaders->FileHeader.Machine;
		}
	}
	else
	{
		hr = E_POINTER;
	}

	return hr;
}