// ElfFile.cpp
//------------------------------------------------------------------------------
//      Copyright (c) 2003, Ingo Weinhold
//
//      Permission is hereby granted, free of charge, to any person obtaining a
//      copy of this software and associated documentation files (the "Software"),
//      to deal in the Software without restriction, including without limitation
//      the rights to use, copy, modify, merge, publish, distribute, sublicense,
//      and/or sell copies of the Software, and to permit persons to whom the
//      Software is furnished to do so, subject to the following conditions:
//
//      The above copyright notice and this permission notice shall be included in
//      all copies or substantial portions of the Software.
//
//      THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//      IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//      FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//      AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//      LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
//      FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
//      DEALINGS IN THE SOFTWARE.
//
//      File Name:              ElfFile.cpp
//      Author:                 Ingo Weinhold (bonefish@users.sf.net)
//      Description:    Implementation of classes for accessing ELF file,
//                                      or more precisely for iterating through their relocation
//                                      sections.
//------------------------------------------------------------------------------

#include <new>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "ElfFile.h"

// sanity bounds
static const uint32     kMaxELFHeaderSize                       = sizeof(Elf_Ehdr) + 32;

// read_exactly
static
status_t
read_exactly(BPositionIO &file, off_t position, void *buffer, size_t size,
                         const char *errorMessage = NULL)
{
        status_t error = B_OK;
        ssize_t read = file.ReadAt(position, buffer, size);
        if (read < 0)
                error = read;
        else if ((size_t)read != size)
                error = B_ERROR;
        if (error != B_OK && errorMessage)
                puts(errorMessage);
        return error;
}


// ElfSection

class SymbolPatcher::ElfSection {
public:
                                                                ElfSection();
                                                                ~ElfSection();

                        void                            SetTo(ElfFile* file, Elf_Shdr* header);
                        void                            Unset();
                        bool                            IsInitialized() const   { return fHeader; }

                        ElfFile*                        GetFile() const;
                        Elf_Shdr*                       GetHeader() const               { return fHeader; }
                        const char*                     GetName() const;
                        uint8*                          GetData() const                 { return fData; }
                        size_t                          GetSize() const;
                        Elf_Word                        GetType() const;
                        Elf_Word                        GetLink() const;
                        Elf_Word                        GetInfo() const;
                        size_t                          GetEntrySize() const;
                        int32                           CountEntries() const;

                        status_t                        Load();
                        void                            Unload();

                        void                            Dump();

private:
                        ElfFile*                        fFile;
                        Elf_Shdr*                       fHeader;
                        uint8*                          fData;
};

// constructor
ElfSection::ElfSection()
        : fFile(NULL),
          fHeader(NULL),
          fData(NULL)
{
}

// destructor
ElfSection::~ElfSection()
{
        Unset();
}

// SetTo
void
ElfSection::SetTo(ElfFile* file, Elf_Shdr* header)
{
        Unset();
        fFile = file;
        fHeader = header;
}

// Unset
void
ElfSection::Unset()
{
        Unload();
        fFile = NULL;
        fHeader = NULL;
}

// GetFile
ElfFile*
ElfSection::GetFile() const
{
        return fFile;
}

// GetName
const char*
ElfSection::GetName() const
{
        const char* name = NULL;
        if (fHeader && fFile) {
                size_t size = 0;
                const char* nameSection = fFile->GetSectionHeaderStrings(&size);
                if (nameSection && fHeader->sh_name < size)
                        name = nameSection + fHeader->sh_name;
        }
        return name;
}

// GetSize
size_t
ElfSection::GetSize() const
{
        return fHeader->sh_size;
}

// GetType
Elf_Word
ElfSection::GetType() const
{
        return fHeader->sh_type;
}

// GetLink
Elf_Word
ElfSection::GetLink() const
{
        return fHeader->sh_link;
}

// GetInfo
Elf_Word
ElfSection::GetInfo() const
{
        return fHeader->sh_info;
}

// GetEntrySize
size_t
ElfSection::GetEntrySize() const
{
        return fHeader->sh_entsize;
}

// CountEntries
int32
ElfSection::CountEntries() const
{
        int32 count = 0;
        if (fHeader) {
                if (GetEntrySize() == 0)
                        return 0;
                count = GetSize() / GetEntrySize();
        }
        return count;
}

// Load
status_t
ElfSection::Load()
{
        status_t error = B_ERROR;
        if (fHeader && !fData && fHeader->sh_type != SHT_NULL
                 && fHeader->sh_type != SHT_NOBITS) {
                BFile* file = fFile->GetFile();
                // allocate memory
                fData = new uint8[fHeader->sh_size];
                if (!fData)
                        return B_NO_MEMORY;
                // read the data
                error = read_exactly(*file, fHeader->sh_offset, fData,
                                                         fHeader->sh_size, "Failed to read section!\n");
                if (error != B_OK)
                        Unload();
        }
        return error;
}

