The Mini Shell Extension Framework – Part III

//
// (C) Copyright by Victor Derks <vba64@xs4all.nl>
//
// See README.TXT for the details of the software licence.
//
#pragma once


#include <map>
#include <vector>
#include <algorithm>
#include "updateregistry.h"
#include "strutil.h"
#include "catchhandler.h"
#include "shelluuids.h"


namespace MSF
{

// Note: Due to a bug in explorer it will not release the interface if the
//       folder is the 'desktop'. The workaround is to return 0 columns in that case.
//       This will prevent the resource leak. The drawback is that there are no
//       columns available if the 'desktop' folder is viewed in 'detailed' mode.

template <typename T, bool t_bEnableDesktopBugWorkAround = true>
class ATL_NO_VTABLE IColumnProviderImpl : public IColumnProvider
{
public:
	typedef std::map<CStringW, std::vector<CString> > mapCacheInfo;

	struct CCompareShColumnId
	{
		bool operator()(const SHCOLUMNID& left, const SHCOLUMNID& right) const
		{
			return memcmp(&left, &right, sizeof(SHCOLUMNID)) < 0;
		}
	};

	typedef std::map<SHCOLUMNID, unsigned int, CCompareShColumnId> mapColumnIdToIndex;

	template <bool t_bEnableDesktopBugWorkAround>
	class CDesktopBugWorkAround
	{
	public:
		void Initialize(const wchar_t* /*wszFolder*/)
		{
		}

		bool operator()() const throw()
		{
			return false;
		}
	};

	template <>
	class CDesktopBugWorkAround<true>
	{
	public:

		static CStringW GetFolderPath(int nFolder)
		{
			TCHAR tszFolderPath[MAX_PATH];

			if (SHGetSpecialFolderPath(NULL, tszFolderPath, nFolder, false))
			{
				return CStringW(CT2W(tszFolderPath));
			}
			else
			{
				return CStringW();
			}
		}


		// Purpose: returns true if the wszFolder equals the 'all users' or user
		//          desktop folder. 
		static bool IsDesktopPath(const wchar_t* wszFolder)
		{
			return GetFolderPath(CSIDL_COMMON_DESKTOPDIRECTORY) == wszFolder ||
				GetFolderPath(CSIDL_DESKTOPDIRECTORY) == wszFolder;
		}


		void Initialize(const wchar_t* wszFolder)
		{
			m_bHideDesktopColumns = IsDesktopPath(wszFolder);
		}


		bool operator()() const throw()
		{
			return m_bHideDesktopColumns;
		}

	private:

		bool m_bHideDesktopColumns;
	};


	static HRESULT WINAPI UpdateRegistry(UINT nResId, BOOL bRegister, const wchar_t* wszDescription) throw()
	{
		return UpdateRegistryFromResource(nResId, bRegister, wszDescription, T::GetObjectCLSID());
	}


	IColumnProviderImpl() :
		m_pCachedInfo(NULL)
	{
		ATLTRACE2(atlTraceCOM, 0, _T("IColumnProviderImpl::IColumnProviderImpl (instance=%p)\n"), this);
	}


	~IColumnProviderImpl() throw()
	{
		ATLTRACE2(atlTraceCOM, 0, _T("IColumnProviderImpl::~IColumnProviderImpl (instance=%p)\n"), this);
	}


	// IColumnProvider
	STDMETHOD(Initialize)(const SHCOLUMNINIT* psci)
	{
		try
		{
			ATLTRACE2(atlTraceCOM, 0,
				_T("IColumnProviderImpl::Initialize, i=%p, tid=%d, dwFlags=%d, wszFolder=%s\n"),
				this, GetCurrentThreadId(), psci->dwFlags, psci->wszFolder);

			// Clear internal caching variables.
			m_pCachedInfo = NULL;
			m_strCachedFilename.Empty();
			m_mapCacheInfo.clear();

			m_desktopbugworkaround.Initialize(psci->wszFolder);

			m_columninfos.clear();
			m_columnidtoindex.clear();

			// Note: OnInitialize needs to be implemented by the derived class.
			static_cast<T*>(this)->OnInitialize(psci->wszFolder);

			return S_OK;
		}
		MSF_COM_CATCH_HANDLER()
	}


