mzexe.h

#ifndef MZEXE_H
#define MZEXE_H
 
#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;
 
        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(Linker::Reader& rd);
 
            void WriteFile(Linker::Writer& wr) const;
 
            void Dump(Dumper::Dumper& dump, offset_t file_offset) const;
        };
 
        std::unique_ptr<PIF> pif;
 
        std::shared_ptr<Linker::Image> 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(Linker::Reader& rd) override;
 
        using Linker::Format::WriteFile;
        offset_t WriteFile(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"};
 
            MZOptionCollector()
            {
                InitializeFields(header_align, file_align);
            }
        };
 
        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 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(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(Linker::Writer& wr);
 
        ~MZStubWriter()
        {
            if(stream.is_open())
            {
                stream.close();
            }
        }
    };
 
}
 
#endif /* MZEXE_H */