1// Licensed to the .NET Foundation under one or more agreements.
2// The .NET Foundation licenses this file to you under the MIT license.
3// See the LICENSE file in the project root for more information.
4#ifndef PEWriter_H
5#define PEWriter_H
6
7#include <crtwrap.h>
8
9#include <windows.h>
10
11#include "ceegen.h"
12
13#include "pesectionman.h"
14
15class PEWriter;
16class PEWriterSection;
17class PEDecoder;
18struct entry;
19struct _IMAGE_SECTION_HEADER;
20
21#define SUBSECTION_ALIGN 16
22
23/***************************************************************/
24// PEWriter is derived from PESectionManager. While the base class just
25// manages the sections, PEWriter can actually write them out.
26
27class PEWriter : public PESectionMan
28{
29public:
30
31 // See ICeeFileGen.h for definition of createFlags
32 HRESULT Init(PESectionMan *pFrom, DWORD createFlags, LPCWSTR seedFileName = NULL);
33 HRESULT Cleanup();
34
35 // Finds section with given name. returns 0 if not found
36 virtual PESection* getSection(const char* name);
37
38 // Create a new section
39 virtual HRESULT newSection(const char* name, PESection **section,
40 unsigned flags=sdNone, unsigned estSize=0x10000,
41 unsigned estRelocs=0x100);
42
43 HRESULT link();
44 HRESULT fixup(CeeGenTokenMapper *pMapper);
45 HRESULT write(__in LPCWSTR fileName);
46 HRESULT write(void **ppImage);
47
48 // calling these functions is optional
49 DWORD getSectionAlignment();
50 void setSectionAlignment(DWORD);
51 DWORD getFileAlignment();
52 void setFileAlignment(DWORD);
53 DWORD getImageBase32();
54 void setImageBase32(DWORD);
55 UINT64 getImageBase64();
56 void setImageBase64(UINT64);
57 void stripRelocations(bool val); // default = false
58
59 void getHeaderInfo(PIMAGE_NT_HEADERS *ppNtHeaders, PIMAGE_SECTION_HEADER *ppSections, ULONG *pNumSections);
60
61 // these affect the charactertics in the NT file header
62 void setCharacteristics(unsigned mask);
63 void clearCharacteristics(unsigned mask);
64
65 // these affect the charactertics in the NT optional header
66 void setDllCharacteristics(unsigned mask);
67 void clearDllCharacteristics(unsigned mask);
68
69 // sets the SubSystem field in the optional header
70 void setSubsystem(unsigned subsystem, unsigned major, unsigned minor);
71
72 // specify the entry point as an offset into the text section. The
73 // method will convert into an RVA from the base
74 void setEntryPointTextOffset(unsigned entryPoint);
75
76 HRESULT setDirectoryEntry(PEWriterSection *section, ULONG entry, ULONG size, ULONG offset=0);
77
78
79 // get the RVA for the IL section
80 ULONG getIlRva();
81
82 // set the RVA for the IL section by supplying offset to the IL section
83 void setIlRva(DWORD offset);
84
85 unsigned getSubsystem();
86
87 size_t getImageBase();
88
89 void setEnCRvaBase(ULONG dataBase, ULONG rdataBase);
90
91 HRESULT getFileTimeStamp(DWORD *pTimeStamp);
92
93 IMAGE_NT_HEADERS32* ntHeaders32() { return (IMAGE_NT_HEADERS32*) m_ntHeaders; }
94 IMAGE_NT_HEADERS64* ntHeaders64() { return (IMAGE_NT_HEADERS64*) m_ntHeaders; }
95
96 bool isPE32() // true -> created a PE 32-bit PE file
97 // false -> created a PE+ 64-bit PE file
98 { return (m_ntHeaders->OptionalHeader.Magic == VAL16(IMAGE_NT_OPTIONAL_HDR32_MAGIC)); }
99
100 bool isI386() // true -> target machine is i386
101 { return (m_ntHeaders->FileHeader.Machine == VAL16(IMAGE_FILE_MACHINE_I386)); }
102
103 bool isIA64() // true -> target machine is ia64
104 { return (m_ntHeaders->FileHeader.Machine == VAL16(IMAGE_FILE_MACHINE_IA64)); }
105
106 bool isAMD64() // true -> target machine is ia64
107 { return (m_ntHeaders->FileHeader.Machine == VAL16(IMAGE_FILE_MACHINE_AMD64)); }
108
109 bool isARM() // true -> target machine is ARM
110 { return (m_ntHeaders->FileHeader.Machine == VAL16(IMAGE_FILE_MACHINE_ARMNT)); }
111
112 bool createCorMainStub() // do we need the CorExeMain/CorDllMain stubs?
113 { return m_createCorMainStub; }
114
115private:
116 DWORD m_ilRVA;
117 BOOL m_encMode;
118 ULONG m_dataRvaBase;
119 ULONG m_rdataRvaBase;
120 ULONG m_codeRvaBase;
121 DWORD m_peFileTimeStamp;
122
123 HANDLE m_file;
124
125 // "Seed" file information. The new file data will be "appended" to the seed file
126 // These are valid only if m_hSeedFile is valid
127
128 HANDLE m_hSeedFile;
129 HANDLE m_hSeedFileMap;
130 PEDecoder * m_pSeedFileDecoder;
131 IMAGE_NT_HEADERS * m_pSeedFileNTHeaders;
132 unsigned m_iSeedSections;
133 IMAGE_SECTION_HEADER*m_pSeedSections;
134 IMAGE_SECTION_HEADER*m_pSeedSectionToAdd; // used only by newSection()
135
136 PEWriterSection **getSectStart() {
137 return (PEWriterSection**)sectStart;
138 }
139
140 PEWriterSection **getSectCur() {
141 return (PEWriterSection**)sectCur;
142 }
143
144 COUNT_T getSectCount() {
145 return COUNT_T(sectCur - sectStart);
146 }
147
148
149 IMAGE_DOS_HEADER m_dosHeader;
150 IMAGE_NT_HEADERS * m_ntHeaders;
151 DWORD m_ntHeadersSize; // Size of IMAGE_NT_HEADERS (not including section headers)
152
153 unsigned virtualPos;
154 unsigned filePos;
155
156 PEWriterSection * reloc;
157 PEWriterSection * strtab;
158
159 IMAGE_SECTION_HEADER *headers, *headersEnd;
160
161 struct directoryEntry {
162 PEWriterSection *section;
163 ULONG offset;
164 ULONG size;
165 };
166
167 // Directory entries in the file header
168 directoryEntry * pEntries;
169 USHORT cEntries;
170
171 bool m_createCorMainStub;
172
173 // Helpers for link()
174 HRESULT linkSortSections(entry * entries,
175 unsigned * piEntries, // OUT
176 unsigned * piUniqueSections); // OUT
177 HRESULT linkSortHeaders(entry * entries, unsigned iEntries, unsigned iUniqueSections);
178 HRESULT linkPlaceSections(entry * entries, unsigned iEntries);
179 void setSectionIndex(IMAGE_SECTION_HEADER * h, unsigned sectionIndex);
180
181 HRESULT Open(__in LPCWSTR fileName);
182 HRESULT Write(const void *data, int size);
183 HRESULT Seek(int offset);
184 HRESULT Pad(int align);
185 HRESULT Close();
186};
187
188// This class encapsulates emitting the base relocs into the
189// .reloc section of the PE file, for the case where the image
190// does not get loaded at its preferred base address.
191
192class PERelocSection
193{
194 private:
195 PEWriterSection * section;
196 unsigned relocPage;
197 unsigned relocSize;
198 DWORD * relocSizeAddr;
199 unsigned pages;
200
201#ifdef _DEBUG
202 unsigned lastRVA;
203#endif
204
205 public:
206 PERelocSection(PEWriterSection *pBaseReloc);
207 void AddBaseReloc(unsigned rva, int type, unsigned short highAdj=0);
208 void Finish(bool isPE32);
209};
210
211class PEWriterSection : public PESection {
212
213public:
214
215 PEWriterSection(const char* name, unsigned flags,
216 unsigned estSize, unsigned estRelocs)
217 : PESection(name, flags, estSize, estRelocs) {}
218
219 virtual HRESULT applyRelocs(IMAGE_NT_HEADERS * pNtHeaders,
220 PERelocSection * relocSection,
221 CeeGenTokenMapper * pTokenMapper,
222 DWORD rdataRvaBase,
223 DWORD dataRvaBase,
224 DWORD textRvaBase);
225
226 virtual HRESULT write (HANDLE file);
227 virtual unsigned writeMem (void ** pMem);
228 virtual bool isSeedSection() { return false; }
229
230};
231
232// This is for sections from the seed file. Their order needs to be maintained and
233// they need to be written to the output file.
234
235class PESeedSection : public PEWriterSection {
236
237public:
238
239 PESeedSection(PEDecoder * peDecoder, IMAGE_SECTION_HEADER * seedSection);
240
241 // PESection methods
242
243 unsigned dataLen() { return m_pSeedSectionHeader->SizeOfRawData; }
244 HRESULT applyRelocs(CeeGenTokenMapper *pTokenMapper) { return S_OK; }
245 char* getBlock(unsigned len, unsigned align) { _ASSERTE(!"PESeedSection"); return NULL; }
246 HRESULT truncate(unsigned newLen) { _ASSERTE(!"PESeedSection"); return E_FAIL; }
247 void writeSectReloc(unsigned val, CeeSection& relativeTo,
248 CeeSectionRelocType reloc,
249 CeeSectionRelocExtra *extra) { _ASSERTE(!"PESeedSection"); return; }
250 HRESULT addSectReloc(unsigned offset, CeeSection& relativeTo,
251 CeeSectionRelocType reloc,
252 CeeSectionRelocExtra *extra) { _ASSERTE(!"PESeedSection"); return E_FAIL; }
253 HRESULT addSectReloc(unsigned offset, PESection *relativeTo,
254 CeeSectionRelocType reloc,
255 CeeSectionRelocExtra *extra) { _ASSERTE(!"PESeedSection"); return E_FAIL; }
256 HRESULT addBaseReloc(unsigned offset, CeeSectionRelocType reloc,
257 CeeSectionRelocExtra *extra) { _ASSERTE(!"PESeedSection"); return E_FAIL; }
258// unsigned char *name();
259// unsigned flags();
260// unsigned getBaseRVA();
261 int getDirEntry() { _ASSERTE(!"PESeedSection"); return 0; }
262 HRESULT directoryEntry(unsigned num) { _ASSERTE(!"PESeedSection"); return E_FAIL; }
263 char * computePointer(unsigned offset) const { _ASSERTE(!"PESeedSection"); return NULL; }
264 BOOL containsPointer(__in char *ptr) const { _ASSERTE(!"PESeedSection"); return FALSE; }
265 unsigned computeOffset(__in char *ptr) const { _ASSERTE(!"PESeedSection"); return 0; }
266 HRESULT cloneInstance(PESection *destination) { _ASSERTE(!"PESeedSection"); return E_FAIL; }
267
268 // PEWriterSection
269
270 HRESULT applyRelocs(IMAGE_NT_HEADERS * pNtHeaders,
271 PERelocSection * relocSection,
272 CeeGenTokenMapper * pTokenMapper,
273 DWORD rdataRvaBase,
274 DWORD dataRvaBase,
275 DWORD textRvaBase) { return S_OK; }
276
277 HRESULT write(HANDLE file);
278 unsigned writeMem(void ** pMem);
279 bool isSeedSection() { return true; }
280
281protected:
282
283 PEDecoder * m_pSeedFileDecoder;
284 IMAGE_SECTION_HEADER * m_pSeedSectionHeader;
285
286};
287
288inline DWORD PEWriter::getSectionAlignment() {
289 return VAL32(m_ntHeaders->OptionalHeader.FileAlignment);
290}
291
292inline void PEWriter::setSectionAlignment(DWORD SectionAlignment) {
293
294 _ASSERTE(m_hSeedFile == INVALID_HANDLE_VALUE);
295 m_ntHeaders->OptionalHeader.SectionAlignment = VAL32(SectionAlignment);
296}
297
298inline DWORD PEWriter::getFileAlignment() {
299 return m_ntHeaders->OptionalHeader.FileAlignment;
300}
301
302inline void PEWriter::setFileAlignment(DWORD fileAlignment) {
303 _ASSERTE(m_hSeedFile == INVALID_HANDLE_VALUE);
304 m_ntHeaders->OptionalHeader.FileAlignment = VAL32(fileAlignment);
305}
306
307inline unsigned PEWriter::getSubsystem() {
308 return VAL16(m_ntHeaders->OptionalHeader.Subsystem);
309}
310
311inline void PEWriter::setSubsystem(unsigned subsystem, unsigned major, unsigned minor) {
312 m_ntHeaders->OptionalHeader.Subsystem = VAL16(subsystem);
313 m_ntHeaders->OptionalHeader.MajorSubsystemVersion = VAL16(major);
314 m_ntHeaders->OptionalHeader.MinorSubsystemVersion = VAL16(minor);
315}
316
317inline void PEWriter::setCharacteristics(unsigned mask) {
318 m_ntHeaders->FileHeader.Characteristics |= VAL16(mask);
319}
320
321inline void PEWriter::clearCharacteristics(unsigned mask) {
322 m_ntHeaders->FileHeader.Characteristics &= VAL16(~mask);
323}
324
325inline void PEWriter::setDllCharacteristics(unsigned mask) {
326 m_ntHeaders->OptionalHeader.DllCharacteristics |= VAL16(mask);
327}
328
329inline void PEWriter::clearDllCharacteristics(unsigned mask) {
330 m_ntHeaders->OptionalHeader.DllCharacteristics &= VAL16(~mask);
331}
332
333inline void PEWriter::setEntryPointTextOffset(unsigned offset) {
334 m_ntHeaders->OptionalHeader.AddressOfEntryPoint = VAL32(offset);
335}
336
337#endif
338