EncoreBootImage.h

/*
 * File:    EncoreBootImage.h
 *
 * Copyright (c) Freescale Semiconductor, Inc. All rights reserved.
 * See included license file for license details.
 */
#if !defined(_EncoreBootImage_h_)
#define _EncoreBootImage_h_

#include <list>
#include <vector>
#include <string>
#include <iostream>
#include <fstream>
#include <string.h>
#include "BootImage.h"
#include "rijndael.h"
#include "smart_ptr.h"
#include "AESKey.h"
#include "StExecutableImage.h"

namespace elftosb
{
typedef uint8_t cipher_block_t[16];

typedef uint8_t sha1_digest_t[20];

typedef uint32_t section_id_t;

inline size_t sizeOfCipherBlocks(unsigned count)
{
    return sizeof(cipher_block_t) * count;
}

inline size_t numberOfCipherBlocks(size_t s)
{
    return (s + sizeof(cipher_block_t) - 1) / sizeof(cipher_block_t);
}

inline size_t sizeInCipherBlocks(size_t s)
{
    return (unsigned)sizeOfCipherBlocks(numberOfCipherBlocks(s));
}

inline size_t sizeOfPaddingForCipherBlocks(size_t s)
{
    return sizeInCipherBlocks(s) - s;
}

class EncoreBootImage : public BootImage
{
public:
    enum
    {
        ROM_DISPLAY_PROGRESS = (1 << 0), 
        ROM_VERBOSE_PROGRESS = (1 << 1)  
    };

    enum
    {
        ROM_IMAGE_HEADER_SIGNATURE =
            'STMP', 
        ROM_IMAGE_HEADER_SIGNATURE2 =
            'sgtl',                       
        ROM_BOOT_IMAGE_MAJOR_VERSION = 1, 
        ROM_BOOT_IMAGE_MINOR_VERSION = 2  
    };

    enum
    {
        BOOT_IMAGE_MINIMUM_SECTION_ALIGNMENT = sizeof(cipher_block_t)
    };

// All of these structures are packed to byte alignment in order to match
// the structure on disk.
#pragma pack(1)

    struct boot_image_header_t
    {
        union
        {
            sha1_digest_t m_digest; 
            struct
            {
                cipher_block_t m_iv; 
                uint8_t m_extra[4];  
            };
        };
        uint8_t m_signature[4];       
        uint8_t m_majorVersion;       
        uint8_t m_minorVersion;       
        uint16_t m_flags;             
        uint32_t m_imageBlocks;       
        uint32_t m_firstBootTagBlock; 
        section_id_t m_firstBootableSectionID; 
        uint16_t m_keyCount;                   
        uint16_t m_keyDictionaryBlock;         
        uint16_t m_headerBlocks;               
        uint16_t m_sectionCount;               
        uint16_t m_sectionHeaderSize;          
        uint8_t m_padding0[2];                 
        uint8_t m_signature2[4];      
        uint64_t m_timestamp;         
        version_t m_productVersion;   
        version_t m_componentVersion; 
        uint16_t m_driveTag;          
        uint8_t m_padding1[6];        
    };

    struct dek_dictionary_entry_t
    {
        cipher_block_t m_mac; 
        aes128_key_t m_dek;   
    };

    struct dek_dictionary_t
    {
        dek_dictionary_entry_t m_entries[1];
    };

    enum
    {
        ROM_SECTION_BOOTABLE = (1 << 0), 
        ROM_SECTION_CLEARTEXT =
            (1 << 1) 
    };

    struct section_header_t
    {
        uint32_t m_tag;    
        uint32_t m_offset; 
        uint32_t m_length; 
        uint32_t m_flags;  
    };

    struct section_table_t
    {
        section_header_t m_sections[1]; 
    };

    struct boot_command_t
    {
        uint8_t m_checksum; 
        uint8_t m_tag;      
        uint16_t m_flags;   
        uint32_t m_address; 
        uint32_t m_count;   
        uint32_t m_data;    
    };

#pragma pack()

