PatchSolution.cpp 2.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162
  1. #include "PatchSolutions.hpp"
  2. namespace nkg {
  3. void PatchSolution::SearchFreeSpace(std::map<Elf64_Off, Elf64_Xword>& SpaceMap, const Elf64Interpreter& Image) {
  4. static auto lpfnUpdateMap = [](std::map<Elf64_Off, Elf64_Xword>& SpaceMap, Elf64_Off offset, Elf64_Xword size) {
  5. auto start = offset;
  6. auto end = start + size;
  7. while (size) {
  8. auto space = SpaceMap.upper_bound(offset);
  9. auto space_start = space->first;
  10. auto space_end = space_start + space->second;
  11. if (space != SpaceMap.end() && space_start < end) { // implicit condition: start < space_start
  12. if (space_end <= end) {
  13. SpaceMap.erase(space);
  14. } else {
  15. auto node = SpaceMap.extract(space);
  16. node.key() = end;
  17. SpaceMap.insert(std::move(node));
  18. }
  19. } else if (space != SpaceMap.begin()) {
  20. --space;
  21. space_start = space->first; // space_start <= start
  22. space_end = space_start + space->second;
  23. if (start < space_end) {
  24. space->second = start - space_start;
  25. if (space->second == 0) {
  26. SpaceMap.erase(space);
  27. }
  28. if (end < space_end) {
  29. SpaceMap.emplace(end, space_end - end);
  30. }
  31. } else {
  32. break;
  33. }
  34. } else {
  35. break;
  36. }
  37. }
  38. };
  39. lpfnUpdateMap(SpaceMap, 0, sizeof(Elf64_Ehdr));
  40. lpfnUpdateMap(SpaceMap, Image.ElfHeader()->e_phoff, Image.ElfHeader()->e_phentsize * Image.ElfHeader()->e_phnum);
  41. lpfnUpdateMap(SpaceMap, Image.ElfHeader()->e_shoff, Image.ElfHeader()->e_shentsize * Image.ElfHeader()->e_shnum);
  42. for (size_t i = 0; i < Image.NumberOfElfProgramHeaders(); ++i) {
  43. auto seg_hdr = Image.ElfProgramHeader(i);
  44. if (seg_hdr->p_type != PT_NULL) {
  45. lpfnUpdateMap(SpaceMap, seg_hdr->p_offset, seg_hdr->p_filesz);
  46. }
  47. }
  48. for (size_t i = 0; i < Image.NumberOfElfSectionHeaders(); ++i) {
  49. auto sec_hdr = Image.ElfSectionHeader(i);
  50. if (sec_hdr->sh_type != SHT_NULL && sec_hdr->sh_type != SHT_NOBITS) {
  51. lpfnUpdateMap(SpaceMap, sec_hdr->sh_offset, sec_hdr->sh_size);
  52. }
  53. }
  54. }
  55. }