brintos

brintos / llvm-project-archived public Read only

0
0
Text · 12.4 KiB · c5ab732 Raw
386 lines · cpp
1#include <getopt.h>2#include <mach-o/loader.h>3#include <stdio.h>4#include <stdlib.h>5#include <string.h>6#include <sys/stat.h>7 8#include <iostream>9#include <optional>10#include <string>11#include <vector>12 13using namespace std;14 15[[noreturn]] void print_help(void) {16  fprintf(stderr, "Append an LC_NOTE to a corefile.  Usage: \n");17  fprintf(stderr, "   -i|--input <corefile>\n");18  fprintf(stderr, "   -o|--output <corefile>\n");19  fprintf(stderr, "   -n|--name <LC_NOTE name>\n");20  fprintf(21      stderr,22      "   -r|--remove-dups  remove existing LC_NOTEs with this same name\n");23  fprintf(stderr, "  One of:\n");24  fprintf(stderr, "   -f|--file <file to embed as LC_NOTE payload>\n");25  fprintf(stderr, "   -s|--str <string to embed as LC_NOTE payload>\n");26  exit(1);27}28 29void parse_args(int argc, char **argv, string &infile, string &outfile,30                string &note_name, vector<uint8_t> &payload,31                bool &remove_dups) {32  const char *const short_opts = "i:o:n:f:s:hr";33  const option long_opts[] = {{"input", required_argument, nullptr, 'i'},34                              {"output", required_argument, nullptr, 'o'},35                              {"name", required_argument, nullptr, 'n'},36                              {"file", required_argument, nullptr, 'f'},37                              {"str", required_argument, nullptr, 's'},38                              {"remove-dups", no_argument, nullptr, 'r'},39                              {"help", no_argument, nullptr, 'h'},40                              {nullptr, no_argument, nullptr, 0}};41  optional<string> infile_str, outfile_str, name_str, payload_file_str,42      payload_str;43  remove_dups = false;44  while (true) {45    const auto opt = getopt_long(argc, argv, short_opts, long_opts, nullptr);46    if (opt == -1)47      break;48    switch (opt) {49    case 'i':50      infile_str = optarg;51      break;52    case 'o':53      outfile_str = optarg;54      break;55    case 'n':56      name_str = optarg;57      break;58    case 'f':59      payload_file_str = optarg;60      break;61    case 's':62      payload_str = optarg;63      break;64    case 'r':65      remove_dups = true;66      break;67    case 'h':68      print_help();69    }70  }71 72  if (!infile_str || !outfile_str || !name_str ||73      (!payload_file_str && !payload_str))74    print_help();75 76  infile = *infile_str;77  outfile = *outfile_str;78  note_name = *name_str;79  if (payload_str) {80    payload.resize(payload_str->size(), 0);81    memcpy(payload.data(), payload_str->c_str(), payload_str->size());82  } else {83    struct stat sb;84    if (stat(payload_file_str->c_str(), &sb)) {85      fprintf(stderr, "File '%s' does not exist.\n", payload_file_str->c_str());86      exit(1);87    }88    payload.resize(sb.st_size, 0);89    FILE *f = fopen(payload_file_str->c_str(), "r");90    fread(payload.data(), 1, sb.st_size, f);91    fclose(f);92  }93}94 95struct all_image_infos_header {96  uint32_t version;         // currently 197  uint32_t imgcount;        // number of binary images98  uint64_t entries_fileoff; // file offset in the corefile of where the array of99                            // struct entry's begin.100  uint32_t entry_size;      // size of 'struct entry'.101  uint32_t unused;          // set to 0102};103 104struct image_entry {105  uint64_t filepath_offset; // corefile offset of the c-string filepath,106                            // if available, else this should be set107                            // to UINT64_MAX.108  uuid_t uuid;              // uint8_t[16].  should be set to all zeroes if109                            // uuid is unknown.110  uint64_t111      load_address; // virtual addr of mach-o header, UINT64_MAX if unknown.112  uint64_t seg_addrs_offset; // corefile offset to the array of struct113                             // segment_vmaddr's, UINT64_MAX if none.114  uint32_t segment_count; // The number of segments for this binary, 0 if none.115  uint32_t116      executing; // Set to 0 if executing status is unknown by corefile117                 // creator.118                 // Set to 1 if this binary was executing on any thread,119                 // so it can be force-loaded by the corefile reader.120                 // Set to 2 if this binary was not executing on any thread.121};122 123int count_lc_notes_with_name(FILE *in, std::string name) {124  fseeko(in, 0, SEEK_SET);125 126  uint8_t magic[4];127  if (fread(magic, 1, 4, in) != 4) {128    printf("Failed to read magic number\n");129    return 0;130  }131  uint8_t magic_32_le[] = {0xce, 0xfa, 0xed, 0xfe};132  uint8_t magic_64_le[] = {0xcf, 0xfa, 0xed, 0xfe};133 134  if (memcmp(magic, magic_32_le, 4) != 0 &&135      memcmp(magic, magic_64_le, 4) != 0) {136    return 0;137  }138 139  fseeko(in, 0, SEEK_SET);140 141  int number_of_load_cmds = 0;142  size_t size_of_mach_header = 0;143  if (memcmp(magic, magic_64_le, 4) == 0) {144    struct mach_header_64 mh;145    size_of_mach_header = sizeof(mh);146    if (fread(&mh, sizeof(mh), 1, in) != 1) {147      fprintf(stderr, "unable to read mach header\n");148      return 0;149    }150    number_of_load_cmds = mh.ncmds;151  } else {152    struct mach_header mh;153    size_of_mach_header = sizeof(mh);154    if (fread(&mh, sizeof(mh), 1, in) != 1) {155      fprintf(stderr, "unable to read mach header\n");156      return 0;157    }158    number_of_load_cmds = mh.ncmds;159  }160 161  int notes_seen = 0;162  fseeko(in, size_of_mach_header, SEEK_SET);163  for (int i = 0; i < number_of_load_cmds; i++) {164    off_t cmd_start = ftello(in);165    uint32_t cmd, cmdsize;166    fread(&cmd, sizeof(uint32_t), 1, in);167    fread(&cmdsize, sizeof(uint32_t), 1, in);168 169    fseeko(in, cmd_start, SEEK_SET);170    off_t next_cmd = cmd_start + cmdsize;171    if (cmd == LC_NOTE) {172      struct note_command note;173      fread(&note, sizeof(note), 1, in);174      if (strncmp(name.c_str(), note.data_owner, 16) == 0)175        notes_seen++;176    }177    fseeko(in, next_cmd, SEEK_SET);178  }179  return notes_seen;180}181 182void copy_and_add_note(FILE *in, FILE *out, std::string lc_note_name,183                       vector<uint8_t> payload_data, bool remove_dups) {184  int number_of_load_cmds = 0;185  off_t header_start = ftello(in);186 187  int notes_to_remove = 0;188  if (remove_dups)189    notes_to_remove = count_lc_notes_with_name(in, lc_note_name);190  fseeko(in, header_start, SEEK_SET);191 192  uint8_t magic[4];193  if (fread(magic, 1, 4, in) != 4) {194    printf("Failed to read magic number\n");195    return;196  }197  uint8_t magic_32_le[] = {0xce, 0xfa, 0xed, 0xfe};198  uint8_t magic_64_le[] = {0xcf, 0xfa, 0xed, 0xfe};199 200  if (memcmp(magic, magic_32_le, 4) != 0 &&201      memcmp(magic, magic_64_le, 4) != 0) {202    return;203  }204 205  fseeko(in, header_start, SEEK_SET);206 207  off_t end_of_infine_loadcmds;208  size_t size_of_mach_header = 0;209  if (memcmp(magic, magic_64_le, 4) == 0) {210    struct mach_header_64 mh;211    size_of_mach_header = sizeof(mh);212    if (fread(&mh, sizeof(mh), 1, in) != 1) {213      fprintf(stderr, "unable to read mach header\n");214      return;215    }216    number_of_load_cmds = mh.ncmds;217    end_of_infine_loadcmds = sizeof(mh) + mh.sizeofcmds;218 219    mh.ncmds += 1;220    mh.ncmds -= notes_to_remove;221    mh.sizeofcmds += sizeof(struct note_command);222    mh.sizeofcmds -= notes_to_remove * sizeof(struct note_command);223    fseeko(out, header_start, SEEK_SET);224    fwrite(&mh, sizeof(mh), 1, out);225  } else {226    struct mach_header mh;227    size_of_mach_header = sizeof(mh);228    if (fread(&mh, sizeof(mh), 1, in) != 1) {229      fprintf(stderr, "unable to read mach header\n");230      return;231    }232    number_of_load_cmds = mh.ncmds;233    end_of_infine_loadcmds = sizeof(mh) + mh.sizeofcmds;234 235    mh.ncmds += 1;236    mh.ncmds -= notes_to_remove;237    mh.sizeofcmds += sizeof(struct note_command);238    mh.sizeofcmds -= notes_to_remove * sizeof(struct note_command);239    fseeko(out, header_start, SEEK_SET);240    fwrite(&mh, sizeof(mh), 1, out);241  }242 243  off_t start_of_infile_load_cmds = ftello(in);244  fseek(in, 0, SEEK_END);245  off_t infile_size = ftello(in);246 247  // LC_SEGMENT may be aligned to 4k boundaries, let's maintain248  // that alignment by putting 4096 minus the size of the added249  // LC_NOTE load command after the output file's load commands.250  off_t end_of_outfile_loadcmds =251      end_of_infine_loadcmds - (notes_to_remove * sizeof(struct note_command)) +252      4096 - sizeof(struct note_command);253  off_t slide = end_of_outfile_loadcmds - end_of_infine_loadcmds;254 255  off_t all_image_infos_infile_offset = 0;256 257  fseek(in, start_of_infile_load_cmds, SEEK_SET);258  fseek(out, start_of_infile_load_cmds, SEEK_SET);259  // Copy all the load commands from IN to OUT, updating any file offsets by260  // SLIDE.261  for (int cmd_num = 0; cmd_num < number_of_load_cmds; cmd_num++) {262    off_t cmd_start = ftello(in);263    uint32_t cmd, cmdsize;264    fread(&cmd, sizeof(uint32_t), 1, in);265    fread(&cmdsize, sizeof(uint32_t), 1, in);266 267    fseeko(in, cmd_start, SEEK_SET);268    off_t next_cmd = cmd_start + cmdsize;269 270    switch (cmd) {271    case LC_SEGMENT: {272      struct segment_command segcmd;273      fread(&segcmd, sizeof(segcmd), 1, in);274      segcmd.fileoff += slide;275      fwrite(&segcmd, cmdsize, 1, out);276    } break;277    case LC_SEGMENT_64: {278      struct segment_command_64 segcmd;279      fread(&segcmd, sizeof(segcmd), 1, in);280      segcmd.fileoff += slide;281      fwrite(&segcmd, cmdsize, 1, out);282    } break;283    case LC_NOTE: {284      struct note_command notecmd;285      fread(&notecmd, sizeof(notecmd), 1, in);286      if ((strncmp(lc_note_name.c_str(), notecmd.data_owner, 16) == 0) &&287          remove_dups) {288        fseeko(in, next_cmd, SEEK_SET);289        continue;290      }291      if (strncmp("all image infos", notecmd.data_owner, 16) == 0)292        all_image_infos_infile_offset = notecmd.offset;293      notecmd.offset += slide;294      fwrite(&notecmd, cmdsize, 1, out);295    } break;296    default: {297      vector<uint8_t> buf(cmdsize);298      fread(buf.data(), cmdsize, 1, in);299      fwrite(buf.data(), cmdsize, 1, out);300    }301    }302    fseeko(in, next_cmd, SEEK_SET);303  }304 305  // Now add our additional LC_NOTE load command.306  struct note_command note;307  note.cmd = LC_NOTE;308  note.cmdsize = sizeof(struct note_command);309  memset(&note.data_owner, 0, 16);310  // data_owner may not be nul terminated if all 16 characters311  // are used, intentionally using strncpy here.312  strncpy(note.data_owner, lc_note_name.c_str(), 16);313  note.offset = infile_size + slide;314  note.size = payload_data.size();315  fwrite(&note, sizeof(struct note_command), 1, out);316 317  fseeko(in, end_of_infine_loadcmds, SEEK_SET);318  fseeko(out, end_of_outfile_loadcmds, SEEK_SET);319 320  // Copy the rest of the corefile contents321  vector<uint8_t> data_buf(1024 * 1024);322  while (!feof(in)) {323    size_t read_bytes = fread(data_buf.data(), 1, data_buf.size(), in);324    if (read_bytes > 0) {325      fwrite(data_buf.data(), read_bytes, 1, out);326    } else {327      break;328    }329  }330 331  fwrite(payload_data.data(), payload_data.size(), 1, out);332 333  // The "all image infos" LC_NOTE payload has file offsets hardcoded334  // in it, unfortunately.  We've shifted the contents of the corefile335  // and these offsets need to be updated in the ouput file.336  // Re-copy them into the outfile with corrected file offsets.337  off_t infile_image_entry_base = 0;338  if (all_image_infos_infile_offset != 0) {339    off_t all_image_infos_outfile_offset =340        all_image_infos_infile_offset + slide;341    fseeko(in, all_image_infos_infile_offset, SEEK_SET);342    struct all_image_infos_header header;343    fread(&header, sizeof(header), 1, in);344    infile_image_entry_base = header.entries_fileoff;345    header.entries_fileoff += slide;346    fseeko(out, all_image_infos_outfile_offset, SEEK_SET);347    fwrite(&header, sizeof(header), 1, out);348 349    for (int i = 0; i < header.imgcount; i++) {350      off_t infile_entries_fileoff = header.entries_fileoff - slide;351      off_t outfile_entries_fileoff = header.entries_fileoff;352 353      struct image_entry ent;354      fseeko(in, infile_entries_fileoff + (header.entry_size * i), SEEK_SET);355      fread(&ent, sizeof(ent), 1, in);356      ent.filepath_offset += slide;357      ent.seg_addrs_offset += slide;358      fseeko(out, outfile_entries_fileoff + (header.entry_size * i), SEEK_SET);359      fwrite(&ent, sizeof(ent), 1, out);360    }361  }362}363 364int main(int argc, char **argv) {365  string infile, outfile, name;366  vector<uint8_t> payload;367  bool remove_dups;368  parse_args(argc, argv, infile, outfile, name, payload, remove_dups);369 370  FILE *in = fopen(infile.c_str(), "r");371  if (!in) {372    fprintf(stderr, "Unable to open %s for reading\n", infile.c_str());373    exit(1);374  }375  FILE *out = fopen(outfile.c_str(), "w");376  if (!out) {377    fprintf(stderr, "Unable to open %s for reading\n", outfile.c_str());378    exit(1);379  }380 381  copy_and_add_note(in, out, name, payload, remove_dups);382 383  fclose(in);384  fclose(out);385}386