    enum
    {
        ROM_NOP_CMD = 0x00,        
        ROM_TAG_CMD = 0x01,        
        ROM_LOAD_CMD = 0x02,       
        ROM_FILL_CMD = 0x03,       
        ROM_JUMP_CMD = 0x04,       
        ROM_CALL_CMD = 0x05,       
        ROM_MODE_CMD = 0x06,       
        ROM_ERASE_CMD = 0x07,      
        ROM_RESET_CMD = 0x08,      
        ROM_MEM_ENABLE_CMD = 0x09, 
        ROM_PROG_CMD = 0x0a        
    };

    enum
    {
        ROM_LAST_TAG = (1 << 0) 
    };

    enum
    {
        ROM_LOAD_DCD = (1 << 0) 
    };

    enum
    {
        ROM_FILL_BYTE = 0,      
        ROM_FILL_HALF_WORD = 1, 
        ROM_FILL_WORD = 2       
    };

    enum
    {
        ROM_HAB_EXEC = (1 << 0),    
        ROM_JUMP_SP_MASK = (1 << 1) 
    };

    enum
    {
        ROM_MEM_CTRL_MASK = 0x0f, 
        ROM_MEM_CTRL_SHIFT = 8    
    };

    enum
    {
        ROM_MEM_SPACE_MASK = 0x0f, 
        ROM_MEM_SPACE_SHIFT = 8    
    };

    enum
    {
        ROM_ERASE_ALL_MASK = (1 << 0),          
        ROM_ERASE_ALL_UNSECURE_MASK = (1 << 1), 
    };

    enum
    {
        ROM_PROG_8BYTE_MASK = (1 << 0) 
    };

    enum
    {
        ROM_MEM_ID_QSPI = 1, 
    };

    enum
    {
        ROM_MEM_SPACE_IFR = 4, 
    };

public:
    // Forward declaration.
    class Section;

    class CipherBlockGenerator
    {
    public:

        virtual unsigned getBlockCount() const { return 0; }
        virtual unsigned getBlocks(unsigned offset, unsigned maxCount, cipher_block_t *data) { return 0; }

        virtual void debugPrint() const {}
    };

    class BootCommand : public CipherBlockGenerator
    {
    public:
        static BootCommand *createFromData(const cipher_block_t *blocks, unsigned count, unsigned *consumed);

    public:
        BootCommand()
            : CipherBlockGenerator()
        {
        }

        virtual ~BootCommand() {}
        virtual void initFromData(const cipher_block_t *blocks, unsigned count, unsigned *consumed) = 0;


        virtual uint8_t getTag() const = 0;

        virtual void fillCommandHeader(boot_command_t &header) = 0;

        virtual uint8_t calculateChecksum(const boot_command_t &header);


        virtual unsigned getBlockCount() const;

        virtual unsigned getBlocks(unsigned offset, unsigned maxCount, cipher_block_t *data);


        virtual unsigned getDataBlockCount() const { return 0; }
        virtual unsigned getDataBlocks(unsigned offset, unsigned maxCount, cipher_block_t *data) { return 0; }

    protected:
        enum
        {
            CMD_TAG_FIELD = 1,
            CMD_FLAGS_FIELD = 2,
            CMD_ADDRESS_FIELD = 4,
            CMD_COUNT_FIELD = 8,
            CMD_DATA_FIELD = 16
        };

        void validateHeader(const boot_command_t *modelHeader, const boot_command_t *testHeader, unsigned whichFields);
    };

    class NopCommand : public BootCommand
    {
    public:
        NopCommand()
            : BootCommand()
        {
        }

        virtual void initFromData(const cipher_block_t *blocks, unsigned count, unsigned *consumed);


        virtual uint8_t getTag() const { return ROM_NOP_CMD; }
        virtual void fillCommandHeader(boot_command_t &header);

        virtual void debugPrint() const;
    };

    class TagCommand : public BootCommand
    {
    public:
        TagCommand()
            : BootCommand()
        {
        }

