dwarf_finder.cpp

#include <iostream>
#include <string>
#include <vector>
#include <stdexcept>
#include <elfutils/libdw.h>
#include <elfutils/libdwfl.h>
#include <dwarf.h>
#include <fcntl.h>
#include <unistd.h>

class DwarfStructOffsetFinder {
private:
    Dwarf* dwarf;

    bool find_struct_die(Dwarf_Die& result, const std::string& struct_name) {
        Dwarf_Off offset = 0;
        Dwarf_Off next_cu_offset;
        size_t cu_header_size;
        int cu_count = 0;

        while (dwarf_nextcu(dwarf, offset, &next_cu_offset, &cu_header_size, nullptr, nullptr, nullptr) == 0) {
            Dwarf_Die cu_die;
            if (dwarf_offdie(dwarf, offset + cu_header_size, &cu_die) != nullptr) {
                if (find_struct_die_in_cu(cu_die, result, struct_name)) {
                    return true;
                }
            }
            offset = next_cu_offset;
        }
        std::cout << "Searched " << cu_count << " compilation units." << std::endl;
        return false;
    }

    bool find_struct_die_in_cu(Dwarf_Die& die, Dwarf_Die& result, const std::string& struct_name) {
        int tag = dwarf_tag(&die);
        if (tag == DW_TAG_structure_type || tag == DW_TAG_class_type || tag == DW_TAG_union_type) {
            if (check_die_name(die, struct_name)) {
                result = die;
                return true;
            }
        } else if (tag == DW_TAG_typedef) {
            Dwarf_Attribute attr;
            Dwarf_Die type_die;
            if (dwarf_attr(&die, DW_AT_type, &attr) != nullptr &&
                dwarf_formref_die(&attr, &type_die) != nullptr) {
                if (check_die_name(die, struct_name)) {
                    // If the typedef name matches, check if it refers to a struct
                    int type_tag = dwarf_tag(&type_die);
                    if (type_tag == DW_TAG_structure_type || type_tag == DW_TAG_class_type || type_tag == DW_TAG_union_type) {
                        result = type_die;
                        return true;
                    }
                }
                // Recurse into the type die
                return find_struct_die_in_cu(type_die, result, struct_name);
            }
        }

        Dwarf_Die child;
        if (dwarf_child(&die, &child) == 0) {
            do {
                if (find_struct_die_in_cu(child, result, struct_name)) {
                    return true;
                }
            } while (dwarf_siblingof(&child, &child) == 0);
        }
        return false;
    }

    bool check_die_name(Dwarf_Die& die, const std::string& name) {
        const char* die_name = dwarf_diename(&die);
        if (die_name != nullptr) {
            if (name == die_name ||
                ("struct " + name) == die_name ||
                ("class " + name) == die_name ||
                ("union " + name) == die_name) {
                return true;
            }
        } else {
            Dwarf_Attribute attr;
            if (dwarf_attr(&die, DW_AT_name, &attr) != nullptr) {
                const char* attr_name = dwarf_formstring(&attr);
                if (attr_name != nullptr) {
                    std::cout << "Found type (from attribute): " << attr_name << std::endl;
                    if (name == attr_name ||
                        ("struct " + name) == attr_name ||
                        ("class " + name) == attr_name ||
                        ("union " + name) == attr_name) {
                        return true;
                    }
                }
            }
        }
        return false;
    }

    bool find_member_die(Dwarf_Die& result, Dwarf_Die& struct_die, const std::vector<std::string>& field_path, Dwarf_Off& offset) {

            Dwarf_Die current_die = struct_die;
            for (const auto& field_name : field_path) {
                Dwarf_Die child;
                bool found = false;

                if (dwarf_child(&current_die, &child) == 0) {
                    do {
                        if (dwarf_tag(&child) == DW_TAG_member) {
                            const char* name = dwarf_diename(&child);
                            if (name != nullptr && field_name == name) {
                                Dwarf_Attribute attr_mem;
                                Dwarf_Word member_offset;
                                if (dwarf_attr(&child, DW_AT_data_member_location, &attr_mem) != nullptr &&
                                    dwarf_formudata(&attr_mem, &member_offset) == 0) {
                                    offset += member_offset;
                                }

                                Dwarf_Die type_die;
                                Dwarf_Attribute attr_type;
                                if (dwarf_attr(&child, DW_AT_type, &attr_type) != nullptr &&
                                     dwarf_formref_die(&attr_type, &type_die) != nullptr) {
                                     current_die = type_die;
                                } else {
                                     current_die = child;
                                }
                                found = true;
                                break;
                            }
                        } else if (dwarf_tag(&child) == DW_TAG_structure_type) {
                            // Handle anonymous nested struct
                            if (dwarf_diename(&child) == nullptr) {
                                current_die = child;
                                found = true;
                                break;
                            }
                        }
                    } while (dwarf_siblingof(&child, &child) == 0);
                }
                 if (!found) {
                    return false;
                 }
            }

            result = current_die;
            return true;
        }

public:
    DwarfStructOffsetFinder(const std::string& filename) : dwarf(nullptr) {
        int fd = open(filename.c_str(), O_RDONLY);
        if (fd < 0) {
            throw std::runtime_error("Failed to open file: " + filename);
        }

        dwarf = dwarf_begin(fd, DWARF_C_READ);
        if (dwarf == nullptr) {
            close(fd);
            throw std::runtime_error("Failed to initialize DWARF: " + std::string(dwarf_errmsg(-1)));
        }
    }

    ~DwarfStructOffsetFinder() {
        if (dwarf != nullptr) {
            dwarf_end(dwarf);
        }
    }

    Dwarf_Off get_member_offset(const std::string& struct_name, const std::string& field_description) {
        std::vector<std::string> field_path;
        size_t start = 0, end = 0;
        while ((end = field_description.find('.', start)) != std::string::npos) {
            field_path.push_back(field_description.substr(start, end - start));
            start = end + 1;
        }
        field_path.push_back(field_description.substr(start));

        Dwarf_Die struct_die, member_die;
        if (!find_struct_die(struct_die, struct_name)) {
            throw std::runtime_error("Struct not found: " + struct_name);
        }

        Dwarf_Off offset = 0;
        if (!find_member_die(member_die, struct_die, field_path, offset)) {
            throw std::runtime_error("Member not found: " + field_description);
        }

        return offset;
    }
};

int main(int argc, char* argv[]) {
    if (argc != 4) {
        std::cerr << "Usage: " << argv[0] << " <executable> <struct_name> <field_description>" << std::endl;
        return 1;
    }

    try {
        DwarfStructOffsetFinder finder(argv[1]);
        Dwarf_Off offset = finder.get_member_offset(argv[2], argv[3]);
        std::cout << "Offset of " << argv[3] << " in struct " << argv[2] << ": " << offset << " bytes" << std::endl;
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return 1;
    }

    return 0;
}