Icon Handler with ATL

One of the exercises I gave at the recent COM Programming class was to build an Icon Handler that integrates with the Windows Shell, where DLLs should have an icon based on their “bitness” – whether they’re 64-bit or 32-bit Portable Executable (PE).

The Shell provides many opportunities for extensibility. An Icon Handler is one of the simplest, but still requires writing a full-fledged COM component that implements certain interfaces that the shell expects. Here is the result of using the Icon Handler DLL, showing the folders c:\Windows\System32 and c:\Windows\SysWow64 (large icons for easier visibility).

Let’s see how to build such an icon handler. The full code is at zodiacon/DllIconHandler.

The first step is to create a new ATL project in Visual Studio. I’ll be using Visual Studio 2022, but any recent version would work essentially the same way (e.g. VS 2019, or 2017). Locate the ATL project type by searching in the terrible new project dialog introduced in VS 2019 and still horrible in VS 2022.

ATL (Active Template Library) is certainly not the only way to build COM components. “Pure” C++ would work as well, but ATL provides all the COM boilerplate such as the required exported functions, class factories, IUnknown implementations, etc. Since ATL is fairly “old”, it lacks the elegance of other libraries such as WRL and WinRT, as it doesn’t take advantage of C++11 and later features. Still, ATL has withstood the test of time, is robust, and full featured when it comes to COM, something I can’t say for these other alternatives.

If you can’t locate the ATL project, you may not have ATL installed propertly. Make sure the C++ Desktop development workload is installed using the Visual Studio Installer.

Click Next and select a project name and location:

Click Create to launch the ATL project wizard. Leave all defaults (Dynamic Link Library) and click OK. Shell extensions of all kinds must be DLLs, as these are loaded by Explorer.exe. It’s not ideal in terms of Explorer’s stability, as aun unhandled exception can bring down the entire process, but this is necessary to get good performance, as no inter-process calls are made.

Two projects are created, named DllIconHandler and DllIconHandlerPS. The latter is a proxy/stub DLL that maybe useful if cross-apartment COM calls are made. This is not needed for shell extensions, so the PS project should simply be removed from the solution.

A detailed discussion of COM is way beyond the scope of this post.

The remaining project contains the COM DLL required code, such as the mandatory exported function, DllGetClassObject, and the other optional but recommended exports (DllRegisterServer, DllUnregisterServer, DllCanUnloadNow and DllInstall). This is one of the nice benefits of working with ATL project for COM component development: all the COM boilerplate is implemented by ATL.

The next step is to add a COM class that will implement our icon handler. Again, we’ll turn to a wizard provided by Visual Studio that provides the fundamentals. Right-click the project and select Add Item… (don’t select Add Class as it’s not good enough). Select the ATL node on the left and ATL Simple Object on the right. Set the name to something like IconHandler:

Click Add. The ATL New Object wizard opens up. The name typed in the Add New Item dialog is used as a basis for generating names for source code elements (like the C++ class) and COM elements (that would be written into the IDL file and the resulting type library). Since we’re not going to define a new interface (we need to implement explorer-defined interfaces), there is no real need to tweak anything. You can click Finish to generate the class.

Three files are added with this last step: IconHandler.h, IconHandler.cpp and IconHandler.rgs. The C++ source files role is obvious – implementing the Icon Handler. The rgs file contains a script in an ATL-provided “language” indicating what information to write to the Registry when this DLL is registered (and what to remove if it’s unregistered).

The IDL (Interface Definition Language) file has also been modified, adding the definitions of the wizard generated interface (which we don’t need) and the coclass. We’ll leave the IDL alone, as we do need it to generate the type library of our component because the ATL registration code uses it internally.

If you look in IconHandler.h, you’ll see that the class implements the IIconHandler empty interface generated by the wizard that we don’t need. It even derives from IDispatch:

class ATL_NO_VTABLE CIconHandler :
	public CComObjectRootEx<CComSingleThreadModel>,
	public CComCoClass<CIconHandler, &CLSID_IconHandler>,
	public IDispatchImpl<IIconHandler, &IID_IIconHandler, &LIBID_DLLIconHandlerLib, /*wMajor =*/ 1, /*wMinor =*/ 0> {

We can leave the IDispatchImpl-inheritance, since it’s harmless. But it’s useless as well, so let’s delete it and also delete the interfaces IIconHandler and IDispatch from the interface map located further down:

class ATL_NO_VTABLE CIconHandler :
	public CComObjectRootEx<CComSingleThreadModel>,
	public CComCoClass<CIconHandler, &CLSID_IconHandler> {
public:
	BEGIN_COM_MAP(CIconHandler)
	END_COM_MAP()

(I have rearranged the code a bit). Now we need to add the interfaces we truly have to implement for an icon handler: IPersistFile and IExtractIcon. To get their definitions, we’ll add an #include for <shlobj_core.h> (this is documented in MSDN). We add the interfaces to the inheritance hierarchy, the COM interface map, and use the Visual Studio feature to add the interface members for us by right-clicking the class name (CIconHandler), pressing Ctrl+. (dot) and selecting Implement all pure virtuals of CIconHandler. The resulting class header looks something like this (some parts omitted for clarity) (I have removed the virtual keyword as it’s inherited and doesn’t have to be specified in derived types):

class ATL_NO_VTABLE CIconHandler :
	public CComObjectRootEx<CComSingleThreadModel>,
	public CComCoClass<CIconHandler, &CLSID_IconHandler>,
	public IPersistFile,
	public IExtractIcon {
public:
	BEGIN_COM_MAP(CIconHandler)
		COM_INTERFACE_ENTRY(IPersistFile)
		COM_INTERFACE_ENTRY(IExtractIcon)
	END_COM_MAP()

//...
	// Inherited via IPersistFile
	HRESULT __stdcall GetClassID(CLSID* pClassID) override;
	HRESULT __stdcall IsDirty(void) override;
	HRESULT __stdcall Load(LPCOLESTR pszFileName, DWORD dwMode) override;
	HRESULT __stdcall Save(LPCOLESTR pszFileName, BOOL fRemember) override;
	HRESULT __stdcall SaveCompleted(LPCOLESTR pszFileName) override;
	HRESULT __stdcall GetCurFile(LPOLESTR* ppszFileName) override;

	// Inherited via IExtractIconW
	HRESULT __stdcall GetIconLocation(UINT uFlags, PWSTR pszIconFile, UINT cchMax, int* piIndex, UINT* pwFlags) override;
	HRESULT __stdcall Extract(PCWSTR pszFile, UINT nIconIndex, HICON* phiconLarge, HICON* phiconSmall, UINT nIconSize) override;
};

Now for the implementation. The IPersistFile interface seems non-trivial, but fortunately we just need to implement the Load method for an icon handler. This is where we get the file name we need to inspect. To check whether a DLL is 64 or 32 bit, we’ll add a simple enumeration and a helper function to the CIconHandler class:

	enum class ModuleBitness {
		Unknown,
		Bit32,
		Bit64
	};
	static ModuleBitness GetModuleBitness(PCWSTR path);

The implementation of IPersistFile::Load looks something like this:

HRESULT __stdcall CIconHandler::Load(LPCOLESTR pszFileName, DWORD dwMode) {
    ATLTRACE(L"CIconHandler::Load %s\n", pszFileName);

    m_Bitness = GetModuleBitness(pszFileName);
    return S_OK;
}

The method receives the full path of the DLL we need to examine. How do we know that only DLL files will be delivered? This has to do with the registration we’ll make for the icon handler. We’ll register it for DLL file extensions only, so that other file types will not be provided. Calling GetModuleBitness (shown later) performs the real work of determining the DLL’s bitness and stores the result in m_Bitness (a data member of type ModuleBitness).

All that’s left to do is tell explorer which icon to use. This is the role of IExtractIcon. The Extract method can be used to provide an icon handle directly, which is useful if the icon is “dynamic” – perhaps generated by different means in each case. In this example, we just need to return one of two icons which have been added as resources to the project (you can find those in the project source code. This is also an opportunity to provide your own icons).

For our case, it’s enough to return S_FALSE from Extract that causes explorer to use the information returned from GetIconLocation. Here is its implementation:

HRESULT __stdcall CIconHandler::GetIconLocation(UINT uFlags, PWSTR pszIconFile, UINT cchMax, int* piIndex, UINT* pwFlags) {
    if (s_ModulePath[0] == 0) {
        ::GetModuleFileName(_AtlBaseModule.GetModuleInstance(), 
            s_ModulePath, _countof(s_ModulePath));
        ATLTRACE(L"Module path: %s\n", s_ModulePath);
    }
    if (s_ModulePath[0] == 0)
        return S_FALSE;

    if (m_Bitness == ModuleBitness::Unknown)
        return S_FALSE;

    wcscpy_s(pszIconFile, wcslen(s_ModulePath) + 1, s_ModulePath);
    ATLTRACE(L"CIconHandler::GetIconLocation: %s bitness: %d\n", 
        pszIconFile, m_Bitness);
    *piIndex = m_Bitness == ModuleBitness::Bit32 ? 0 : 1;
    *pwFlags = GIL_PERINSTANCE;

    return S_OK;
}

The method’s purpose is to return the current (our icon handler DLL) module’s path and the icon index to use. This information is enough for explorer to load the icon itself from the resources. First, we get the module path to where our DLL has been installed. Since this doesn’t change, it’s only retrieved once (with GetModuleFileName) and stored in a static variable (s_ModulePath).

If this fails (unlikely) or the bitness could not be determined (maybe the file was not a PE at all, but just had such an extension), then we return S_FALSE. This tells explorer to use the default icon for the file type (DLL). Otherwise, we store 0 or 1 in piIndex, based on the IDs of the icons (0 corresponds to the lower of the IDs).

Finally, we need to set a flag inside pwFlags to indicate to explorer that this icon extraction is required for every file (GIL_PERINSTANCE). Otherwise, explorer calls IExtractIcon just once for any DLL file, which is the opposite of what we want.

The final piece of the puzzle (in terms of code) is how to determine whether a PE is 64 or 32 bit. This is not the point of this post, as any custom algorithm can be used to provide different icons for different files of the same type. For completeness, here is the code with comments:

CIconHandler::ModuleBitness CIconHandler::GetModuleBitness(PCWSTR path) {
    auto bitness = ModuleBitness::Unknown;
    //
    // open the DLL as a data file
    //
    auto hFile = ::CreateFile(path, GENERIC_READ, FILE_SHARE_READ, 
        nullptr, OPEN_EXISTING, 0, nullptr);
    if (hFile == INVALID_HANDLE_VALUE)
        return bitness;

    //
    // create a memory mapped file to read the PE header
    //
    auto hMemMap = ::CreateFileMapping(hFile, nullptr, PAGE_READONLY, 0, 0, nullptr);
    ::CloseHandle(hFile);
    if (!hMemMap)
        return bitness;

    //
    // map the first page (where the header is located)
    //
    auto p = ::MapViewOfFile(hMemMap, FILE_MAP_READ, 0, 0, 1 << 12);
    if (p) {
        auto header = ::ImageNtHeader(p);
        if (header) {
            auto machine = header->FileHeader.Machine;
            bitness = header->Signature == IMAGE_NT_OPTIONAL_HDR64_MAGIC ||
                machine == IMAGE_FILE_MACHINE_AMD64 || machine == IMAGE_FILE_MACHINE_ARM64 ?
                ModuleBitness::Bit64 : ModuleBitness::Bit32;
        }
        ::UnmapViewOfFile(p);
    }
    ::CloseHandle(hMemMap);
    
    return bitness;
}

To make all this work, there is still one more concern: registration. Normal COM registration is necessary (so that the call to CoCreateInstance issued by explorer has a chance to succeed), but not enough. Another registration is needed to let explorer know that this icon handler exists, and is to be used for files with the extension “DLL”.

Fortunately, ATL provides a convenient mechanism to add Registry settings using a simple script-like configuration, which does not require any code. The added keys/values have been placed in DllIconHandler.rgs like so:

HKCR
{
	NoRemove DllFile
	{
		NoRemove ShellEx
		{
			IconHandler = s '{d913f592-08f1-418a-9428-cc33db97ed60}'
		}
	}
}

This sets an icon handler in HKEY_CLASSES_ROOT\DllFile\ShellEx, where the IconHandler value specifies the CLSID of our component. You can find the CLSID in the IDL file where the coclass element is defined:

[
	uuid(d913f592-08f1-418a-9428-cc33db97ed60)
]
coclass IconHandler {

Replace your own CLSID if you’re building this project from scratch. Registration itself is done with the RegSvr32 built-in tool. With an ATL project, a successful build also causes RegSvr32 to be invoked on the resulting DLL, thus performing registration. The default behavior is to register in HKEY_CLASSES_ROOT which uses HKEY_LOCAL_MACHINE behind the covers. This requires running Visual Studio elevated (or an elevated command window if called from outside VS). It will register the icon handler for all users on the machine. If you prefer to register for the current user only (which uses HKEY_CURRENT_USER and does not require running elevated), you can set the per-user registration in VS by going to project properties, clinking on the Linker element and setting per-user redirection:

If you’re registering from outside VS, the per-user registration is achieved with:

regsvr32 /n /i:user <dllpath>

This is it! The full source code is available here.