	// Purpose: GetColumnInfo is called by the shell to retrieve the column names.
	// By calling this repeatedly the shell can detect how many columns there are.
	STDMETHOD(GetColumnInfo)(DWORD dwIndex, SHCOLUMNINFO* psci)
	{
		ATLTRACE2(atlTraceCOM, 0,
			_T("IColumnProviderImpl::GetColumnInfo, i=%p, tid=%d, dwIndex=%d\n"),
			this, GetCurrentThreadId(), dwIndex);

		if (m_desktopbugworkaround() || dwIndex >= m_columninfos.size())
			return S_FALSE; // tell the shell there are no more columns.

		*psci = m_columninfos[dwIndex];
		return S_OK;
	}


	STDMETHOD(GetItemData)(const SHCOLUMNID* pscid, const SHCOLUMNDATA* pscd, VARIANT* pvarData)
	{
		ATLTRACE2(atlTraceCOM, 0, _T("IColumnProviderImpl::GetItemData, i=%p, tid=%d, f=%d, c=%d, file=%s\n"),
			this, GetCurrentThreadId(), pscd->dwFlags, pscid->pid, CW2T(pscd->wszFile));

		try
		{
			if (!IsSupportedItem(*pscd))
			{
				VariantInit(pvarData); // must initialise out args as we return a success code.
				return S_FALSE;
			}

			const wchar_t* wszFilename = PathFindFileNameW(pscd->wszFile);

			bool bFlushCache = IsBitSet(pscd->dwFlags, SHCDF_UPDATEITEM);
			const std::vector<CString>* pCachedInfo;

			if (bFlushCache)
			{
				pCachedInfo = NULL;
			}
			else
			{
				// User of this interface (explorer.exe) are expected to ask 
				// all active item data per file.
				pCachedInfo = FindInLastUsedCache(wszFilename);
			}

			if (pCachedInfo == NULL)
			{
				pCachedInfo = GetAndCacheFileInfo(wszFilename, pscd->wszFile, bFlushCache);
			}

			pvarData->bstrVal = (*pCachedInfo)[GetIndex(*pscid)].AllocSysString();
			pvarData->vt = VT_BSTR;

			return S_OK;
		}
		MSF_COM_CATCH_HANDLER()
	}

protected:

	// Purpose: provides a default GUID for the columns.
	const GUID& GetStandardFormatIdentifier() const
	{
		return T::GetObjectCLSID();
	}
	

	void RegisterColumn(const GUID& fmtid, DWORD pid, const wchar_t* wszTitle, UINT cChars, DWORD fmt = LVCFMT_LEFT, DWORD csFlags = SHCOLSTATE_TYPE_STR)
	{
		SHCOLUMNID columnid;

		columnid.pid   = pid;
		columnid.fmtid = fmtid;

		SHCOLUMNINFO columninfo;

		columninfo.scid = columnid;

		ATLASSERT(wcslen(wszTitle) < MAX_COLUMN_NAME_LEN && "wszTitle is too long");
		wcscpy(columninfo.wszTitle, wszTitle);
		columninfo.cChars = cChars;
		columninfo.fmt = fmt;
		columninfo.csFlags = csFlags | SHCOLSTATE_EXTENDED | SHCOLSTATE_SECONDARYUI;
		// Note: VT_LPSTR/VT_BSTR works ok. Other types seems to have issues with sorting.
		columninfo.vt = VT_BSTR;
		columninfo.wszDescription[0] = 0; // not used by the shell.

		m_columninfos.push_back(columninfo);
		m_columnidtoindex[columnid] = static_cast<unsigned int>(m_columninfos.size() - 1);
	}


	void RegisterColumn(const GUID& fmtid, DWORD pid, UINT nResourceID, UINT cChars, DWORD fmt = LVCFMT_LEFT, DWORD csFlags = SHCOLSTATE_TYPE_STR)
	{
		RegisterColumn(pid, fmtid, LoadStringW(nResourceID), cChars, fmt, csFlags);
	}


