mzexe.h

#ifndef MZEXE_H
#define MZEXE_H
 
#include <array>
#include <iomanip>
#include <set>
#include <string>
#include <vector>
#include "../common.h"
#include "../linker/format.h"
#include "../linker/module.h"
#include "../linker/options.h"
#include "../linker/reader.h"
#include "../linker/section.h"
#include "../linker/segment.h"
#include "../linker/segment_manager.h"
#include "../linker/writer.h"
#include "../dumper/dumper.h"
 
namespace Microsoft
{
    class MZFormat : public virtual Linker::SegmentManager
    {
    public:
        /* * * General members * * */
 
        enum magic_type
        {
            MAGIC_MZ = 1,
            MAGIC_ZM,
            MAGIC_DL,
        };
 
        char signature[2] = { 'M', 'Z' }; /* TODO: make parameter */
 
        uint16_t last_block_size = 0;
        uint16_t file_size_blocks = 0; /* TODO: consider making file size a parameter */
 
        offset_t ImageSize() const override;
 
        uint16_t relocation_count = 0;
        uint16_t header_size_paras = 0; /* TODO: make header size a parameter */
        uint16_t min_extra_paras = 0;
        uint16_t max_extra_paras = 0;
        uint16_t ss = 0;
        uint16_t sp = 0;
        uint16_t checksum = 0; /* TODO: fill when writing */
        uint16_t ip = 0;
        uint16_t cs = 0;
        uint16_t relocation_offset = 0; /* TODO: make parameter */
 
        uint16_t overlay_number = 0;
        uint16_t data_segment = 0;
 
        uint16_t behavior_bits = 0;
        uint16_t oem_id = 0;
        uint16_t oem_info = 0;
        uint32_t new_header_offset = 0;
        uint32_t win386_new_header_offset = 0;
 
        struct Relocation
        {
            uint16_t segment;
            uint16_t offset;
 
            Relocation(uint16_t segment, uint16_t offset)
                : segment(segment), offset(offset)
            {
            }
 
            static Relocation FromLinear(uint32_t address);
 
            uint32_t GetOffset() const;
 
            bool operator ==(const Relocation& other) const;
 
            bool operator <(const Relocation& other) const;
        };
 
        std::vector<Relocation> relocations;
 
        struct PIF
        {
            static constexpr uint32_t MAGIC_BEGIN = 0x0013EDC1;
            static constexpr uint32_t MAGIC_END = 0xEDC10013;
            static constexpr size_t SIZE = 19;
 
            /* TODO: requires testing */
            uint16_t maximum_extra_paragraphs;
            uint16_t minimum_extra_paragraphs;
            uint8_t flags;
            uint8_t lowest_used_interrupt;
            uint8_t highest_used_interrupt;
            uint8_t com_port_usage;
            uint8_t lpt_port_usage;
            uint8_t screen_usage;
 
            void SetDefaults();
 
            void ReadFile(const std::shared_ptr<Linker::Reader>& rd);
 
            void WriteFile(const std::shared_ptr<Linker::Writer>& wr) const;
 
            void Dump(Dumper::Dumper& dump, offset_t file_offset) const;
        };
 
        std::unique_ptr<PIF> pif;
 
        std::shared_ptr<Linker::Contents> image;
 
        magic_type GetSignature() const;
 
        void SetSignature(magic_type magic);
 
        void Clear() override;
 
        ~MZFormat()
        {
            Clear();
        }
 
        void SetFileSize(uint32_t size);
 
        uint32_t GetHeaderSize() const;
 
        uint32_t GetPifOffset() const;
 
        void ReadFile(const std::shared_ptr<Linker::Reader>& rd) override;
 
        using Linker::Format::WriteFile;
        offset_t WriteFile(const std::shared_ptr<Linker::Writer>& wr) const override;
 
        void Dump(Dumper::Dumper& dump) const override;
 
        void CalculateValues() override;
 
        /* * * Writer members * * */
 
        class MZOptionCollector : public Linker::OptionCollector
        {
        public:
            Linker::Option<std::optional<offset_t>> header_align{"header_align", "Aligns the end of the header to a specific boundary, must be power of 2"};
            Linker::Option<std::optional<offset_t>> file_align{"file_align", "Aligns the end of the file to a specific boundary, must be power of 2"};
            Linker::Option<offset_t> stack{"stack", "Specify the stack size"};
 
            MZOptionCollector()
            {
                InitializeFields(header_align, file_align, stack);
            }
        };
 
        enum memory_model_t
        {
            MODEL_DEFAULT,
            MODEL_TINY,
            MODEL_SMALL,
            MODEL_COMPACT,
            MODEL_LARGE,
        };
        memory_model_t memory_model = MODEL_DEFAULT;
 
        /* filled in automatically */
        uint16_t stack_size = 0;
 
        uint16_t extra_paras = 0;
 
        uint32_t zero_fill = 0;
 
        uint32_t option_header_align = 0x10;
 
        uint32_t option_file_align = 1;
 
        bool FormatSupportsSegmentation() const override;
 
        bool FormatIs16bit() const override;
 
        bool FormatIsProtectedMode() const override;
 
        unsigned FormatAdditionalSectionFlags(std::string section_name) const override;
 
        static std::vector<Linker::OptionDescription<void>> MemoryModelNames;
 
        std::vector<Linker::OptionDescription<void>> GetMemoryModelNames() override;
 
        void SetModel(std::string model) override;
 
        std::shared_ptr<Linker::OptionCollector> GetOptions() override;
 
        void SetOptions(std::map<std::string, std::string>& options) override;
 
        void OnNewSegment(std::shared_ptr<Linker::Segment> segment) override;
 
        void CreateDefaultSegments();
 
        std::unique_ptr<Script::List> GetScript(Linker::Module& module);
 