        TagCommand(const Section &section);

        virtual void initFromData(const cipher_block_t *blocks, unsigned count, unsigned *consumed);


        virtual uint8_t getTag() const { return ROM_TAG_CMD; }
        virtual void fillCommandHeader(boot_command_t &header);


        inline void setSectionIdentifier(uint32_t identifier) { m_sectionIdentifier = identifier; }
        inline void setSectionLength(uint32_t length) { m_sectionLength = length; }
        inline void setSectionFlags(uint32_t flags) { m_sectionFlags = flags; }
        inline void setLast(bool isLast) { m_isLast = isLast; }

        virtual void debugPrint() const;

    protected:
        uint32_t m_sectionIdentifier; 
        uint32_t m_sectionLength;     
        uint32_t m_sectionFlags;      
        bool m_isLast;                
    };

    class LoadCommand : public BootCommand
    {
    public:
        LoadCommand();

        LoadCommand(uint32_t address, const uint8_t *data, uint32_t length);

        virtual void initFromData(const cipher_block_t *blocks, unsigned count, unsigned *consumed);


        virtual uint8_t getTag() const { return ROM_LOAD_CMD; }
        virtual void fillCommandHeader(boot_command_t &header);

        inline void setDCD(bool isDCD) { m_loadDCD = isDCD; }


        inline void setLoadAddress(uint32_t address) { m_address = address; }
        inline uint32_t getLoadAddress() const { return m_address; }


        void setData(const uint8_t *data, uint32_t length);

        inline uint8_t *getData() { return m_data; }
        inline const uint8_t *getData() const { return m_data; }
        inline uint32_t getLength() const { return m_length; }


        virtual unsigned getDataBlockCount() const;

        virtual unsigned getDataBlocks(unsigned offset, unsigned maxCount, cipher_block_t *data);

        virtual void debugPrint() const;

    protected:
        smart_array_ptr<uint8_t> m_data; 
        uint8_t m_padding[15];           
        unsigned m_padCount;             
        uint32_t m_length;               
        uint32_t m_address;              
        bool m_loadDCD;                  

        void fillPadding();
        uint32_t calculateCRC() const;
    };

    class FillCommand : public BootCommand
    {
    public:
        FillCommand();

        virtual void initFromData(const cipher_block_t *blocks, unsigned count, unsigned *consumed);


        virtual uint8_t getTag() const { return ROM_FILL_CMD; }
        virtual void fillCommandHeader(boot_command_t &header);


        inline void setAddress(uint32_t address) { m_address = address; };
        inline uint32_t getAddress() const { return m_address; }
        inline void setFillCount(uint32_t count) { m_count = count; }
        inline uint32_t getFillCount() const { return m_count; }


        void setPattern(uint8_t pattern);
        void setPattern(uint16_t pattern);
        void setPattern(uint32_t pattern);

        inline uint32_t getPattern() const { return m_pattern; }

        virtual void debugPrint() const;

    protected:
        uint32_t m_address; 
        uint32_t m_count;   
        uint32_t m_pattern; 
    };

    class ModeCommand : public BootCommand
    {
    public:
        ModeCommand()
            : BootCommand()
            , m_mode(0)
        {
        }

        virtual void initFromData(const cipher_block_t *blocks, unsigned count, unsigned *consumed);


        virtual uint8_t getTag() const { return ROM_MODE_CMD; }
        virtual void fillCommandHeader(boot_command_t &header);


        inline void setBootMode(uint32_t mode) { m_mode = mode; }
        inline uint32_t getBootMode() const { return m_mode; }

        virtual void debugPrint() const;

    protected:
        uint32_t m_mode; 
    };

    class EraseCommand : public BootCommand
    {
    public:
        EraseCommand()
            : BootCommand()
            , m_doEraseAll(false)
            , m_doEraseAllUnsecure(false)
            , m_startAddress(0)
            , m_byteCount(0)
            , m_memoryId(0)
        {
        }

        virtual void initFromData(const cipher_block_t *blocks, unsigned count, unsigned *consumed);