// Unload
void
ElfSection::Unload()
{
        if (fData) {
                delete[] fData;
                fData = NULL;
        }
}

// Dump
void
ElfSection::Dump()
{
        printf("section %32s: size: %lu\n", GetName(), GetSize());
}


// ElfSymbol

// constructor
ElfSymbol::ElfSymbol(ElfSection* section, int32 index)
        : fSection(section),
          fIndex(index),
          fSymbol(NULL)
{
}

// destructor
ElfSymbol::~ElfSymbol()
{
        Unset();
}

// SetTo
void
ElfSymbol::SetTo(ElfSection* section, int32 index)
{
        Unset();
        fSection = section;
        fIndex = index;
}

// Unset
void
ElfSymbol::Unset()
{
        fSection = NULL;
        fIndex = -1;
        fSymbol = NULL;
}

// GetSymbolStruct
const Elf_Sym*
ElfSymbol::GetSymbolStruct()
{
        Elf_Sym* symbol = fSymbol;
        if (!symbol && fSection && fSection->GetData()) {
                size_t symbolSize = fSection->GetEntrySize();
                if (symbolSize == 0)
                        return NULL;
                int32 symbolCount = fSection->GetSize() / symbolSize;
                if (fIndex >= 0 && fIndex < symbolCount)
                        symbol = (Elf_Sym*)(fSection->GetData() + fIndex * symbolSize);
        }
        return symbol;
}

// GetName
const char*
ElfSymbol::GetName()
{
        const char* name = NULL;
        if (const Elf_Sym* symbol = GetSymbolStruct()) {
                size_t size = 0;
                const char* data = fSection->GetFile()->GetStringSectionStrings(
                        fSection->GetLink(), &size);
                if (data && symbol->st_name < size)
                        name = data + symbol->st_name;
        }
        return name;
}

// GetBinding
uint32
ElfSymbol::GetBinding()
{
        uint32 binding = STB_LOCAL;
        if (const Elf_Sym* symbol = GetSymbolStruct())
                binding = ELF_ST_BIND(symbol->st_info);
        return binding;
}

// GetType
uint32
ElfSymbol::GetType()
{
        uint32 type = STT_NOTYPE;
        if (const Elf_Sym* symbol = GetSymbolStruct())
                type = ELF_ST_TYPE(symbol->st_info);
        return type;
}

// GetTargetSectionIndex
uint32
ElfSymbol::GetTargetSectionIndex()
{
        uint32 index = SHN_UNDEF;
        if (const Elf_Sym* symbol = GetSymbolStruct())
                index = symbol->st_shndx;
        return index;
}


// ElfRelocation

// constructor
ElfRelocation::ElfRelocation(ElfSection* section, int32 index)
        : fSection(section),
          fIndex(index),
          fRelocation(NULL)
{
}

// destructor
ElfRelocation::~ElfRelocation()
{
        Unset();
}

// SetTo
void
ElfRelocation::SetTo(ElfSection* section, int32 index)
{
        Unset();
        fSection = section;
        fIndex = index;
}

// Unset
void
ElfRelocation::Unset()
{
        fSection = NULL;
        fIndex = -1;
        fRelocation = NULL;
}

// GetRelocationStruct
Elf_Rel*
ElfRelocation::GetRelocationStruct()
{
        Elf_Rel* relocation = fRelocation;
        if (!relocation && fSection) {
                if (!fSection->GetData()) {
                        if (fSection->Load() != B_OK)
                                return NULL;
                }
                size_t entrySize = fSection->GetEntrySize();
                if (entrySize == 0 || entrySize < sizeof(Elf_Rel))
                        return NULL;
                int32 entryCount = fSection->GetSize() / entrySize;
                if (fIndex >= 0 && fIndex < entryCount) {
                        relocation = (Elf_Rel*)(fSection->GetData()
                                                 + fIndex * entrySize);
                }
        }
        return relocation;
}

// GetType
uint32
ElfRelocation::GetType()
{
        uint32 type = R_NONE;
        if (Elf_Rel* relocation = GetRelocationStruct())
                type = ELF_R_TYPE(relocation->r_info);
        return type;
}

// GetSymbolIndex
uint32
ElfRelocation::GetSymbolIndex()
{
        uint32 index = 0;
        if (Elf_Rel* relocation = GetRelocationStruct())
                index = ELF_R_SYM(relocation->r_info);
        return index;
}

