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 ¬e_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(¬e, 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(¬ecmd, 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(¬ecmd, 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(¬e.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(¬e, 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