        virtual uint8_t getTag() const { return ROM_ERASE_CMD; }
        virtual void fillCommandHeader(boot_command_t &header);


        void setEraseAll(bool doEraseAll) { m_doEraseAll = doEraseAll; }
        bool getEraseAll() const { return m_doEraseAll; }


        void setEraseAllUnsecure(bool doEraseAllUnsecure) { m_doEraseAllUnsecure = doEraseAllUnsecure; }
        bool getEraseAllUnsecure() const { return m_doEraseAllUnsecure; }


        void setAddressRange(uint32_t startAddress, uint32_t count);
        void getAddressRange(uint32_t *startAddress, uint32_t *count) const;


        void setMemoryId(uint32_t memId) { m_memoryId = memId; }
        uint32_t getMemoryId() const { return m_memoryId; }

        virtual void debugPrint() const;

    protected:
        bool m_doEraseAll;
        bool m_doEraseAllUnsecure;
        uint32_t m_startAddress;
        uint32_t m_byteCount;
        uint32_t m_memoryId;
    };

    class ProgramCommand : public BootCommand
    {
    public:
        ProgramCommand()
            : BootCommand()
            , m_isEightByte(false)
            , m_index(0)
            , m_dataWord1(0)
            , m_dataWord2(0)
            , m_memSpace(0)
        {
        }

        virtual void initFromData(const cipher_block_t *blocks, unsigned count, unsigned *consumed);


        virtual uint8_t getTag() const { return ROM_PROG_CMD; }
        virtual void fillCommandHeader(boot_command_t &header);


        void setIsEightByte(bool doIt) { m_isEightByte = doIt; }
        bool isEightByte() const { return m_isEightByte; }


        void setMemSpace(uint32_t id) { m_memSpace = id; }
        uint32_t getMemSpace() const { return m_memSpace; }


        void setIndex(uint32_t index) { m_index = index; }
        uint32_t getIndex() const { return m_index; }


        void setDataWord1(uint32_t dataWord) { m_dataWord1 = dataWord; }
        uint32_t getDataWord1() const { return m_dataWord1; }


        void setDataWord2(uint32_t dataWord) { m_dataWord2 = dataWord; }
        uint32_t getDataWord2() const { return m_dataWord2; }

        virtual void debugPrint() const;

    protected:
        bool m_isEightByte;   
        uint32_t m_index;     
        uint32_t m_dataWord1; 
        uint32_t m_dataWord2; 
        uint32_t m_memSpace;  
    };

    class JumpCommand : public BootCommand
    {
    public:
        JumpCommand()
            : BootCommand()
            , m_address(0)
            , m_argument(0)
            , m_isHAB(false)
            , m_ivtSize(0)
            , m_stackPointer(0)
            , m_isStackPointerSet(false)
        {
        }

        virtual void initFromData(const cipher_block_t *blocks, unsigned count, unsigned *consumed);


        virtual uint8_t getTag() const { return ROM_JUMP_CMD; }
        virtual void fillCommandHeader(boot_command_t &header);


        inline void setAddress(uint32_t address) { m_address = address; }
        inline uint32_t getAddress() const { return m_address; }
        inline void setArgument(uint32_t argument) { m_argument = argument; }
        inline uint32_t getArgument() const { return m_argument; }
        inline void setIsHAB(bool isHAB) { m_isHAB = isHAB; }
        inline bool isHAB() const { return m_isHAB; }
        inline void setIVTSize(uint32_t ivtSize) { m_ivtSize = ivtSize; }
        inline uint32_t getIVTSize() const { return m_ivtSize; }
        inline void setStackPointer(uint32_t sp) { m_stackPointer = sp; }
        inline uint32_t getStackPointer() { return m_stackPointer; }
        inline void setIsStackPointerSet(bool isSP) { m_isStackPointerSet = isSP; }
        inline bool IsStackPointerSet() const { return m_isStackPointerSet; }

        virtual void debugPrint() const;

