OLD | NEW |
| (Empty) |
1 // Copyright 2005-2009 Google Inc. | |
2 // | |
3 // Licensed under the Apache License, Version 2.0 (the "License"); | |
4 // you may not use this file except in compliance with the License. | |
5 // You may obtain a copy of the License at | |
6 // | |
7 // http://www.apache.org/licenses/LICENSE-2.0 | |
8 // | |
9 // Unless required by applicable law or agreed to in writing, software | |
10 // distributed under the License is distributed on an "AS IS" BASIS, | |
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
12 // See the License for the specific language governing permissions and | |
13 // limitations under the License. | |
14 // ======================================================================== | |
15 // | |
16 // Utility functions related to PE files (executables) | |
17 | |
18 #include "omaha/base/pe_utils.h" | |
19 | |
20 #include "omaha/base/debug.h" | |
21 #include "omaha/base/error.h" | |
22 #include "omaha/base/scoped_any.h" | |
23 | |
24 namespace omaha { | |
25 | |
26 namespace { | |
27 | |
28 // Not really necessary as long as running on x86 architecture throughout, but | |
29 // what the hell | |
30 #if defined(BIG_ENDIAN) | |
31 uint32 GetUint32LE(void const * const p) { | |
32 uint8 const * const pu = reinterpret_cast<uint8 const * const>(p); | |
33 uint32 i = pu[0] | pu[1]<<8 | pu[2]<<16 | pu[3]<<24; | |
34 return i; | |
35 } | |
36 | |
37 void PutUint32LE(uint32 i, void * const p) { | |
38 uint8 * const pu = reinterpret_cast<uint8 * const>(p); | |
39 pu[0] = i & 0xff; | |
40 pu[1] = (i >> 8) & 0xff; | |
41 pu[2] = (i >> 16) & 0xff; | |
42 pu[3] = (i >> 24) & 0xff; | |
43 } | |
44 #else // LITTLE_ENDIAN | |
45 inline uint32 GetUint32LE(void const * const p) { | |
46 uint32 const * const pu = reinterpret_cast<uint32 const * const>(p); | |
47 return *pu; | |
48 } | |
49 | |
50 inline void PutUint32LE(uint32 i, void * const p) { | |
51 uint32 * const pu = reinterpret_cast<uint32 * const>(p); | |
52 *pu = i; | |
53 } | |
54 #endif | |
55 | |
56 // Magic PE constants | |
57 const uint32 kPEHeaderOffset = 60; | |
58 const uint32 kPEHeaderChecksumOffset = 88; | |
59 const uint32 kPEHeaderSizeMin = 160; | |
60 const char magic_EXE_header[] = "MZ\0\0"; | |
61 const char magic_PE_header[] = "PE\0\0"; | |
62 | |
63 } // namespace | |
64 | |
65 HRESULT SetPEChecksum(const TCHAR *filename, uint32 checksum) { | |
66 // Write the checksum field of the Windows NT-specific "optional" header. | |
67 // Use Windows API calls rather than C library calls so that it will be | |
68 // really a small routine when used in the stub executable. | |
69 | |
70 ASSERT(filename, (L"")); | |
71 | |
72 scoped_hfile file(::CreateFile(filename, | |
73 GENERIC_READ | GENERIC_WRITE, | |
74 0, | |
75 NULL, | |
76 OPEN_EXISTING, | |
77 FILE_ATTRIBUTE_NORMAL, | |
78 NULL)); | |
79 if (!file) | |
80 return HRESULTFromLastError(); | |
81 | |
82 size_t size = ::GetFileSize(get(file), NULL); | |
83 if (size == INVALID_FILE_SIZE) | |
84 return HRESULTFromLastError(); | |
85 | |
86 scoped_file_mapping mapping(::CreateFileMapping(get(file), | |
87 NULL, | |
88 PAGE_READWRITE, | |
89 0, | |
90 0, | |
91 NULL)); | |
92 if (!mapping) | |
93 return HRESULTFromLastError(); | |
94 | |
95 scoped_file_view file_data(::MapViewOfFile(get(mapping), | |
96 FILE_MAP_WRITE, | |
97 0, | |
98 0, | |
99 size)); | |
100 if (!file_data) | |
101 return HRESULTFromLastError(); | |
102 | |
103 uint8 * image = reinterpret_cast<uint8 *>(get(file_data)); | |
104 | |
105 return SetPEChecksumToBuffer(image, size, checksum); | |
106 } | |
107 | |
108 HRESULT GetPEChecksum(const TCHAR *filename, uint32 * checksum) { | |
109 // Read the checksum field out of the Windows NT-specific "optional" header. | |
110 // Use Windows API calls rather than C library calls so that it will be | |
111 // really a small routine when used in the stub executable. | |
112 | |
113 ASSERT(filename, (L"")); | |
114 ASSERT(checksum, (L"")); | |
115 | |
116 scoped_hfile file(::CreateFile(filename, | |
117 GENERIC_READ, | |
118 FILE_SHARE_READ, | |
119 NULL, | |
120 OPEN_EXISTING, | |
121 FILE_ATTRIBUTE_READONLY, | |
122 NULL)); | |
123 if (!file) | |
124 return HRESULTFromLastError(); | |
125 | |
126 size_t size = ::GetFileSize(get(file), NULL); | |
127 if (size == INVALID_FILE_SIZE) | |
128 return HRESULTFromLastError(); | |
129 | |
130 scoped_file_mapping mapping(::CreateFileMapping(get(file), | |
131 NULL, | |
132 PAGE_READONLY, | |
133 0, | |
134 0, | |
135 NULL)); | |
136 if (!mapping) | |
137 return HRESULTFromLastError(); | |
138 | |
139 scoped_file_view file_data(::MapViewOfFile(get(mapping), | |
140 FILE_MAP_READ, | |
141 0, | |
142 0, | |
143 size)); | |
144 if (!file_data) | |
145 return HRESULTFromLastError(); | |
146 | |
147 uint8 * image = reinterpret_cast<uint8 *>(get(file_data)); | |
148 | |
149 return GetPEChecksumFromBuffer(image, size, checksum); | |
150 } | |
151 | |
152 HRESULT SetPEChecksumToBuffer(uint8 *buffer, size_t size, uint32 checksum) { | |
153 // Sanity checks | |
154 if (size < 64) { | |
155 ASSERT(false, (L"File too short to be valid executable")); | |
156 return E_FAIL; | |
157 } | |
158 | |
159 uint32 x = GetUint32LE(magic_EXE_header); | |
160 if (::memcmp(buffer, &x, 2)) { | |
161 ASSERT(false, (L"Missing executable's magic number")); | |
162 return E_FAIL; | |
163 } | |
164 | |
165 uint32 peheader = GetUint32LE(buffer + kPEHeaderOffset); | |
166 if (size < peheader + kPEHeaderSizeMin) { | |
167 ASSERT(false, (L"Too small given PE header size")); | |
168 return E_FAIL; | |
169 } | |
170 | |
171 x = GetUint32LE(magic_PE_header); | |
172 if (::memcmp(buffer + peheader, &x, 4)) { | |
173 ASSERT(false, (L"Missing PE header magic number")); | |
174 return E_FAIL; | |
175 } | |
176 | |
177 // Finally, write the checksum | |
178 PutUint32LE(checksum, &x); | |
179 ::memcpy(buffer + peheader + kPEHeaderChecksumOffset, &x, 4); | |
180 | |
181 return S_OK; | |
182 } | |
183 | |
184 HRESULT GetPEChecksumFromBuffer(const unsigned char *buffer, | |
185 size_t size, | |
186 uint32 *checksum) { | |
187 // Sanity checks | |
188 if (size < 64) { | |
189 ASSERT(false, (L"File too short to be valid executable")); | |
190 return E_FAIL; | |
191 } | |
192 | |
193 uint32 x = GetUint32LE(magic_EXE_header); | |
194 if (::memcmp(buffer, &x, 2)) { | |
195 ASSERT(false, (L"Missing executable's magic number")); | |
196 return E_FAIL; | |
197 } | |
198 | |
199 uint32 peheader = GetUint32LE(buffer + kPEHeaderOffset); | |
200 if (size < peheader + kPEHeaderSizeMin) { | |
201 ASSERT(false, (L"Too small given PE header size")); | |
202 return E_FAIL; | |
203 } | |
204 | |
205 x = GetUint32LE(magic_PE_header); | |
206 if (::memcmp(buffer + peheader, &x, 4)) { | |
207 ASSERT(false, (L"Missing PE header magic number")); | |
208 return E_FAIL; | |
209 } | |
210 | |
211 // Finally, read the checksum | |
212 | |
213 *checksum = GetUint32LE(buffer + peheader + kPEHeaderChecksumOffset); | |
214 | |
215 return S_OK; | |
216 } | |
217 | |
218 } // namespace omaha | |
219 | |
OLD | NEW |