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//
5// ZapWriter.h
6//
7
8//
9// Infrastructure for writing PE files. (Not NGEN specific)
10//
11// ======================================================================================
12
13
14
15#ifndef __ZAPWRITER_H__
16#define __ZAPWRITER_H__
17
18#include "zapnodetype.h"
19
20class ZapWriter;
21class ZapHeap;
22
23// This is maximum size of anything in the image written by ZapWriter. Used for overflow checking.
24#define ZAPWRITER_MAX_SIZE 0x3FFFFFFF
25
26// All ZapNodes should be allocated from ZapHeap returned by ZapWriter::GetHeap()
27void *operator new(size_t size, ZapHeap * pZapHeap);
28void *operator new[](size_t size, ZapHeap * pZapHeap);
29
30//
31// ZapHeap does not support deallocation. Empty operators delete avoids deallocating memory
32// if the constructor fails
33//
34inline void operator delete(void *, ZapHeap * pZapHeap)
35{
36 // Memory allocated by ZapHeap is never freed
37}
38inline void operator delete[](void *, ZapHeap * pZapHeap)
39{
40 // Memory allocated by ZapHeap is never freed
41}
42
43
44//------------------------------------------------------------------------------------------------------
45// ZapNode is the basic building block of the native image. Every ZapNode must know how to persist itself.
46//
47// The basic contract for a ZapNode is that it understands its allocations requirements (size and alignment),
48// and knows how to save itself (given a ZapWriter). At some point a ZapNode is given a location in the
49// executable (an RVA), which it is responsible remembering.
50//
51// See file:../../doc/BookOfTheRuntime/NGEN/NGENDesign.doc for an overview.
52//
53class ZapNode
54{
55 friend class ZapWriter;
56
57 DWORD m_RVA;
58
59public:
60 void SetRVA(DWORD dwRVA)
61 {
62 _ASSERTE(m_RVA == 0 || m_RVA == (DWORD)-1);
63 m_RVA = dwRVA;
64 }
65
66 ZapNode()
67 {
68 // All ZapNodes are expected to be allocate from ZapWriter::GetHeap() that returns zero filled memory
69 _ASSERTE(m_RVA == 0);
70 }
71
72 // This constructor should be used to allocate temporary ZapNodes on the stack only
73 ZapNode(DWORD rva)
74 : m_RVA(rva)
75 {
76 }
77
78 virtual ~ZapNode()
79 {
80 }
81
82 // Returns the size of the node in the image. All nodes that are written into the image should override this method.
83 virtual DWORD GetSize()
84 {
85#if defined(_MSC_VER) //UNREACHABLE doesn't work in GCC, when the method has a non-void return
86 UNREACHABLE();
87#else
88 _ASSERTE(!"Unreachable");
89 return 0;
90#endif
91 }
92
93 // Alignment for this node.
94 virtual UINT GetAlignment()
95 {
96 return 1;
97 }
98
99 // Returns the type of the ZapNode. All nodes should override this method.
100 virtual ZapNodeType GetType()
101 {
102 return ZapNodeType_Unknown;
103 }
104
105 // Assign RVA to this node. dwPos is current RVA, returns updated current RVA.
106 virtual DWORD ComputeRVA(ZapWriter * pZapWriter, DWORD dwPos);
107
108 // All nodes that are written into the image should override this method. The implementation should write exactly GetSize() bytes
109 // using ZapWriter::Write method
110 virtual void Save(ZapWriter * pZapWriter)
111 {
112 UNREACHABLE();
113 }
114
115 // Returns the RVA of the node. Valid only after ComputeRVA phase
116 DWORD GetRVA()
117 {
118 _ASSERTE(m_RVA != 0 && m_RVA != (DWORD)-1);
119 return m_RVA;
120 }
121
122 // Returns whether the node was placed into a virtual section
123 BOOL IsPlaced()
124 {
125 return m_RVA != 0;
126 }
127};
128
129//---------------------------------------------------------------------------------------
130// Virtual section of PE image.
131class ZapVirtualSection : public ZapNode
132{
133 friend class ZapWriter;
134
135 DWORD m_dwAlignment;
136
137 SArray<ZapNode *> m_Nodes;
138
139 // State initialized once the section is placed
140 DWORD m_dwSize;
141
142 DWORD m_dwSectionType;
143
144 BYTE m_defaultFill;
145
146 ZapVirtualSection(DWORD dwAlignment)
147 : m_dwAlignment(dwAlignment)
148 {
149 }
150
151public:
152 virtual DWORD GetSize()
153 {
154 return m_dwSize;
155 }
156
157 virtual ZapNodeType GetType()
158 {
159 return ZapNodeType_VirtualSection;
160 }
161
162 DWORD GetSectionType()
163 {
164 return m_dwSectionType;
165 }
166
167 void SetSectionType(DWORD dwSectionType)
168 {
169 _ASSERTE((dwSectionType & IBCTypeReservedFlag) != 0 || !"IBCType flag is not specified");
170 _ASSERTE((dwSectionType & RangeTypeReservedFlag) != 0 || !"RangeType flag is not specified");
171 _ASSERTE((dwSectionType & VirtualSectionTypeReservedFlag) != 0 || !"VirtualSectionType flag is not specified");
172 _ASSERTE((dwSectionType & VirtualSectionTypeReservedFlag) < CORCOMPILE_SECTION_TYPE_COUNT || !"Invalid VirtualSectionType flag");
173 m_dwSectionType = dwSectionType;
174 }
175
176 void SetDefaultFill(BYTE fill)
177 {
178 m_defaultFill = fill;
179 }
180
181 void Place(ZapNode * pNode)
182 {
183 _ASSERTE(!pNode->IsPlaced());
184 m_Nodes.Append(pNode);
185 pNode->SetRVA((DWORD)-1);
186 }
187
188 COUNT_T GetNodeCount()
189 {
190 return m_Nodes.GetCount();
191 }
192
193 ZapNode * GetNode(COUNT_T iNode)
194 {
195 return m_Nodes[iNode];
196 }
197};
198
199//---------------------------------------------------------------------------------------
200// The named physical section of the PE Image. It contains one or more virtual sections.
201class ZapPhysicalSection : public ZapNode
202{
203 friend class ZapWriter;
204
205 SArray<ZapVirtualSection *> m_Sections;
206
207 LPCSTR m_pszName;
208 DWORD m_dwCharacteristics;
209
210 // Number of zero filled sections (zero filled sections are always last in m_Sections array)
211 COUNT_T m_nBssSections;
212
213 // State initialized once the section is placed
214 DWORD m_dwSize;
215 DWORD m_dwFilePos;
216 DWORD m_dwSizeOfRawData;
217
218 ZapPhysicalSection(LPCSTR pszName, DWORD dwCharacteristics)
219 : m_pszName(pszName),
220 m_dwCharacteristics(dwCharacteristics)
221 {
222 }
223
224public:
225 ~ZapPhysicalSection()
226 {
227 for (COUNT_T iVirtualSection = 0; iVirtualSection < m_Sections.GetCount(); iVirtualSection++)
228 {
229 ZapVirtualSection * pVirtualSection = m_Sections[iVirtualSection];
230 pVirtualSection->~ZapVirtualSection();
231 }
232 }
233
234 virtual DWORD GetSize()
235 {
236 return m_dwSize;
237 }
238
239 virtual ZapNodeType GetType()
240 {
241 return ZapNodeType_PhysicalSection;
242 }
243
244 DWORD GetFilePos()
245 {
246 _ASSERTE(m_dwFilePos != 0);
247 return m_dwFilePos;
248 }
249
250 COUNT_T GetVirtualSectionCount()
251 {
252 return m_Sections.GetCount();
253 }
254
255 ZapVirtualSection * GetVirtualSection(COUNT_T iSection)
256 {
257 return m_Sections[iSection];
258 }
259
260};
261
262//---------------------------------------------------------------------------------------
263//
264// The ZapWriter
265//
266// Notice that ZapWriter implements IStream that can be passed to APIs that write to stream
267//
268// The main API in a ZapWriter is (not suprisingly) the code:ZapWriter.Write method.
269//
270// Relocations are handled by a higher level object, code:ZapImage, which knows about all the sections of a
271// ngen image and how to do relections. Every ZapWriter has an associated ZapImage which you get to by
272// calling code:ZapImage.GetImage.
273//
274class ZapWriter : public IStream
275{
276 ZapHeap * m_pHeap;
277
278 SArray<ZapPhysicalSection *> m_Sections;
279
280 ZapNode * m_DirectoryEntries[IMAGE_NUMBEROF_DIRECTORY_ENTRIES];
281 DWORD m_dwTimeDateStamp;
282 ULONGLONG m_BaseAddress;
283 ULONGLONG m_SizeOfStackReserve;
284 ULONGLONG m_SizeOfStackCommit;
285 USHORT m_Subsystem;
286 USHORT m_DllCharacteristics;
287 BOOL m_isDll;
288 DWORD m_FileAlignment;
289
290 // Current state of the writer for debug checks
291 INDEBUG(BOOL m_fSaving;)
292
293 DWORD m_dwCurrentRVA;
294 BOOL m_fWritingRelocs; // Set to true once we start reloc sections at the end of the file
295
296 void SaveContent();
297
298 DWORD GetSizeOfNTHeaders();
299 void SaveHeaders();
300
301 // Simple buffered writer
302 void InitializeWriter(IStream * pStream);
303
304 IStream * m_pStream;
305 PBYTE m_pBuffer;
306 ULONG m_nBufferPos;
307 INDEBUG(DWORD m_dwWriterFilePos;)
308
309 //
310 // NT Headers
311 //
312
313 BOOL Is64Bit()
314 {
315#ifdef _TARGET_64BIT_
316 return TRUE;
317#else // !_TARGET_64BIT_
318 return FALSE;
319#endif // !_TARGET_64BIT_
320 }
321
322 USHORT GetMachine()
323 {
324 return IMAGE_FILE_MACHINE_NATIVE_NI;
325 }
326
327 void SaveDosHeader();
328 void SaveSignature();
329 void SaveFileHeader();
330 void SaveOptionalHeader();
331 void SaveSections();
332
333 // IStream support - the only actually implemented method is IStream::Write
334
335 // IUnknown methods
336 STDMETHODIMP_(ULONG) AddRef()
337 {
338 return 1;
339 }
340
341 STDMETHODIMP_(ULONG) Release()
342 {
343 return 1;
344 }
345
346 STDMETHODIMP QueryInterface(REFIID riid, LPVOID *ppv)
347 {
348 HRESULT hr = S_OK;
349 if (IsEqualIID(riid, IID_IUnknown) || IsEqualIID(riid, IID_IStream)) {
350 *ppv = static_cast<IStream *>(this);
351 }
352 else {
353 *ppv = NULL;
354 hr = E_NOINTERFACE;
355 }
356 return hr;
357 }
358
359 // ISequentialStream methods:
360 STDMETHODIMP Read(void *pv, ULONG cb, ULONG *pcbRead)
361 {
362 _ASSERTE(false);
363 return E_NOTIMPL;
364 }
365
366 STDMETHODIMP Write(void const *pv, ULONG cb, ULONG *pcbWritten);
367
368 // IStream methods:
369 STDMETHODIMP Seek(LARGE_INTEGER dlibMove, DWORD dwOrigin, ULARGE_INTEGER *plibNewPosition)
370 {
371 // IMetaDataEmit::SaveToStream calls Seek(0) but ignores the returned error
372 //_ASSERTE(false);
373 return E_NOTIMPL;
374 }
375
376 STDMETHODIMP SetSize(ULARGE_INTEGER libNewSize)
377 {
378 _ASSERTE(false);
379 return E_NOTIMPL;
380 }
381
382 STDMETHODIMP CopyTo(IStream *pstm, ULARGE_INTEGER cb, ULARGE_INTEGER *pcbRead, ULARGE_INTEGER *pcbWritten)
383 {
384 _ASSERTE(false);
385 return E_NOTIMPL;
386 }
387
388 STDMETHODIMP Commit(DWORD grfCommitFlags)
389 {
390 _ASSERTE(false);
391 return E_NOTIMPL;
392 }
393
394 STDMETHODIMP Revert()
395 {
396 _ASSERTE(false);
397 return E_NOTIMPL;
398 }
399
400 STDMETHODIMP LockRegion(ULARGE_INTEGER libOffset, ULARGE_INTEGER cb, DWORD dwLockType)
401 {
402 _ASSERTE(false);
403 return E_NOTIMPL;
404 }
405
406 STDMETHODIMP UnlockRegion(ULARGE_INTEGER libOffset, ULARGE_INTEGER cb, DWORD dwLockType)
407 {
408 _ASSERTE(false);
409 return E_NOTIMPL;
410 }
411
412 STDMETHODIMP Stat(STATSTG *pstatstg, DWORD grfStatFlag)
413 {
414 _ASSERTE(false);
415 return E_NOTIMPL;
416 }
417
418 STDMETHODIMP Clone(IStream **ppIStream)
419 {
420 _ASSERTE(false);
421 return E_NOTIMPL;
422 }
423
424public:
425 ZapWriter();
426 ~ZapWriter();
427
428 void Initialize();
429
430 // Create new section in the PE file. The sections will be saved in the order they are created.
431 ZapPhysicalSection * NewPhysicalSection(LPCSTR pszName, DWORD dwCharacteristics)
432 {
433 _ASSERTE(!IsSaving());
434 ZapPhysicalSection * pSection = new (GetHeap()) ZapPhysicalSection(pszName, dwCharacteristics);
435 m_Sections.Append(pSection);
436 return pSection;
437 }
438
439 // Create new virtual section within the physical section. The sections will be saved in the order they are created.
440 // The default virtual section alignment is 16.
441 ZapVirtualSection * NewVirtualSection(ZapPhysicalSection * pPhysicalSection, DWORD dwAlignment = 16, ZapVirtualSection * pInsertAfter = NULL)
442 {
443 _ASSERTE(!IsSaving());
444 ZapVirtualSection * pSection = new (GetHeap()) ZapVirtualSection(dwAlignment);
445 if (pInsertAfter != NULL)
446 {
447 // pInsertAfter is workaround to get decent layout with the current scheme of virtual sections. It should not be necessary
448 // once we have better layout algorithms in place.
449 for (COUNT_T iSection = 0; iSection < pPhysicalSection->m_Sections.GetCount(); iSection++)
450 {
451 if (pPhysicalSection->m_Sections[iSection] == pInsertAfter)
452 {
453 pPhysicalSection->m_Sections.Insert(pPhysicalSection->m_Sections+(iSection+1));
454 pPhysicalSection->m_Sections[iSection+1] = pSection;
455 return pSection;
456 }
457 }
458 _ASSERTE(false);
459 }
460
461 pPhysicalSection->m_Sections.Append(pSection);
462 return pSection;
463 }
464
465 void MarkBssSection(ZapPhysicalSection * pPhysicalSection, ZapVirtualSection * pSection)
466 {
467 _ASSERTE(!IsSaving());
468 _ASSERTE(pPhysicalSection->m_Sections[pPhysicalSection->m_Sections.GetCount() - 1] == pSection);
469 pPhysicalSection->m_nBssSections++;
470 }
471
472 void Append(ZapVirtualSection * pVirtualSection, ZapNode * pNode)
473 {
474 _ASSERTE(!IsSaving());
475 pVirtualSection->m_Nodes.Append(pNode);
476 }
477
478 // Set the directory entry in the image to match the given ZapNode
479 void SetDirectoryEntry(DWORD entry, ZapNode * pNode)
480 {
481 _ASSERTE(!IsSaving());
482 _ASSERTE(entry < IMAGE_NUMBEROF_DIRECTORY_ENTRIES);
483 _ASSERTE(m_DirectoryEntries[entry] == NULL);
484 m_DirectoryEntries[entry] = pNode;
485 }
486
487 // Set the timedate stamp of the image
488 void SetTimeDateStamp(DWORD dwTimeDateStamp)
489 {
490 _ASSERTE(!IsSaving());
491 m_dwTimeDateStamp = dwTimeDateStamp;
492 }
493
494 // Set the base address of the image
495 void SetBaseAddress(ULONGLONG baseAddress)
496 {
497 _ASSERTE(!IsSaving());
498 m_BaseAddress = baseAddress;
499 }
500
501 ULONGLONG GetBaseAddress()
502 {
503 _ASSERTE(m_BaseAddress != 0);
504 return m_BaseAddress;
505 }
506
507 void SetSizeOfStackReserve(ULONGLONG sizeOfStackReserve)
508 {
509 _ASSERTE(!IsSaving());
510 m_SizeOfStackReserve = sizeOfStackReserve;
511 }
512
513 void SetSizeOfStackCommit(ULONGLONG sizeOfStackCommit)
514 {
515 _ASSERTE(!IsSaving());
516 m_SizeOfStackCommit = sizeOfStackCommit;
517 }
518
519 void SetSubsystem(USHORT subsystem)
520 {
521 _ASSERTE(!IsSaving());
522 m_Subsystem = subsystem;
523 }
524
525 void SetDllCharacteristics(USHORT dllCharacteristics)
526 {
527 _ASSERTE(!IsSaving());
528 m_DllCharacteristics = dllCharacteristics;
529 }
530
531 void SetIsDll(BOOL isDLL)
532 {
533 m_isDll = isDLL;
534 }
535
536 void SetFileAlignment(DWORD fileAlignment)
537 {
538 m_FileAlignment = fileAlignment;
539 }
540
541 // Compute RVAs for everything in the file
542 void ComputeRVAs();
543
544 // Save the content into stream
545 void Save(IStream * pStream);
546
547 // Get the heap. The lifetime of this heap is same as the lifetime of the ZapWriter. All ZapNodes should
548 // be allocated from this heap.
549 ZapHeap * GetHeap()
550 {
551 return m_pHeap;
552 }
553
554 COUNT_T GetPhysicalSectionCount()
555 {
556 return m_Sections.GetCount();
557 }
558
559 ZapPhysicalSection * GetPhysicalSection(COUNT_T iSection)
560 {
561 return m_Sections[iSection];
562 }
563
564#ifdef _DEBUG
565 // Certain methods can be called only during the save phase
566 BOOL IsSaving()
567 {
568 return m_fSaving;
569 }
570#endif
571
572 DWORD GetCurrentRVA()
573 {
574 _ASSERTE(IsSaving());
575 return m_dwCurrentRVA;
576 }
577
578
579 // This is the main entrypoint used to write the image. Every implementation of ZapNode::Save will call this method.
580 void Write(PVOID p, DWORD dwSize);
581
582 // Writes padding
583 void WritePad(DWORD size, BYTE fill = 0);
584
585 // Flush any buffered data
586 void FlushWriter();
587
588 BOOL IsWritingRelocs()
589 {
590 return m_fWritingRelocs;
591 }
592
593 void SetWritingRelocs()
594 {
595 m_fWritingRelocs = TRUE;
596 }
597
598 // Convenience helper to initialize IMAGE_DATA_DIRECTORY
599 static void SetDirectoryData(IMAGE_DATA_DIRECTORY * pDir, ZapNode * pZapNode);
600};
601
602//---------------------------------------------------------------------------------------
603// ZapBlob
604//
605// Generic node for unstructured sequence of bytes.
606// Includes SHash support (ZapBlob::SHashTraits)
607//
608class ZapBlob : public ZapNode
609{
610 DWORD m_cbSize;
611
612protected:
613 ZapBlob(SIZE_T cbSize)
614 : m_cbSize((DWORD)cbSize)
615 {
616 if (cbSize > ZAPWRITER_MAX_SIZE)
617 ThrowHR(COR_E_OVERFLOW);
618 }
619
620public:
621 class SHashKey
622 {
623 PBYTE m_pData;
624 SIZE_T m_cbSize;
625
626 public:
627 SHashKey(PVOID pData, SIZE_T cbSize)
628 : m_pData((PBYTE)pData), m_cbSize(cbSize)
629 {
630 }
631
632 PBYTE GetData() const
633 {
634 return m_pData;
635 }
636
637 SIZE_T GetBlobSize() const
638 {
639 return m_cbSize;
640 }
641 };
642
643 class SHashTraits : public DefaultSHashTraits<ZapBlob *>
644 {
645 public:
646 typedef const ZapBlob::SHashKey key_t;
647
648 static key_t GetKey(element_t e)
649 {
650 LIMITED_METHOD_CONTRACT;
651 return key_t(e->GetData(), e->GetBlobSize());
652 }
653 static BOOL Equals(key_t k1, key_t k2)
654 {
655 LIMITED_METHOD_CONTRACT;
656 if (k1.GetBlobSize() != k2.GetBlobSize())
657 return FALSE;
658 return memcmp(k1.GetData(), k2.GetData(), k1.GetBlobSize()) == 0;
659 }
660 static count_t Hash(key_t k)
661 {
662 LIMITED_METHOD_CONTRACT;
663 count_t hash = 5381 + (count_t)(k.GetBlobSize() << 7);
664
665 PBYTE pbData = k.GetData();
666 PBYTE pbDataEnd = pbData + k.GetBlobSize();
667
668 for (/**/ ; pbData < pbDataEnd; pbData++)
669 {
670 hash = ((hash << 5) + hash) ^ *pbData;
671 }
672 return hash;
673 }
674 };
675
676 virtual PBYTE GetData()
677 {
678 return (PBYTE)(this + 1);
679 }
680
681 // Used to shrink the size of the blob
682 void AdjustBlobSize(SIZE_T cbSize)
683 {
684 _ASSERTE(cbSize <= m_cbSize);
685 _ASSERTE(cbSize != 0);
686 m_cbSize = (DWORD)cbSize;
687 }
688
689 // Raw size of the blob
690 DWORD GetBlobSize()
691 {
692 return m_cbSize;
693 }
694
695 virtual DWORD GetSize()
696 {
697 return m_cbSize;
698 }
699
700 virtual ZapNodeType GetType()
701 {
702 return ZapNodeType_Blob;
703 }
704
705 virtual void Save(ZapWriter * pZapWriter);
706
707 // Create new zap blob node. The node *does* own copy of the memory.
708 static ZapBlob * NewBlob(ZapWriter * pWriter, PVOID pData, SIZE_T cbSize);
709
710 // Create new aligned zap blob node.
711 static ZapBlob * NewAlignedBlob(ZapWriter * pWriter, PVOID pData, SIZE_T cbSize, SIZE_T cbAlignment);
712};
713
714class ZapBlobPtr : public ZapBlob
715{
716 PBYTE m_pData;
717
718public:
719 ZapBlobPtr(PVOID pData, SIZE_T cbSize)
720 : ZapBlob(cbSize), m_pData((PBYTE)pData)
721 {
722 }
723
724 virtual PBYTE GetData()
725 {
726 return m_pData;
727 }
728};
729
730class ZapDummyNode : public ZapNode
731{
732 DWORD m_cbSize;
733
734public:
735 ZapDummyNode(DWORD cbSize)
736 : m_cbSize(cbSize)
737 {
738 }
739
740 virtual DWORD GetSize()
741 {
742 return m_cbSize;
743 }
744};
745
746#endif // __ZAPWRITER_H__
747