    protected:
        uint32_t m_address;      
        uint32_t m_argument;     
        uint32_t m_stackPointer; 
        bool m_isHAB; 
        bool m_isStackPointerSet; 
        uint32_t m_ivtSize;       
    };

    class CallCommand : public JumpCommand
    {
    public:
        CallCommand()
            : JumpCommand()
        {
        }

        virtual uint8_t getTag() const { return ROM_CALL_CMD; }
        virtual void debugPrint() const;
    };

    class ResetCommand : public BootCommand
    {
    public:
        ResetCommand()
            : BootCommand()
        {
        }

        virtual void initFromData(const cipher_block_t *blocks, unsigned count, unsigned *consumed);


        virtual uint8_t getTag() const { return ROM_RESET_CMD; }
        virtual void fillCommandHeader(boot_command_t &header);

        virtual void debugPrint() const;

    protected:
    };

    class MemEnableCommand : public BootCommand
    {
    public:
        MemEnableCommand()
            : BootCommand()
            , m_startAddress(0)
            , m_byteCount(0)
            , m_memControllerId(0)
        {
        }

        virtual void initFromData(const cipher_block_t *blocks, unsigned count, unsigned *consumed);


        virtual uint8_t getTag() const { return ROM_MEM_ENABLE_CMD; }
        virtual void fillCommandHeader(boot_command_t &header);


        void setAddressRange(uint32_t startAddress, uint32_t count);
        void getAddressRange(uint32_t *startAddress, uint32_t *count) const;


        void setMemControllerId(uint32_t memControllerId) { m_memControllerId = memControllerId; }
        uint32_t getMemControllerId() const { return m_memControllerId; }

        virtual void debugPrint() const;

    protected:
        uint32_t m_startAddress;
        uint32_t m_byteCount;
        uint32_t m_memControllerId;
    };

    class Section : public CipherBlockGenerator
    {
    public:
        Section()
            : CipherBlockGenerator()
            , m_identifier(0)
            , m_image(0)
            , m_alignment(BOOT_IMAGE_MINIMUM_SECTION_ALIGNMENT)
            , m_flags(0)
            , m_leaveUnencrypted(false)
        {
        }

        Section(uint32_t identifier)
            : CipherBlockGenerator()
            , m_identifier(identifier)
            , m_image(0)
            , m_alignment(BOOT_IMAGE_MINIMUM_SECTION_ALIGNMENT)
            , m_flags(0)
            , m_leaveUnencrypted(false)
        {
        }


        inline void setIdentifier(uint32_t identifier) { m_identifier = identifier; }
        inline uint32_t getIdentifier() const { return m_identifier; }


        virtual void setFlags(uint32_t flags) { m_flags = flags; }
        virtual uint32_t getFlags() const { return m_flags | (m_leaveUnencrypted ? ROM_SECTION_CLEARTEXT : 0); }
        virtual void fillSectionHeader(section_header_t &header);


        void setImage(EncoreBootImage *image) { m_image = image; }
        EncoreBootImage *getImage() const { return m_image; }


        void setAlignment(unsigned alignment);

        unsigned getAlignment() const { return m_alignment; }
        unsigned getPadBlockCountForOffset(unsigned offset);


        void setLeaveUnencrypted(unsigned flag) { m_leaveUnencrypted = flag; }
        bool getLeaveUnencrypted() const { return m_leaveUnencrypted; }

    protected:
        uint32_t m_identifier;    
        EncoreBootImage *m_image; 
        unsigned m_alignment;     
        uint32_t m_flags;         
        bool m_leaveUnencrypted;  
    };

    class BootSection : public Section
    {
    public:
        typedef std::list<BootCommand *> command_list_t;
        typedef command_list_t::iterator iterator_t;
        typedef command_list_t::const_iterator const_iterator_t;

    public:
        BootSection()
            : Section()
        {
        }

        BootSection(uint32_t identifier)
            : Section(identifier)
        {
        }

        virtual ~BootSection();

        virtual void fillFromData(const cipher_block_t *blocks, unsigned count);