// GetOffset
Elf_Addr
ElfRelocation::GetOffset()
{
        Elf_Addr offset = 0;
        if (Elf_Rel* relocation = GetRelocationStruct())
                offset = relocation->r_offset;
        return offset;
}

// GetSymbol
status_t
ElfRelocation::GetSymbol(ElfSymbol* symbol)
{
        status_t error = B_BAD_VALUE;
        if (symbol && fSection) {
                uint32 index = GetSymbolIndex();
                if (ElfSection* symbols
                        = fSection->GetFile()->SectionAt(fSection->GetLink(), true)) {
                        symbol->SetTo(symbols, index);
                        if (symbol->GetSymbolStruct())
                                error = B_OK;
                }
        }
        return error;
}


// ElfRelocationIterator

// constructor
ElfRelocationIterator::ElfRelocationIterator(ElfFile* file)
        : fFile(file),
          fSectionIndex(-1),
          fEntryIndex(-1)
{
}

// destructor
ElfRelocationIterator::~ElfRelocationIterator()
{
}

// GetNext
bool
ElfRelocationIterator::GetNext(ElfRelocation* relocation)
{
        bool result = false;
        if (fFile && relocation) {
                // set to possible entry
                ElfSection* section = NULL;
                if (fSectionIndex < 0) {
                        fSectionIndex = 0;
                        fEntryIndex = 0;
                        section = _FindNextSection();
                } else {
                        fEntryIndex++;
                        section = fFile->SectionAt(fSectionIndex);
                }
                // find next valid entry
                while (section && fEntryIndex >= section->CountEntries()) {
                        fSectionIndex++;
                        section = _FindNextSection();
                        fEntryIndex = 0;
                }
                // set result
                if (section) {
                        relocation->SetTo(section, fEntryIndex);
                        result = true;
                }
        }
        return result;
}

// _FindNextSection
ElfSection*
ElfRelocationIterator::_FindNextSection()
{
        if (fFile) {
                for (; fSectionIndex < fFile->CountSections(); fSectionIndex++) {
                        ElfSection* section = fFile->SectionAt(fSectionIndex);
                        if (section && (section->GetType() == SHT_REL || section->GetType() == SHT_RELA))
                                return section;
                }
        }
        return NULL;
}


// ElfFile

// constructor
ElfFile::ElfFile()
        : fFile(),
          fSectionHeaders(NULL),
          fSections(NULL),
          fSectionCount(0),
          fSectionHeaderSize(0)
{
}

// destructor
ElfFile::~ElfFile()
{
        Unset();
}

// SetTo
status_t
ElfFile::SetTo(const char *filename)
{
        Unset();
        status_t error = _SetTo(filename);
        if (error)
                Unset();
        return error;
}

// Unset
void
ElfFile::Unset()
{
        // delete sections
        if (fSections) {
                delete[] fSections;
                fSections = NULL;
        }
        // delete section headers
        if (fSectionHeaders) {
                delete[] fSectionHeaders;
                fSectionHeaders = NULL;
        }
        fSectionCount = 0;
        fSectionHeaderSize = 0;
        fFile.Unset();
}

// Unload
void
ElfFile::Unload()
{
        for (int i = 0; i < fSectionCount; i++)
                fSections[i].Unload();
}

// GetSectionHeaderStrings
const char*
ElfFile::GetSectionHeaderStrings(size_t* size)
{
        return GetStringSectionStrings(fHeader.e_shstrndx, size);
}

// GetStringSectionStrings
const char*
ElfFile::GetStringSectionStrings(int32 index, size_t* _size)
{
        const char* data = NULL;
        size_t size = 0;
        if (ElfSection* section = SectionAt(index, true)) {
                data = (const char*)section->GetData();
                size = (data ? section->GetSize() : 0);
        }
        // set results
        if (_size)
                *_size = size;
        return data;
}

// SectionAt
ElfSection*
ElfFile::SectionAt(int32 index, bool load)
{
        ElfSection* section = NULL;
        if (fSections && index >= 0 && index < fSectionCount) {
                section = fSections + index;
                if (load && !section->GetData()) {
                        if (section->Load() != B_OK) {
                                section = NULL;
                                printf("Failed to load section %" B_PRId32 "\n", index);
                        }
                }
        }
        return section;
}

// Dump
void
ElfFile::Dump()
{
        printf("%" B_PRId32 " sections\n", fSectionCount);
        for (int i = 0; i < fSectionCount; i++)
                fSections[i].Dump();
}