	void RegisterColumn(const wchar_t* wszTitle, UINT cChars, DWORD fmt = LVCFMT_LEFT, DWORD csFlags = SHCOLSTATE_TYPE_STR)
	{
		RegisterColumn(GetStandardFormatIdentifier(), static_cast<DWORD>(m_columninfos.size()),
			wszTitle, cChars, fmt, csFlags);
	}


	void RegisterColumn(UINT nResourceID, UINT cChars, DWORD fmt = LVCFMT_LEFT, DWORD csFlags = SHCOLSTATE_TYPE_STR)
	{
		RegisterColumn(LoadStringW(nResourceID), cChars, fmt, csFlags);
	}


	void RegisterExtension(CStringW strExtension)
	{
		m_extensions.push_back(strExtension.MakeLower());
	}


	bool IsSupportedItem(const SHCOLUMNDATA& scd) const
	{
		// Only support online files.
		if (scd.dwFileAttributes & static_cast<const T*>(this)->GetFileAttributeMask())
			return false;

		// Check file extension.
		return IsSupportedFileType(scd.pwszExt);
	}


	bool IsSupportedFileType(CStringW strExtension) const
	{
		strExtension.MakeLower();

		std::vector<CStringW>::const_iterator result =
			std::find(m_extensions.begin(), m_extensions.end(), strExtension);

		return result != m_extensions.end();
	}


	unsigned int GetIndex(const SHCOLUMNID& scid) const
	{
		mapColumnIdToIndex::const_iterator it = m_columnidtoindex.find(scid);
		RaiseExceptionIf(it == m_columnidtoindex.end());

		return it->second;
	}

	
	DWORD GetFileAttributeMask() const
	{
		return FILE_ATTRIBUTE_DIRECTORY | FILE_ATTRIBUTE_OFFLINE;
	}


private:

	const std::vector<CString>* GetAndCacheFileInfo(const CStringW& strFilename, const wchar_t* wszFile, bool bFlushCache)
	{
		const std::vector<CString>* pCachedInfo;

		if (bFlushCache)
		{
			pCachedInfo = NULL;
		}
		else
		{
			pCachedInfo = FindInCache(strFilename);
		}

		if (pCachedInfo == NULL)
		{
			// New file info must be stored in the cache.
			vector<CString> strColumnInfos;

			// Note: GetAllColumnInfo must be implemented by the derived class.
			static_cast<T*>(this)->GetAllColumnInfo(CString(CW2CT(wszFile)), strColumnInfos);

			ATLASSERT(strColumnInfos.size() == m_columninfos.size() && "Missing info");

			pCachedInfo = InsertInCache(strFilename, strColumnInfos);
		}

		return InsertInLastUsedCache(strFilename, pCachedInfo);
	}


	const std::vector<CString>* FindInCache(const CStringW& strFilename) const
	{
		mapCacheInfo::const_iterator it = m_mapCacheInfo.find(strFilename);
		if (it != m_mapCacheInfo.end())
		{
			return &(*it).second;
		}
		else
		{
			return NULL;
		}
	}


	const std::vector<CString>* InsertInCache(const CStringW& strFilename, const std::vector<CString>& strColumnInfos)
	{
		return &(m_mapCacheInfo[strFilename] = strColumnInfos);
	}


	const std::vector<CString>* InsertInLastUsedCache(const CStringW& strFilename, const std::vector<CString>* pCachedInfo)
	{
		m_strCachedFilename = strFilename;
		m_pCachedInfo = pCachedInfo;

		return m_pCachedInfo;
	}


	const std::vector<CString>* FindInLastUsedCache(const wchar_t* wszFilename) const
	{
		if (m_pCachedInfo == NULL)
			return NULL; // last used cache is empty.

		if (m_strCachedFilename != wszFilename)
			return NULL; // last used cache was for a different file.
		
		return m_pCachedInfo;
	}

	// Member variables.
	std::vector<SHCOLUMNINFO>   m_columninfos;
	std::vector<CStringW>       m_extensions;
	mapColumnIdToIndex          m_columnidtoindex;

	CStringW                    m_strCachedFilename;
	const std::vector<CString>* m_pCachedInfo;
	mapCacheInfo                m_mapCacheInfo;

	CDesktopBugWorkAround<t_bEnableDesktopBugWorkAround> m_desktopbugworkaround;
};


} // namespace MSF