1 : // Copyright 2012 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 : // Declares an image layout builder, a utility class for constructing valid
16 : // ImageLayout objects.
17 :
18 : #ifndef SYZYGY_PE_IMAGE_LAYOUT_BUILDER_H_
19 : #define SYZYGY_PE_IMAGE_LAYOUT_BUILDER_H_
20 :
21 : #include <windows.h>
22 : #include <winnt.h>
23 : #include <vector>
24 :
25 : #include "syzygy/block_graph/block_graph.h"
26 : #include "syzygy/block_graph/ordered_block_graph.h"
27 : #include "syzygy/pe/image_layout.h"
28 : #include "syzygy/pe/pe_file_parser.h"
29 :
30 : namespace pe {
31 :
32 : // A helper class that assists in assigning address space to PE image sections,
33 : // building self-consistent PE image headers etc.
34 : class ImageLayoutBuilder {
35 : public:
36 : typedef block_graph::BlockGraph BlockGraph;
37 : typedef block_graph::OrderedBlockGraph OrderedBlockGraph;
38 : typedef core::RelativeAddress RelativeAddress;
39 :
40 : // Constructs a new image layout builder that populates the provided image
41 : // layout. The image layout must outlive the builder.
42 : explicit ImageLayoutBuilder(ImageLayout* image_layout);
43 :
44 : // Sets the padding. If this is non-zero, blank space will be left after
45 : // each block that is laid out. The contents of this space are dictated by
46 : // the PEFileWriter (which uses 0xcc for code sections, and 0x00 for data
47 : // sections).
48 E : void set_padding(size_t padding) { padding_ = padding; }
49 :
50 : // Accessors.
51 E : ImageLayout* image_layout() { return image_layout_; }
52 E : BlockGraph* block_graph() { return image_layout_->blocks.graph(); }
53 E : BlockGraph::Block* dos_header_block() { return dos_header_block_; }
54 E : BlockGraph::Block* nt_headers_block() { return nt_headers_block_; }
55 : size_t padding() const { return padding_; }
56 : const ImageLayout* image_layout() const { return image_layout_; }
57 : const BlockGraph* block_graph() const {
58 : return image_layout_->blocks.graph();
59 : }
60 : const BlockGraph::Block* dos_header_block() const {
61 : return dos_header_block_;
62 : }
63 : const BlockGraph::Block* nt_headers_block() const {
64 : return nt_headers_block_;
65 : }
66 :
67 : // Lays out the image headers, and sets the file and section alignment using
68 : // the values from the header.
69 : // @param dos_header_block must be a block that's a valid DOS header
70 : // and stub. This block must also refer to the NT headers block,
71 : // which in turn must contain valid NT headers.
72 : // @returns true iff the dos_header_block is valid.
73 : // @pre OpenSection and LayoutBlock must not have been called.
74 : bool LayoutImageHeaders(BlockGraph::Block* dos_header_block);
75 :
76 : // Opens a new section for writing. If another section is already open, closes
77 : // it first.
78 : // @param name the name of the section.
79 : // @param characteristics the section characteristics.
80 : // @param section a pointer to the section information.
81 : // @pre LayoutImageHeaders must have been called.
82 : bool OpenSection(const char* name, uint32 characteristics);
83 : bool OpenSection(const BlockGraph::Section* section);
84 :
85 : // Advances the layout cursor by the provided distance.
86 : // @param distance the distance by which to advance the output cursor.
87 : void AdvanceCursor(size_t distance);
88 :
89 : // Advances the layout cursor by the minimum distance necessary to satisfy
90 : // the requested alignment.
91 : // @param alignment the alignment to satisfy.
92 : // @returns the distance that the output cursor was advanced.
93 : size_t AlignCursor(size_t alignment);
94 :
95 : // Lays out the provided block using the blocks internal alignment.
96 : bool LayoutBlock(BlockGraph::Block* block);
97 :
98 : // Lays out the provided block using the provided alignment.
99 : bool LayoutBlock(size_t alignment, BlockGraph::Block* block);
100 :
101 : // Closes the initialized data portion of the section that is currently
102 : // being written. If this is not explicitly called for a section it will be
103 : // automatically determined based on block contents.
104 : // @pre OpenSection must already have been called.
105 : void CloseExplicitSectionData();
106 :
107 : // Closes the section that is currently being written.
108 : // @returns true on success, false otherwise.
109 : // @pre a section must currently be open.
110 : bool CloseSection();
111 :
112 : // Creates sections and lays out blocks using the provided ordered block
113 : // graph as a template. Lays out all sections except for the reloc section,
114 : // which must be the last section if it is present.
115 : // @param obg the ordered block graph to layout, which must be for the same
116 : // block-graph as used in the constructor.
117 : // @returns true on success, false otherwise.
118 : // @pre LayoutImageHeaders has been called.
119 : bool LayoutOrderedBlockGraph(const OrderedBlockGraph& obg);
120 :
121 : // Finalizes the image layout. This builds the relocs, finalizes the headers,
122 : // and does any other PE touch-ups that are required to make the image
123 : // self-consistent. This may remove and/or modify blocks in the block-graph.
124 : // @returns true on success, false otherwise.
125 : bool Finalize();
126 :
127 : private:
128 : // Lays out a block at the current cursor location.
129 : bool LayoutBlockImpl(BlockGraph::Block* block);
130 :
131 : // Ensure that the Safe SEH Table is sorted.
132 : bool SortSafeSehTable();
133 : // Allocates and populates a new relocations section containing
134 : // relocations for all absolute references in address_space_.
135 : bool CreateRelocsSection();
136 : // Write the NT headers and section headers to the image.
137 : // After this is done, the image is "baked", and everything except for
138 : // the image checksum should be up to date.
139 : bool FinalizeHeaders();
140 : // Ensure that the image layout has the same number of blocks as the
141 : // block-graph. The relocs blocks that are in the block-graph but not in the
142 : // image layout will be removed. If there are extra blocks from other sections
143 : // in the block-graph an error will be returned.
144 : // @returns true if the block-graph and the image layout are consistent,
145 : // false otherwise.
146 : bool ReconcileBlockGraphAndImageLayout();
147 :
148 : // The image layout we're building.
149 : ImageLayout* image_layout_;
150 :
151 : // The padding we're using.
152 : size_t padding_;
153 :
154 : // The current location of the output cursor, as well as information regarding
155 : // the current section in progress.
156 : RelativeAddress cursor_;
157 : RelativeAddress section_start_;
158 : RelativeAddress section_auto_init_end_;
159 : RelativeAddress section_init_end_;
160 :
161 : // The blocks that describe the DOS header and the NT headers.
162 : BlockGraph::Block* dos_header_block_;
163 : BlockGraph::Block* nt_headers_block_;
164 : };
165 :
166 : } // namespace pe
167 :
168 : #endif // SYZYGY_PE_IMAGE_LAYOUT_BUILDER_H_
|