// _SetTo
status_t
ElfFile::_SetTo(const char *filename)
{
        if (!filename)
                return B_BAD_VALUE;
        // open file
        status_t error = fFile.SetTo(filename, B_READ_ONLY);
        // get the file size
        off_t fileSize = 0;
        error = fFile.GetSize(&fileSize);
        if (error != B_OK) {
                printf("Failed to get file size!\n");
                return error;
        }
        // read ELF header
        error = read_exactly(fFile, 0, &fHeader, sizeof(Elf_Ehdr),
                                                 "Failed to read ELF object header!\n");
        if (error != B_OK)
                return error;
        // check the ident field
        // magic
        if (fHeader.e_ident[EI_MAG0] != ELFMAG0
                || fHeader.e_ident[EI_MAG1] != ELFMAG1
                || fHeader.e_ident[EI_MAG2] != ELFMAG2
                || fHeader.e_ident[EI_MAG3] != ELFMAG3) {
                printf("Bad ELF file magic!\n");
                return B_BAD_VALUE;
        }
        // class
        if (fHeader.e_ident[EI_CLASS] != ELFCLASS) {
                printf("Wrong ELF class!\n");
                return B_BAD_VALUE;
        }
        // check data encoding (endianess)
        if (fHeader.e_ident[EI_DATA] != ELFDATA2LSB) {
                printf("Wrong data encoding!\n");
                return B_BAD_VALUE;
        }
        // version
        if (fHeader.e_ident[EI_VERSION] != EV_CURRENT) {
                printf("Wrong data encoding!\n");
                return B_BAD_VALUE;
        }
        // get the header values
        uint32 headerSize                               = fHeader.e_ehsize;
        uint32 sectionHeaderTableOffset = fHeader.e_shoff;
        uint32 sectionHeaderSize                = fHeader.e_shentsize;
        uint32 sectionHeaderCount               = fHeader.e_shnum;
        // check the sanity of the header values
        // ELF header size
        if (headerSize < sizeof(Elf_Ehdr) || headerSize > kMaxELFHeaderSize) {
                printf("Invalid ELF header: invalid ELF header size: %" B_PRIu32 ".", headerSize);
                return B_BAD_VALUE;
        }
        // section header table offset
        if (sectionHeaderTableOffset == 0) {
                printf("ELF file has no section header table!\n");
                return B_BAD_VALUE;
        }
        uint32 sectionHeaderTableSize = 0;
        if (sectionHeaderTableOffset < headerSize
                || sectionHeaderTableOffset > fileSize) {
                printf("Invalid ELF header: invalid section header table offset: %" B_PRIu32 ".",
                           sectionHeaderTableOffset);
                return B_BAD_VALUE;
        }
        // section header table offset
        sectionHeaderTableSize = sectionHeaderSize * sectionHeaderCount;
        if (sectionHeaderSize < sizeof(Elf_Shdr)
                || sectionHeaderTableOffset + sectionHeaderTableSize > fileSize) {
                printf("Invalid ELF header: section header table exceeds file: %" B_PRIu32 ".",
                           sectionHeaderTableOffset + sectionHeaderTableSize);
                return B_BAD_VALUE;
        }
        // allocate memory for the section header table and read it
        fSectionHeaders = new(std::nothrow) uint8[sectionHeaderTableSize];
        fSectionCount = sectionHeaderCount;
        fSectionHeaderSize = sectionHeaderSize;
        if (!fSectionHeaders)
                return B_NO_MEMORY;
        error = read_exactly(fFile, sectionHeaderTableOffset, fSectionHeaders,
                                                 sectionHeaderTableSize,
                                                 "Failed to read section headers!\n");
        if (error != B_OK)
                return error;
        // allocate memory for the section pointers
        fSections = new(std::nothrow) ElfSection[fSectionCount];
        if (!fSections)
                return B_NO_MEMORY;
        // init the sections
        for (int i = 0; i < fSectionCount; i++)
                fSections[i].SetTo(this, _SectionHeaderAt(i));
        return error;
}

// _SectionHeaderAt
Elf_Shdr*
ElfFile::_SectionHeaderAt(int32 index)
{
        Elf_Shdr* header = NULL;
        if (fSectionHeaders && index >= 0 && index < fSectionCount)
                header = (Elf_Shdr*)(fSectionHeaders + index * fSectionHeaderSize);
        return header;
}

// _LoadSection
status_t
ElfFile::_LoadSection(int32 index)
{
        status_t error = B_OK;
        if (fSections && index >= 0 && index < fSectionCount) {
                ElfSection& section = fSections[index];
                error = section.Load();
        } else
                error = B_BAD_VALUE;
        return error;
}