1*cdf0e10cSrcweir #undef UNICODE
2*cdf0e10cSrcweir #undef _UNICODE
3*cdf0e10cSrcweir 
4*cdf0e10cSrcweir #pragma once
5*cdf0e10cSrcweir 
6*cdf0e10cSrcweir #ifdef _MSC_VER
7*cdf0e10cSrcweir #pragma warning(push, 1) /* disable warnings within system headers */
8*cdf0e10cSrcweir #endif
9*cdf0e10cSrcweir #define WIN32_LEAN_AND_MEAN
10*cdf0e10cSrcweir #include <windows.h>
11*cdf0e10cSrcweir #include <msiquery.h>
12*cdf0e10cSrcweir #include <imagehlp.h>
13*cdf0e10cSrcweir #include <tchar.h>
14*cdf0e10cSrcweir #include <strsafe.h>
15*cdf0e10cSrcweir #ifdef _MSC_VER
16*cdf0e10cSrcweir #pragma warning(pop)
17*cdf0e10cSrcweir #endif
18*cdf0e10cSrcweir 
19*cdf0e10cSrcweir #include <malloc.h>
20*cdf0e10cSrcweir #include <time.h>
21*cdf0e10cSrcweir #include <string>
22*cdf0e10cSrcweir #include <hash_map>
23*cdf0e10cSrcweir 
24*cdf0e10cSrcweir const DWORD PE_Signature = 0x00004550;
25*cdf0e10cSrcweir typedef std::pair< std::string, bool > StringPair;
26*cdf0e10cSrcweir typedef	std::hash_map< std::string, bool > ExcludeLibsMap;
27*cdf0e10cSrcweir 
28*cdf0e10cSrcweir #ifdef DEBUG
29*cdf0e10cSrcweir static void OutputDebugStringFormat( LPCSTR pFormat, ... )
30*cdf0e10cSrcweir {
31*cdf0e10cSrcweir 	CHAR    buffer[1024];
32*cdf0e10cSrcweir 	va_list args;
33*cdf0e10cSrcweir 
34*cdf0e10cSrcweir 	va_start( args, pFormat );
35*cdf0e10cSrcweir 	StringCchVPrintfA( buffer, sizeof(buffer), pFormat, args );
36*cdf0e10cSrcweir 	OutputDebugStringA( buffer );
37*cdf0e10cSrcweir }
38*cdf0e10cSrcweir #else
39*cdf0e10cSrcweir static void OutputDebugStringFormat( LPCSTR, ... )
40*cdf0e10cSrcweir {
41*cdf0e10cSrcweir }
42*cdf0e10cSrcweir #endif
43*cdf0e10cSrcweir 
44*cdf0e10cSrcweir static bool IsValidHandle( HANDLE handle )
45*cdf0e10cSrcweir {
46*cdf0e10cSrcweir 	return NULL != handle && INVALID_HANDLE_VALUE != handle;
47*cdf0e10cSrcweir }
48*cdf0e10cSrcweir 
49*cdf0e10cSrcweir static std::string GetMsiProperty(MSIHANDLE handle, const std::string& sProperty)
50*cdf0e10cSrcweir {
51*cdf0e10cSrcweir 	std::string result;
52*cdf0e10cSrcweir     TCHAR		szDummy[1] = TEXT("");
53*cdf0e10cSrcweir     DWORD		nChars = 0;
54*cdf0e10cSrcweir 
55*cdf0e10cSrcweir     if (MsiGetProperty(handle, sProperty.c_str(), szDummy, &nChars) == ERROR_MORE_DATA)
56*cdf0e10cSrcweir     {
57*cdf0e10cSrcweir         DWORD nBytes = ++nChars * sizeof(TCHAR);
58*cdf0e10cSrcweir         LPTSTR buffer = reinterpret_cast<LPTSTR>(_alloca(nBytes));
59*cdf0e10cSrcweir         ZeroMemory( buffer, nBytes );
60*cdf0e10cSrcweir         MsiGetProperty(handle, sProperty.c_str(), buffer, &nChars);
61*cdf0e10cSrcweir         result = buffer;
62*cdf0e10cSrcweir     }
63*cdf0e10cSrcweir     return result;
64*cdf0e10cSrcweir }
65*cdf0e10cSrcweir 
66*cdf0e10cSrcweir static BOOL rebaseImage( const std::string& filePath, LPVOID address )
67*cdf0e10cSrcweir {
68*cdf0e10cSrcweir 	ULONG ulOldImageSize;
69*cdf0e10cSrcweir 	ULONG_PTR lpOldImageBase;
70*cdf0e10cSrcweir 	ULONG ulNewImageSize;
71*cdf0e10cSrcweir 	ULONG_PTR lpNewImageBase = reinterpret_cast<ULONG_PTR>(address);
72*cdf0e10cSrcweir 
73*cdf0e10cSrcweir 	BOOL bResult = ReBaseImage(
74*cdf0e10cSrcweir 		filePath.c_str(),
75*cdf0e10cSrcweir 		"",
76*cdf0e10cSrcweir 		TRUE,
77*cdf0e10cSrcweir 		FALSE,
78*cdf0e10cSrcweir 		FALSE,
79*cdf0e10cSrcweir 		0,
80*cdf0e10cSrcweir 		&ulOldImageSize,
81*cdf0e10cSrcweir 		&lpOldImageBase,
82*cdf0e10cSrcweir 		&ulNewImageSize,
83*cdf0e10cSrcweir 		&lpNewImageBase,
84*cdf0e10cSrcweir 		(ULONG)time(NULL) );
85*cdf0e10cSrcweir 
86*cdf0e10cSrcweir 	return bResult;
87*cdf0e10cSrcweir }
88*cdf0e10cSrcweir 
89*cdf0e10cSrcweir static BOOL rebaseImage( MSIHANDLE /*handle*/, const std::string& sFilePath, LPVOID address )
90*cdf0e10cSrcweir {
91*cdf0e10cSrcweir 	std::string	mystr;
92*cdf0e10cSrcweir 	mystr = "Full file: " + sFilePath;
93*cdf0e10cSrcweir 
94*cdf0e10cSrcweir 	BOOL bResult = rebaseImage( sFilePath, address );
95*cdf0e10cSrcweir 
96*cdf0e10cSrcweir 	if ( !bResult )
97*cdf0e10cSrcweir 	{
98*cdf0e10cSrcweir 		OutputDebugStringFormat( "Rebasing library %s failed", mystr.c_str() );
99*cdf0e10cSrcweir 	}
100*cdf0e10cSrcweir 
101*cdf0e10cSrcweir 	return bResult;
102*cdf0e10cSrcweir }
103*cdf0e10cSrcweir 
104*cdf0e10cSrcweir static BOOL rebaseImagesInFolder( MSIHANDLE handle, const std::string& sPath, LPVOID address, ExcludeLibsMap& rExcludeMap )
105*cdf0e10cSrcweir {
106*cdf0e10cSrcweir 	std::string     sDir     = sPath;
107*cdf0e10cSrcweir 	std::string	    sPattern = sPath + TEXT("*.dll");
108*cdf0e10cSrcweir 	WIN32_FIND_DATA	aFindFileData;
109*cdf0e10cSrcweir 
110*cdf0e10cSrcweir 	HANDLE hFind = FindFirstFile( sPattern.c_str(), &aFindFileData );
111*cdf0e10cSrcweir 	if ( IsValidHandle(hFind) )
112*cdf0e10cSrcweir 	{
113*cdf0e10cSrcweir 		BOOL fSuccess = false;
114*cdf0e10cSrcweir 
115*cdf0e10cSrcweir 		do
116*cdf0e10cSrcweir 		{
117*cdf0e10cSrcweir 			std::string sFileName = aFindFileData.cFileName;
118*cdf0e10cSrcweir 			if ( rExcludeMap.find( sFileName ) == rExcludeMap.end() )
119*cdf0e10cSrcweir 			{
120*cdf0e10cSrcweir 				OutputDebugStringFormat( "Rebase library: %s", sFileName.c_str() );
121*cdf0e10cSrcweir                 std::string	sLibFile = sDir +  sFileName;
122*cdf0e10cSrcweir                 rebaseImage( handle, sLibFile, address );
123*cdf0e10cSrcweir 			}
124*cdf0e10cSrcweir 			else
125*cdf0e10cSrcweir 			{
126*cdf0e10cSrcweir 				OutputDebugStringFormat( "Exclude library %s from rebase", sFileName.c_str() );
127*cdf0e10cSrcweir 			}
128*cdf0e10cSrcweir 
129*cdf0e10cSrcweir 			fSuccess = FindNextFile( hFind, &aFindFileData );
130*cdf0e10cSrcweir 		}
131*cdf0e10cSrcweir 		while ( fSuccess );
132*cdf0e10cSrcweir 
133*cdf0e10cSrcweir 		FindClose( hFind );
134*cdf0e10cSrcweir 	}
135*cdf0e10cSrcweir 
136*cdf0e10cSrcweir 	return ERROR_SUCCESS;
137*cdf0e10cSrcweir }
138*cdf0e10cSrcweir 
139*cdf0e10cSrcweir static BOOL rebaseImages( MSIHANDLE handle, LPVOID pAddress, ExcludeLibsMap& rMap )
140*cdf0e10cSrcweir {
141*cdf0e10cSrcweir 	std::string sInstallPath = GetMsiProperty(handle, TEXT("INSTALLLOCATION"));
142*cdf0e10cSrcweir 
143*cdf0e10cSrcweir 	std::string sBasisDir  = sInstallPath + TEXT("Basis\\program\\");
144*cdf0e10cSrcweir 	std::string sOfficeDir = sInstallPath + TEXT("program\\");
145*cdf0e10cSrcweir 	std::string sUreDir    = sInstallPath + TEXT("URE\\bin\\");
146*cdf0e10cSrcweir 
147*cdf0e10cSrcweir 	BOOL bResult = rebaseImagesInFolder( handle, sBasisDir, pAddress, rMap );
148*cdf0e10cSrcweir 	bResult &= rebaseImagesInFolder( handle, sOfficeDir, pAddress, rMap );
149*cdf0e10cSrcweir 	bResult &= rebaseImagesInFolder( handle, sUreDir, pAddress, rMap );
150*cdf0e10cSrcweir 
151*cdf0e10cSrcweir 	return bResult;
152*cdf0e10cSrcweir }
153*cdf0e10cSrcweir 
154*cdf0e10cSrcweir static BOOL IsServerSystem( MSIHANDLE /*handle*/ )
155*cdf0e10cSrcweir {
156*cdf0e10cSrcweir 	OSVERSIONINFOEX osVersionInfoEx;
157*cdf0e10cSrcweir 	osVersionInfoEx.dwOSVersionInfoSize = sizeof(OSVERSIONINFOEX);
158*cdf0e10cSrcweir 	GetVersionEx(reinterpret_cast<LPOSVERSIONINFO>(&osVersionInfoEx));
159*cdf0e10cSrcweir 
160*cdf0e10cSrcweir 	if ( osVersionInfoEx.wProductType != VER_NT_WORKSTATION )
161*cdf0e10cSrcweir 	{
162*cdf0e10cSrcweir         OutputDebugStringFormat( "Server system detected. No rebase necessary!" );
163*cdf0e10cSrcweir 		return TRUE;
164*cdf0e10cSrcweir 	}
165*cdf0e10cSrcweir 	else
166*cdf0e10cSrcweir     {
167*cdf0e10cSrcweir         OutputDebugStringFormat( "Client system detected. Rebase necessary!" );
168*cdf0e10cSrcweir         return FALSE;
169*cdf0e10cSrcweir     }
170*cdf0e10cSrcweir }
171*cdf0e10cSrcweir 
172*cdf0e10cSrcweir static void InitExcludeFromRebaseList( MSIHANDLE handle, ExcludeLibsMap& rMap )
173*cdf0e10cSrcweir {
174*cdf0e10cSrcweir 	size_t      nPos( 0 );
175*cdf0e10cSrcweir     const TCHAR cDelim = ',';
176*cdf0e10cSrcweir 	std::string sLibsExcluded = GetMsiProperty(handle, TEXT("EXCLUDE_FROM_REBASE"));
177*cdf0e10cSrcweir 
178*cdf0e10cSrcweir     while ( nPos < sLibsExcluded.size() )
179*cdf0e10cSrcweir 	{
180*cdf0e10cSrcweir 	    size_t nDelPos = sLibsExcluded.find_first_of( cDelim, nPos );
181*cdf0e10cSrcweir 
182*cdf0e10cSrcweir 		std::string sExcludedLibName;
183*cdf0e10cSrcweir 		if ( nDelPos != std::string::npos )
184*cdf0e10cSrcweir 		{
185*cdf0e10cSrcweir 			sExcludedLibName = sLibsExcluded.substr( nPos, nDelPos - nPos );
186*cdf0e10cSrcweir 		    nPos = nDelPos+1;
187*cdf0e10cSrcweir 		}
188*cdf0e10cSrcweir 		else
189*cdf0e10cSrcweir 		{
190*cdf0e10cSrcweir 			sExcludedLibName = sLibsExcluded.substr( nPos );
191*cdf0e10cSrcweir 			nPos = sLibsExcluded.size();
192*cdf0e10cSrcweir 		}
193*cdf0e10cSrcweir 
194*cdf0e10cSrcweir 		if ( sExcludedLibName.size() > 0 )
195*cdf0e10cSrcweir 		{
196*cdf0e10cSrcweir 			OutputDebugStringFormat( "Insert library %s into exclude from rebase list", sExcludedLibName.c_str() );
197*cdf0e10cSrcweir 			rMap.insert( StringPair( sExcludedLibName, true ));
198*cdf0e10cSrcweir 		}
199*cdf0e10cSrcweir 	}
200*cdf0e10cSrcweir }
201*cdf0e10cSrcweir 
202*cdf0e10cSrcweir extern "C" BOOL __stdcall RebaseLibrariesOnProperties( MSIHANDLE handle )
203*cdf0e10cSrcweir {
204*cdf0e10cSrcweir 	static LPVOID pDefault = reinterpret_cast<LPVOID>(0x10000000);
205*cdf0e10cSrcweir 
206*cdf0e10cSrcweir 	OutputDebugStringFormat( "RebaseLibrariesOnProperties has been called" );
207*cdf0e10cSrcweir 	std::string sDontOptimizeLibs = GetMsiProperty(handle, TEXT("DONTOPTIMIZELIBS"));
208*cdf0e10cSrcweir 	if ( sDontOptimizeLibs.length() > 0 && sDontOptimizeLibs == "1" )
209*cdf0e10cSrcweir 	{
210*cdf0e10cSrcweir         OutputDebugStringFormat( "Don't optimize libraries set. No rebase necessary!" );
211*cdf0e10cSrcweir 		return TRUE;
212*cdf0e10cSrcweir 	}
213*cdf0e10cSrcweir 
214*cdf0e10cSrcweir 	if ( !IsServerSystem( handle ))
215*cdf0e10cSrcweir 	{
216*cdf0e10cSrcweir 		ExcludeLibsMap aExcludeLibsMap;
217*cdf0e10cSrcweir 		InitExcludeFromRebaseList( handle, aExcludeLibsMap );
218*cdf0e10cSrcweir 
219*cdf0e10cSrcweir 		return rebaseImages( handle, pDefault, aExcludeLibsMap );
220*cdf0e10cSrcweir 	}
221*cdf0e10cSrcweir 
222*cdf0e10cSrcweir 	return TRUE;
223*cdf0e10cSrcweir }
224