1 : // Copyright 2012 Google Inc. All Rights Reserved.
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 : // Lays out the provided block using the blocks internal alignment.
86 : bool LayoutBlock(BlockGraph::Block* block);
87 :
88 : // Lays out the provided block using the provided alignment.
89 : bool LayoutBlock(size_t alignment, BlockGraph::Block* block);
90 :
91 : // Closes the initialized data portion of the section that is currently
92 : // being written. If this is not explicitly called for a section it will be
93 : // automatically determined based on block contents.
94 : // @pre OpenSection must already have been called.
95 : void CloseExplicitSectionData();
96 :
97 : // Closes the section that is currently being written.
98 : // @returns true on success, false otherwise.
99 : // @pre a section must currently be open.
100 : bool CloseSection();
101 :
102 : // Creates sections and lays out blocks using the provided ordered block
103 : // graph as a template. Lays out all sections except for the reloc section,
104 : // which must be the last section if it is present.
105 : // @param obg the ordered block graph to layout, which must be for the same
106 : // block-graph as used in the constructor.
107 : // @returns true on success, false otherwise.
108 : // @pre LayoutImageHeaders has been called.
109 : bool LayoutOrderedBlockGraph(const OrderedBlockGraph& obg);
110 :
111 : // Finalizes the image layout. This builds the relocs, finalizes the headers,
112 : // and does any other PE touch-ups that are required to make the image
113 : // self-consistent. This may remove and/or modify blocks in the block-graph.
114 : // @returns true on success, false otherwise.
115 : bool Finalize();
116 :
117 : private:
118 : // Lays out a block at the current cursor location.
119 : bool LayoutBlockImpl(BlockGraph::Block* block);
120 :
121 : // Ensure that the Safe SEH Table is sorted.
122 : bool SortSafeSehTable();
123 : // Allocates and populates a new relocations section containing
124 : // relocations for all absolute references in address_space_.
125 : bool CreateRelocsSection();
126 : // Write the NT headers and section headers to the image.
127 : // After this is done, the image is "baked", and everything except for
128 : // the image checksum should be up to date.
129 : bool FinalizeHeaders();
130 : // Ensure that the image layout has the same number of blocks as the
131 : // block-graph. The relocs blocks that are in the block-graph but not in the
132 : // image layout will be removed. If there are extra blocks from other sections
133 : // in the block-graph an error will be returned.
134 : // @returns true if the block-graph and the image layout are consistent,
135 : // false otherwise.
136 : bool ReconcileBlockGraphAndImageLayout();
137 :
138 : // The image layout we're building.
139 : ImageLayout* image_layout_;
140 :
141 : // The padding we're using.
142 : size_t padding_;
143 :
144 : // The current location of the output cursor, as well as information regarding
145 : // the current section in progress.
146 : RelativeAddress cursor_;
147 : RelativeAddress section_start_;
148 : RelativeAddress section_auto_init_end_;
149 : RelativeAddress section_init_end_;
150 :
151 : // The blocks that describe the DOS header and the NT headers.
152 : BlockGraph::Block* dos_header_block_;
153 : BlockGraph::Block* nt_headers_block_;
154 : };
155 :
156 : } // namespace pe
157 :
158 : #endif // SYZYGY_PE_IMAGE_LAYOUT_BUILDER_H_
|