OLD | NEW |
(Empty) | |
| 1 // Copyright 2015 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 #include "syzygy/experimental/protect/protect_lib/integrity_check_layout_transfo
rm.h" |
| 15 |
| 16 //TODO: remove this include |
| 17 #include <inttypes.h> |
| 18 |
| 19 #include "syzygy/core/address.h" |
| 20 #include "syzygy/experimental/protect/protect_lib/protect_utils.h" |
| 21 #include "syzygy/pe/pe_file_writer.h" |
| 22 #include "syzygy/pe/pe_utils.h" |
| 23 |
| 24 namespace protect { |
| 25 |
| 26 uint8_t IntegrityCheckLayoutTransform::ComputeAggregatedChunksHash( |
| 27 const std::set<uint32_t> chunk_indexes){ |
| 28 uint8_t precomputed_xor = 0; |
| 29 for (auto chunk_it = chunk_indexes.rbegin(); |
| 30 chunk_it != chunk_indexes.rend(); |
| 31 ++chunk_it) { |
| 32 auto chunk_info = (*ic_block_reference_free_chunks)[*chunk_it]; |
| 33 precomputed_xor += chunk_info.hash_ + chunk_info.hash_of_next_instruction_; |
| 34 precomputed_xor = -precomputed_xor; |
| 35 } |
| 36 return precomputed_xor; |
| 37 } |
| 38 |
| 39 uint8_t |
| 40 IntegrityCheckLayoutTransform::ComputeAggregatedBlocksHash(uint64_t bb_id){ |
| 41 uint8_t precomputed_hash = 0; |
| 42 auto checkee_iter = (*this->checker_to_checkee_map_)[bb_id].begin(); |
| 43 for (; checkee_iter != (*this->checker_to_checkee_map_)[bb_id].end(); |
| 44 ++checkee_iter){ |
| 45 precomputed_hash += (*this->precomputed_hashes_)[checkee_iter->first] * |
| 46 checkee_iter->second; |
| 47 } |
| 48 return precomputed_hash; |
| 49 } |
| 50 |
| 51 bool IntegrityCheckLayoutTransform::RecomputePivot( |
| 52 const uint64_t bb_id, |
| 53 const uint8_t precomputed_hash, |
| 54 const uint8_t precomputed_xor, |
| 55 const size_t pivot_offset, |
| 56 const size_t sub_offset, |
| 57 block_graph::BlockGraph::Block *block){ |
| 58 |
| 59 const uint8_t *pivot_byte = block->data() + pivot_offset; |
| 60 DCHECK(*pivot_byte == 0x00); |
| 61 |
| 62 const uint8_t *sub_opcode = block->data() + sub_offset; |
| 63 DCHECK(*sub_opcode == 0x2c); |
| 64 |
| 65 const uint8_t *hash_byte = sub_opcode + 1; |
| 66 //+1 is the sub hash value |
| 67 uint8_t new_bytes[3]; |
| 68 new_bytes[0] = *pivot_byte; |
| 69 new_bytes[1] = *sub_opcode; |
| 70 new_bytes[2] = *hash_byte; |
| 71 |
| 72 uint8_t old_hash = *hash_byte; |
| 73 new_bytes[2] = precomputed_hash + precomputed_xor; // new hash |
| 74 const uint8_t* data = block->data(); |
| 75 size_t data_size = block->data_size(); |
| 76 uint8_t* new_data = new uint8_t[data_size]; |
| 77 memcpy(new_data, data, data_size); |
| 78 |
| 79 // Set new pivot byte. Starts at offset 0 |
| 80 uint8_t new_pivot = old_hash - new_bytes[2]; |
| 81 new_bytes[0] = new_pivot; |
| 82 |
| 83 DCHECK_EQ((uint8_t)(new_pivot + precomputed_hash + precomputed_xor), |
| 84 old_hash); |
| 85 new_data[pivot_offset] = new_bytes[0]; |
| 86 new_data[sub_offset + 1] = new_bytes[2]; |
| 87 |
| 88 block->CopyData(data_size, new_data); |
| 89 delete[] new_data; |
| 90 |
| 91 return true; |
| 92 } |
| 93 |
| 94 bool IntegrityCheckLayoutTransform::PatchPivot(BlockGraph::Label label) { |
| 95 uint64_t bb_id = GetBasicBlockIdByLabel(label, this->id_to_label_); |
| 96 |
| 97 if (bb_id == -1) |
| 98 return true; |
| 99 |
| 100 if ((*this->checker_to_checkee_map_)[bb_id].size() < 1) |
| 101 return true; |
| 102 |
| 103 uint8_t precomputed_hash = ComputeAggregatedBlocksHash(bb_id); |
| 104 uint8_t precomputed_xor = 0; |
| 105 if (*perform_chunk_checks_) { |
| 106 //recompute xor hash |
| 107 auto checkee_chunks_it = ic_chunk_checker_to_checkee_map_->find(bb_id); |
| 108 DCHECK(checkee_chunks_it != ic_chunk_checker_to_checkee_map_->end()); |
| 109 |
| 110 DCHECK_NE(checkee_chunks_it->second.size(), static_cast<uint32_t>(0)); |
| 111 precomputed_xor = ComputeAggregatedChunksHash(checkee_chunks_it->second); |
| 112 } |
| 113 char *buffer = new char[50]; |
| 114 sprintf_s(buffer, 50, "Pivot:%llu", bb_id); |
| 115 // offset of sub instruction after returning from hash function |
| 116 size_t pivot_offset = (*label_name_to_block_)[buffer].second; |
| 117 auto block = (*label_name_to_block_)[buffer].first; |
| 118 |
| 119 sprintf_s(buffer, 50, "sub %llu", bb_id); |
| 120 // offset of sub instruction after returning from hash function |
| 121 auto sub_instr_block = label_name_to_block_->find(buffer); |
| 122 DCHECK(sub_instr_block != label_name_to_block_->end()); |
| 123 size_t sub_offset = sub_instr_block->second.second; |
| 124 delete[] buffer; |
| 125 if (RecomputePivot(bb_id, precomputed_hash, precomputed_xor, |
| 126 pivot_offset, sub_offset, block)){ |
| 127 this->nr_hashes_patched_++; |
| 128 } |
| 129 return true; |
| 130 } |
| 131 |
| 132 int IntegrityCheckLayoutTransform::PatchPrecomputedHashes( |
| 133 const TransformPolicyInterface* policy, |
| 134 BlockGraph::Block* block) { |
| 135 if (!ShouldPostProcessBlock(block, this->id_to_label_)) |
| 136 return 0; |
| 137 |
| 138 // Iterate over every label in the block and patch the pivot |
| 139 auto it = block->labels().begin(); |
| 140 for (; it != block->labels().end(); ++it) { |
| 141 if (!PatchPivot(it->second)) |
| 142 return 0; |
| 143 } |
| 144 |
| 145 return 0; |
| 146 } |
| 147 // This function adjusts the inter-block references that have shifted after |
| 148 // code was inserted |
| 149 bool IntegrityCheckLayoutTransform::CheckHash( |
| 150 BasicCodeBlock* bb, |
| 151 std::vector<uint8_t> new_block_buffer, |
| 152 const core::AbsoluteAddress image_base) { |
| 153 auto inst_iter = bb->instructions().begin(); |
| 154 uint64_t bb_id = (uint64_t)-1; |
| 155 if ((inst_iter != bb->instructions().end()) && (inst_iter->has_label())) { |
| 156 bb_id = GetBasicBlockIdByLabel(inst_iter->label(), this->id_to_label_); |
| 157 if (bb_id != -1) { |
| 158 auto buf_it = new_block_buffer.begin(); |
| 159 std::advance(buf_it, bb->offset()); |
| 160 |
| 161 uint8_t hash = 0; |
| 162 uint32_t block_size = (*this->basic_block_sizes_)[bb_id]; |
| 163 for (uint32_t i = 0; i < block_size; ++i) { |
| 164 hash += *buf_it; |
| 165 if (i % 16 == 0) |
| 166 fprintf(phash, "\n"); |
| 167 else if (i % 8 == 0) |
| 168 fprintf(phash, " "); |
| 169 fprintf(phash, "%02X ", *buf_it); |
| 170 ++buf_it; |
| 171 } |
| 172 // Compute hash of image_base. |
| 173 uint8_t hash_image_base = 0; |
| 174 for (uint8_t i = 0; i < 4; i++) { |
| 175 hash_image_base += image_base.value() >> (i*8); |
| 176 } |
| 177 |
| 178 // For each chunk or checkee subtract hash of image base. |
| 179 uint8_t nr_checkees = (*this->checker_to_checkee_map_)[bb_id].size() + |
| 180 (*ic_chunk_checker_to_checkee_map_)[bb_id].size(); |
| 181 hash -= hash_image_base * nr_checkees; |
| 182 |
| 183 uint8_t precompute_hash = (*this->precomputed_hashes_)[bb_id]; |
| 184 if (precompute_hash != hash) { |
| 185 (*this->precomputed_hashes_)[bb_id] = hash; |
| 186 } |
| 187 |
| 188 fprintf(phash, "\n%s,", bb->subgraph()->original_block()->name().c_str()); |
| 189 fprintf(phash, "%" PRIx64 ",", bb_id); |
| 190 fprintf(phash, "%" PRIx32 ",", |
| 191 bb->subgraph()->original_block()->addr().value() + bb->offset(
)); |
| 192 fprintf(phash, "%" PRIx8 "\n", hash); |
| 193 } //end if |
| 194 } //end if |
| 195 |
| 196 //We need to compute hash of the chunks whose last instruction has absolute |
| 197 // address. If there is no chunk checking this step is not needed. |
| 198 if (!*perform_chunk_checks_) return true; |
| 199 |
| 200 std::string chunk_pointerlabel = "n "; |
| 201 auto end_block = bb->instructions().end(); |
| 202 uint64_t chunk_bb_id; |
| 203 uint32_t chunk_index; |
| 204 uint32_t offset = bb->offset(); |
| 205 for (; inst_iter != end_block; ++inst_iter) |
| 206 { |
| 207 offset += inst_iter->size(); |
| 208 if (!inst_iter->has_label()) continue; |
| 209 |
| 210 if (inst_iter->label().name() |
| 211 .compare(0, chunk_pointerlabel.length(), chunk_pointerlabel) == 0){ |
| 212 // update last visited chunk index |
| 213 GetChunkTokensFromlabel(inst_iter->label().name(), |
| 214 &chunk_bb_id, |
| 215 &chunk_index); |
| 216 |
| 217 size_t unique_key = GetChunkUniqueKey(chunk_bb_id, chunk_index); |
| 218 |
| 219 uint32_t vector_index = (*ic_block_chunk_index_map_)[unique_key]; |
| 220 |
| 221 DCHECK_GE(vector_index, static_cast<uint32_t>(0)); |
| 222 DCHECK_LT(vector_index, ic_block_reference_free_chunks->size()); |
| 223 |
| 224 auto chunk = (*ic_block_reference_free_chunks)[vector_index]; |
| 225 DCHECK(chunk.block_id_ == chunk_bb_id); |
| 226 DCHECK(chunk.chunk_index_ == chunk_index); |
| 227 //we need to recompute chunks whose last instruction has absoloute address |
| 228 if (chunk.next_instruction_size_ == 0) continue; |
| 229 |
| 230 uint32_t chunk_offset = offset + chunk.size_ - inst_iter->size(); |
| 231 |
| 232 auto buf_it = new_block_buffer.begin(); |
| 233 std::advance(buf_it, chunk_offset); |
| 234 |
| 235 uint8_t hash = 0; |
| 236 for (uint32_t i = 0; i < chunk.next_instruction_size_; ++i) { |
| 237 hash += *buf_it; |
| 238 ++buf_it; |
| 239 } |
| 240 |
| 241 chunk.hash_of_next_instruction_ = hash; |
| 242 (*ic_block_reference_free_chunks)[vector_index] = chunk; |
| 243 } |
| 244 } |
| 245 |
| 246 return true; |
| 247 } |
| 248 |
| 249 bool IntegrityCheckLayoutTransform::FixPrecomputedHashes( |
| 250 const TransformPolicyInterface* policy, |
| 251 const core::AbsoluteAddress image_base, |
| 252 BlockGraph::Block* block, |
| 253 std::vector<uint8_t> new_block_buffer) { |
| 254 |
| 255 if (!ShouldPostProcessBlock(block, this->id_to_label_)) |
| 256 return false; |
| 257 |
| 258 // Use the decomposition policy to skip blocks that aren't eligible for |
| 259 // basic-block decomposition. |
| 260 if (!policy->BlockIsSafeToBasicBlockDecompose(block)) |
| 261 return false; |
| 262 |
| 263 // Decompose block to basic blocks. |
| 264 BasicBlockSubGraph *subgraph = new BasicBlockSubGraph(); |
| 265 BasicBlockDecomposer bb_decomposer(block, subgraph); |
| 266 if (!bb_decomposer.Decompose()) |
| 267 return false; |
| 268 |
| 269 BasicBlockSubGraph::BBCollection& basic_blocks = |
| 270 subgraph->basic_blocks(); // set of BB to protect |
| 271 |
| 272 // Iterate over every basic block and recompute the hash |
| 273 for (auto it = basic_blocks.begin(); it != basic_blocks.end(); ++it) { |
| 274 BasicCodeBlock* bb = BasicCodeBlock::Cast(*it); |
| 275 |
| 276 if (bb == NULL) |
| 277 continue; |
| 278 |
| 279 CheckHash(bb, new_block_buffer, image_base); |
| 280 } |
| 281 |
| 282 return true; |
| 283 } |
| 284 |
| 285 bool IntegrityCheckLayoutTransform::TransformImageLayout( |
| 286 const TransformPolicyInterface* policy, |
| 287 const pe::ImageLayout* image_layout, |
| 288 const OrderedBlockGraph* ordered_block_graph) { |
| 289 pe::PEFileWriter writer(*image_layout); |
| 290 |
| 291 if (!writer.ValidateHeaders()) |
| 292 return false; |
| 293 |
| 294 if (!writer.CalculateSectionRanges()) |
| 295 return false; |
| 296 |
| 297 core::AbsoluteAddress* image_base = writer.GetImageBase(); |
| 298 |
| 299 // Create the output buffer, reserving enough room for the whole file. |
| 300 DCHECK(!image_layout->sections.empty()); |
| 301 size_t image_size = writer.GetImageSize(); |
| 302 std::vector<uint8_t> buffer; |
| 303 buffer.reserve(image_size); |
| 304 |
| 305 // Iterate through all blocks in the address space writing them as we go. |
| 306 BlockGraph::AddressSpace::RangeMap::const_iterator block_it2( |
| 307 image_layout->blocks.address_space_impl().ranges().begin()); |
| 308 BlockGraph::AddressSpace::RangeMap::const_iterator block_end( |
| 309 image_layout->blocks.address_space_impl().ranges().end()); |
| 310 |
| 311 BlockGraph::AddressSpace::RangeMap::const_iterator block_it( |
| 312 image_layout->blocks.address_space_impl().ranges().begin()); |
| 313 |
| 314 BlockGraph::SectionId section_id = BlockGraph::kInvalidSectionId; |
| 315 size_t section_index = BlockGraph::kInvalidSectionId; |
| 316 |
| 317 // TODO: remove file |
| 318 phash = fopen("phash.txt", "w"); |
| 319 fprintf(phash, "Block name, BBid, Address, hash\n"); |
| 320 |
| 321 for (; block_it != block_end; ++block_it) { |
| 322 BlockGraph::Block* block = |
| 323 const_cast<BlockGraph::Block*>(block_it->second); |
| 324 |
| 325 // If we're jumping to a new section output the necessary padding. |
| 326 if (block->section() != section_id) { |
| 327 writer.FlushSection(section_index, &buffer); |
| 328 section_id = block->section(); |
| 329 section_index++; |
| 330 DCHECK_GT(image_layout->sections.size(), section_index); |
| 331 } |
| 332 |
| 333 core::FileOffsetAddress size_before(buffer.size()); |
| 334 |
| 335 if (!writer.WriteOneBlock(*image_base, section_index, block, |
| 336 &buffer, &size_before)) { |
| 337 LOG(ERROR) << "Failed to write block \"" << block->name() << "\"."; |
| 338 return false; |
| 339 } |
| 340 |
| 341 // compute new hash value of block |
| 342 auto buf_it = buffer.begin(); |
| 343 std::advance(buf_it, size_before.value()); |
| 344 std::vector<uint8_t> new_block_buffer(buf_it, buffer.end()); |
| 345 |
| 346 // Compute the new hash values inside buffer |
| 347 FixPrecomputedHashes(policy, *image_base, block, new_block_buffer); |
| 348 } // end for |
| 349 |
| 350 fclose(phash); |
| 351 |
| 352 block_it = image_layout->blocks.address_space_impl().ranges().begin(); |
| 353 for (; block_it != block_end; ++block_it) { |
| 354 BlockGraph::Block* block = |
| 355 const_cast<BlockGraph::Block*>(block_it->second); |
| 356 // patch the hash values in-place |
| 357 PatchPrecomputedHashes(policy, block); |
| 358 } |
| 359 |
| 360 return true; |
| 361 } |
| 362 |
| 363 // static vars |
| 364 const char IntegrityCheckLayoutTransform::kTransformName[] = |
| 365 "IntegrityCheckLayoutTransform"; |
| 366 } // namespace protect |
OLD | NEW |