        virtual uint32_t getFlags() const { return Section::getFlags() | ROM_SECTION_BOOTABLE; }


        void addCommand(BootCommand *command) { m_commands.push_back(command); }
        unsigned getCommandCount() const { return (unsigned)m_commands.size(); }
        iterator_t begin() { return m_commands.begin(); }
        iterator_t end() { return m_commands.end(); }
        const_iterator_t begin() const { return m_commands.begin(); }
        const_iterator_t end() const { return m_commands.end(); }


        virtual unsigned getBlockCount() const;

        virtual unsigned getBlocks(unsigned offset, unsigned maxCount, cipher_block_t *data);

        virtual void debugPrint() const;

    protected:
        command_list_t m_commands; 

    protected:
        void deleteCommands();
    };

    class DataSection : public Section
    {
    public:
        DataSection()
            : Section()
            , m_data()
            , m_length(0)
        {
        }

        DataSection(uint32_t identifier)
            : Section(identifier)
            , m_data()
            , m_length(0)
        {
        }

        void setData(const uint8_t *data, unsigned length);

        void setDataNoCopy(const uint8_t *data, unsigned length);


        virtual unsigned getBlockCount() const;

        virtual unsigned getBlocks(unsigned offset, unsigned maxCount, cipher_block_t *data);

        virtual void debugPrint() const;

    protected:
        smart_array_ptr<uint8_t> m_data; 
        unsigned m_length;               
    };

public:
    typedef std::list<Section *> section_list_t;                     
    typedef section_list_t::iterator section_iterator_t;             
    typedef section_list_t::const_iterator const_section_iterator_t; 

    typedef std::vector<AES128Key> key_list_t;               
    typedef key_list_t::iterator key_iterator_t;             
    typedef key_list_t::const_iterator const_key_iterator_t; 

public:
    EncoreBootImage();

    virtual ~EncoreBootImage();


    void addSection(Section *newSection);
    inline unsigned sectionCount() const { return (unsigned)m_sections.size(); }
    inline section_iterator_t beginSection() { return m_sections.begin(); }
    inline section_iterator_t endSection() { return m_sections.end(); }
    inline const_section_iterator_t beginSection() const { return m_sections.begin(); }
    inline const_section_iterator_t endSection() const { return m_sections.end(); }
    section_iterator_t findSection(Section *section);

    uint32_t getSectionOffset(Section *section);


    inline void addKey(const AES128Key &newKey) { m_keys.push_back(newKey); }
    inline unsigned keyCount() const { return (unsigned)m_keys.size(); }
    inline key_iterator_t beginKeys() { return m_keys.begin(); }
    inline key_iterator_t endKeys() { return m_keys.end(); }
    inline const_key_iterator_t beginKeys() const { return m_keys.begin(); }
    inline const_key_iterator_t endKeys() const { return m_keys.end(); }
    inline bool isEncrypted() const { return m_keys.size() != 0; }


    virtual void setProductVersion(const version_t &version);
    virtual void setComponentVersion(const version_t &version);


    inline void setFlags(uint16_t flags) { m_headerFlags = flags; }
    inline uint32_t getFlags() const { return m_headerFlags; }

    virtual void setDriveTag(uint16_t tag) { m_driveTag = tag; }
    uint32_t getImageSize();

    virtual std::string getFileExtension() const { return ".sb"; }

    virtual void writeToStream(std::ostream &stream);

    virtual void debugPrint() const;

protected:
    uint16_t m_headerFlags;       
    version_t m_productVersion;   
    version_t m_componentVersion; 
    uint16_t m_driveTag;          
    section_list_t m_sections;    
    key_list_t m_keys;            
    AES128Key m_sessionKey;       

    void prepareImageHeader(boot_image_header_t &header);
    uint64_t getTimestamp();
    Section *findFirstBootableSection();
    unsigned getPadBlockCountForSection(Section *section, unsigned offset);
};

}; // namespace elftosb

#endif // _EncoreBootImage_h_