        void Link(Linker::Module& module);
 
#if 0
        void LinkLarge(Linker::Module& module);
#endif
 
        void ProcessModule(Linker::Module& module) override;
 
        uint32_t GetDataSize() const;
 
        void GenerateFile(std::string filename, Linker::Module& module) override;
 
        using Linker::OutputFormat::GetDefaultExtension;
        std::string GetDefaultExtension(Linker::Module& module, std::string filename) const override;
    };
 
    class MZSimpleStubWriter
    {
    public:
        std::string filename;
        bool valid = true;
        std::ifstream stream;
 
        MZSimpleStubWriter(std::string filename = "")
            : filename(filename)
        {
        }
 
        offset_t size = -1;
 
        bool OpenAndCheckValidFile();
 
        offset_t GetStubImageSize();
 
        void WriteStubImage(std::ostream& out);
 
        void WriteStubImage(const std::shared_ptr<Linker::Writer>& wr);
 
        ~MZSimpleStubWriter()
        {
            if(stream.is_open())
            {
                stream.close();
            }
        }
    };
 
    class MZStubWriter
    {
    public:
        std::string filename;
        bool valid = true;
        std::ifstream stream;
 
        MZStubWriter(std::string filename = "")
            : filename(filename)
        {
        }
 
        uint32_t original_file_size = -1;
        uint32_t stub_file_size = 0;
        uint16_t stub_reloc_count = 0;
        uint32_t original_header_size = 0;
        uint32_t stub_header_size = 0;
        uint16_t original_reloc_offset = 0;
        uint16_t stub_reloc_offset = 0;
 
        bool OpenAndCheckValidFile();
 
        offset_t GetStubImageSize();
 
        void WriteStubImage(std::ostream& out);
 
        void WriteStubImage(const std::shared_ptr<Linker::Writer>& wr);
 
        ~MZStubWriter()
        {
            if(stream.is_open())
            {
                stream.close();
            }
        }
    };
 
    template <typename T, size_t N, typename Predicate>
        offset_t FindActualSignature(const std::shared_ptr<Linker::Reader>& rd, std::array<T, N>& signature, Predicate predicate, bool search_win386_offset = false)
    {
        offset_t file_offset = rd->Tell();
        rd->ReadData(signature);
        if(file_offset == 0 && !predicate(signature))
        {
            // try to find real file offset
            offset_t file_end = rd->GetImageEnd();
 
            rd->Seek(2);
            uint32_t mz_image_end = rd->ReadUnsigned(2, ::LittleEndian);
            mz_image_end = (uint32_t(rd->ReadUnsigned(2, ::LittleEndian)) << 9) - (-mz_image_end & 0x1FF);
 
            uint32_t win386_header_offset = 0;
            uint32_t ne_header_offset = 0;
            if(file_end >= 0x40)
            {
                if(search_win386_offset)
                {
                    rd->Seek(0x38);
                    win386_header_offset = rd->ReadUnsigned(4, ::LittleEndian);
                }
 
                rd->Seek(0x3C);
                ne_header_offset = rd->ReadUnsigned(4, ::LittleEndian);
            }
 
            if(ne_header_offset != 0 && ne_header_offset + N < file_end)
            {
                rd->Seek(ne_header_offset);
                rd->ReadData(signature);
            }
 
            if(predicate(signature))
            {
                file_offset = rd->Tell() - N;
            }
            else
            {
                if(mz_image_end != 0 && mz_image_end != ne_header_offset && mz_image_end + N < file_end)
                {
                    rd->Seek(mz_image_end);
                    rd->ReadData(signature);
                }
 
                if(predicate(signature))
                {
                    file_offset = rd->Tell() - N;
                }
                else
                {
                    if(win386_header_offset != 0 && win386_header_offset != ne_header_offset && win386_header_offset != mz_image_end && win386_header_offset + N < file_end)
                    {
                        rd->Seek(win386_header_offset);
                        rd->ReadData(signature);
                    }
 
                    if(predicate(signature))
                    {
                        file_offset = rd->Tell() - N;
                    }
                    else
                    {
                        rd->Seek(N);
                    }
                }
            }
        }
 
        return file_offset;
    }
 
    template <typename T, size_t N>
        offset_t FindActualSignature(const std::shared_ptr<Linker::Reader>& rd, std::array<T, N>& signature, const char * expected_signature, bool search_win386_offset = false)
    {
        return FindActualSignature(
            rd,
            signature,
            [expected_signature](std::array<T, N>& signature) { return memcmp(signature.data(), expected_signature, N) == 0; },
            search_win386_offset);
    }
 
    template <typename T, size_t N>
        offset_t FindActualSignature(const std::shared_ptr<Linker::Reader>& rd, std::array<T, N>& signature, const char * expected_signature1, const char * expected_signature2, bool search_win386_offset = false)
    {
        return FindActualSignature(
            rd,
            signature,
            [expected_signature1, expected_signature2](std::array<T, N>& signature) { return memcmp(signature.data(), expected_signature1, N) == 0 || memcmp(signature.data(), expected_signature2, N) == 0; },
            search_win386_offset);
    }
 
    template <typename T, size_t N>
        offset_t FindActualSignature(const std::shared_ptr<Linker::Reader>& rd, std::array<T, N>& signature, const char * expected_signature1, const char * expected_signature2, const char * expected_signature3)
    {
        return FindActualSignature(
            rd,
            signature,
            [expected_signature1, expected_signature2, expected_signature3](std::array<T, N>& signature) { return memcmp(signature.data(), expected_signature1, N) == 0 || memcmp(signature.data(), expected_signature2, N) == 0 || memcmp(signature.data(), expected_signature3, N) == 0; },
            false);
    }
}
 
#endif /* MZEXE_H */