/* Copyright (c) Microsoft Corporation. All rights reserved. */
/*
 * MMLite loader for ARM, AIF image format.
 */

#define OPTIMIZE        /* enable Base.h optimizations, since we're the "base" */

#include <mmlite.h>
#include <assert.h>
#include <base/debugger.h>

#include <diagnostics.h>
#undef DBGLVL
#define DBGLVL 0

extern PIBASERTL pTheBaseRtl;

/*
 * This included file should provide the following functions
 * (or replacements thereof) and define the type of FILE_HANDLE.
 */
#include "loader_io.c"

extern FILE_HANDLE _Open( const _TCHAR * FileName, UINT Mode);
extern INT _Close( FILE_HANDLE File);
extern INT _Read( FILE_HANDLE File, BYTE *Buffer, UINT nBytes);
extern UINT _Seek( FILE_HANDLE File, UINT Position);
extern INT _Stat( FILE_HANDLE File);
/*extern INT _Write( FILE_HANDLE File, BYTE *Buffer, UINT nBytes); */
/*extern INT _Remove( TSTR FileName ); */


/* AIF header, 32 words total
 */
typedef struct _AIF_HEADER {
    UINT32  BlDecompress;
    UINT32  BlSelfReloc;
    UINT32  BlDebugInit;
    UINT32  BlEntryPoint;   /* if non exec, just EntryPoint */
    UINT32  ExitInstruction;
    UINT32  TextSize;
    UINT32  DataSize;
    UINT32  DebugSize;
    UINT32  BssSize;
    UINT32  DebugType;
    UINT32  ImageBase;
    UINT32  WorkSpace;
    UINT32  AddressingMode;
    UINT32  DataBase;
    UINT32  Reserved[2];
#if defined(MMLITE_VERSION)
#define VersionNumber Reserved[1]
#endif
    UINT32  DebugInitInstruction;
    UINT32  ZeroInitCode[15];
} AIF_HEADER;

/* Turn this on to save memory when loading DEBUG images
 * [This should be off in final product, images wont have debugs]
 */
#define SHRINK_DEBUGS 1

#if SHRINK_DEBUGS

/* Sanitize the relocation list of an image,
 * constraining entries within the text&data.
 */
static void FixRelocations( PBYTE RelocData, AIF_HEADER *pAif)
{
    UINT32 *pRelocList, *pNextReloc, Reloc;
    UINT Last;

    /* The reloc code has a fixed size of 46 words.
     * Sanity check against last instruction.
     */
    pRelocList = (UINT32 *)(RelocData + 45*4);
    if (*pRelocList++ != 0xeafffff8 /*b .-22"*/) {
        DBGME(3,printf("FixRelocations: bad reloc code %x\n",pRelocList[-1]));
        return;
    }
    
    /* Valid address range is [Base..Base+Text+Data]
     * Relocations are Base-relative addresses.
     */
    Last = pAif->TextSize + pAif->DataSize;

    /* Build new list by moving up invalid entries.
     * [We can only be shortening it]
     */
    pNextReloc = pRelocList;
    for (;;) {
        Reloc = *pNextReloc++;
        if (Reloc == 0xffffffff)
            break;

        /* Copy if valid entry
         */
        if (Reloc < Last)
            *pRelocList++ = Reloc;
    }
    
    DBGME(0,printf("FixRelocations: fixed %d relocs\n",
                 pRelocList - (UINT32*)(RelocData + 45*4)));

    /* Terminate new list
     */
    *pRelocList = 0xffffffff;
}
#endif /* SHRINK_DEBUGS */

/*
 * Compute the size of an in-memory AIF image.
 * [Used in case the simulator loaded the image, not us]
 */
SCODE SizeofArmImage( AIF_HEADER *pAif, INT *pSize )
{
    UINT MinSize, TrueSize;
    UINT RelocsSize, *pRelocs;
    
#if defined(MMLITE_VERSION)
    /* Is the version compatible with our build ?
     */
    if (pAif->VersionNumber > MMLITE_VERSION) {
        DBGME(3,printf("Error -- Incompatible version!\n"));
        return E_INCOMPATIBLE_VERSION;
    }
#endif

    /* Sanity checks.
     * NOTE: The ARM debugger loader does not load debug and reloc
     * info upon a LoadImage command, sometimes (when ?!?!?).
     * So these tests will have to suffice.
     */
    if ((pAif->ExitInstruction != 0xef000011) ||
        (pAif->AddressingMode != 0x20) ||
        (pAif->ZeroInitCode[0] != 0xe04ec00f))
        return E_INVALID_PARAMETER;

    /* At the very least, the image contains this much stuff
     */
    MinSize = pAif->TextSize + pAif->DataSize + pAif->DebugSize;
    TrueSize = MinSize + pAif->BssSize;

    /* Must look at the relocations. If there is no BSS we shouldnt
     * clobber them, even if yes, after the init this is free memory.
     * Use the "BL Relocs" instruction in the header for this.
     */
    RelocsSize = ((pAif->BlSelfReloc & 0x00ffffff) << 2) + 8 + 4;

    /* The reloc code has a fixed size of 46 words.
     * Sanity check against last instruction.
     */
    pRelocs  = (PUINT)((UINT)pAif + RelocsSize);
    pRelocs += 45;
    if (*pRelocs++ != 0xeafffff8 /*b .-22"*/) {
#if 0
        /* See comment above */
        return E_INVALID_PARAMETER;
#else
        *pSize = pAif->TextSize + pAif->DataSize + pAif->BssSize;
        return NO_ERROR;
#endif
    }

    /* Looks good. The relocs array is terminated by a -1
     */
    for (RelocsSize = 46*sizeof(UINT); *pRelocs != (UINT)-1; pRelocs++)
        RelocsSize += sizeof(UINT);

    /* Now see if the BSS is already bignuf.
     */
    RelocsSize += MinSize + 4;
    *pSize = (RelocsSize > TrueSize) ? RelocsSize : TrueSize;

    return NO_ERROR;
}

#if 1
/* this was a #if defined(ANGEL_SEMIHOST) but the code below does not
 * know how to handle partial reads, like via hostfsd (max 512 bytes).
 */

/*
 * Load an AIF formatted image in memory
 */
SCODE LdrLoadImage(
        PINAMESPACE pNameSpace,
        const _TCHAR * pImage,
        const _TCHAR * pArgs,
        UINT32 Flags,
        IModule **ppIModule)
{
    FILE_HANDLE f;
    INT ImageSize, MemorySize;
    INT RelocationSize;
    BYTE *pMem = NULL;
    PIHEAP pIHeap = CurrentHeap();
    MODULEENTRY EntryPoint;
    AIF_HEADER *pAif;
    PIMODULE pMod;
    PIPROCESS pPrc;
    UINT CntRead;
    SCODE sc = E_INVALID_PARAMETER; /* most likely */

    /* Check the file exists, see how big it is
     */
    f = _Open( pImage, 0x1 );
    if (f == INVALID_FILE_HANDLE) {
        DBGME(3,printf("File %s not found.\n",pImage));
        return E_FILE_NOT_FOUND;
    }

    ImageSize = _Stat( f );
    if (ImageSize < sizeof(AIF_HEADER)) {
        DBGME(3,printf("Image %s too small (%d).\n",pImage,ImageSize));
        goto Exit;
    }

    /* Read in the Header
     */
    pAif = (AIF_HEADER*) pIHeap->v->Alloc(pIHeap, 0, sizeof(AIF_HEADER), 0);
    if (pAif == NULL) {
        DBGME(3,printf("Outtamem %d.\n",sizeof(AIF_HEADER)));
        sc = E_NOT_ENOUGH_MEMORY;
        goto Exit;
    }

    CntRead = _Read( f, (PBYTE)pAif, sizeof(AIF_HEADER));
    if (sizeof(AIF_HEADER) != CntRead) {
        DBGME(3,printf("Short read %x %d\n",sc,CntRead));
        pIHeap->v->Free(pIHeap,0,pAif);
        goto Exit;
    }

#if defined(MMLITE_VERSION)
    /* Is the version compatible with our build ?
     */
    if (pAif->VersionNumber > MMLITE_VERSION) {
        DBGME(3,printf("Error -- Incompatible version!\n"));
        pIHeap->v->Free(pIHeap,0,pAif);
        sc = E_INCOMPATIBLE_VERSION;
        goto Exit;
    }
#endif

    _Seek(f, 0); /* as if nothing happened */

    /* Compute how much memory we need
     */
    MemorySize = pAif->TextSize + pAif->DataSize;
    RelocationSize = ImageSize - MemorySize - pAif->DebugSize;
    MemorySize += (RelocationSize > pAif->BssSize) ?
        RelocationSize : pAif->BssSize;

#if SHRINK_DEBUGS
    /* If we drop DEBUGs, drop BSS as well
     * Read relocs separately.
     */
    if (pAif->DebugSize)
        ImageSize = pAif->TextSize + pAif->DataSize;
#else
    /* Read all of it in, DEBUG included
     */
    if (ImageSize > MemorySize) 
        MemorySize = ImageSize;
#endif

    /* Done with header
     */
    pIHeap->v->Free(pIHeap,0,pAif);

    /* Allocate a descriptor for this image */

    /* Allocate enough memory for image
     */
    pMem = (PBYTE) pIHeap->v->Alloc(pIHeap, 0, MemorySize, 0);
    if (pMem == NULL) {
        DBGME(3,printf("Outtamem %d.\n",MemorySize));
        sc = E_NOT_ENOUGH_MEMORY;
        goto Exit;
    }

    DBGME(0,printf("mem=%x img=%x rel=%x\n",
                 MemorySize,ImageSize,RelocationSize));

    if (RelocationSize == 0) {
        DBGME(3,printf("Image has no relocations !?!?\n"));
        /* Continue anyways and hope for the best
         * BUGBUG Should check if linkbase is ok and reject ifnot
         */
    }

    /* This is the header, again
     */
    pAif = (AIF_HEADER *)pMem;

    /* Read Image in
     */
    CntRead = _Read( f, pMem, ImageSize);
    if (ImageSize != CntRead) {
        DBGME(3,printf("Short read %x %d\n",sc,CntRead));
        goto Exit;
    }

#if SHRINK_DEBUGS
    /* Drop DEBUG info, if any
     * If there are no symbols, the relocations have been
     * already read in above.  Else we gotta do it now.
     */
    if (pAif->DebugSize) {
        UINT RelocOffset;

        /* Distance from file header of reloc info:
         * Offset in Words, 2 words of PC pipeline, 1 word into header
         */
        RelocOffset = ((pAif->BlSelfReloc & 0x00ffffff) << 2) + 8 + 4;
        _Seek(f, RelocOffset);

        /* Where in memory it goes. We put it where the BSS *will be*,
         * knowing that the BSS is zeroed after relocation.
         */
        RelocOffset = pAif->TextSize + pAif->DataSize;
        CntRead = _Read(f, pMem + RelocOffset, RelocationSize);
        if (RelocationSize != CntRead) {
            DBGME(3,printf("Short reloc read %x %d\n",sc,CntRead));
            goto Exit;
        }

        /* One more twist.  The relocation list contains relocations
         * for the DEBUG info too (sic). Make sure the self-relocation
         * code will only patch within the image's boundaries.
         */
        FixRelocations(pMem + RelocOffset, pAif);

        /* Now patch branch inside header
         */
        pAif->BlSelfReloc = 0xeb000000 |  ((RelocOffset - 8 - 4) >> 2);
    }
#endif

    /* No need to zero BSS, program will do it itself
     */

    /* Done with file handle
     */
    _Close(f);
    f = INVALID_FILE_HANDLE;

    /* The first word of an AIF image is the entry point and is
     * a nop instruction.  Replace it with:
     *
     *   stmfd   r13!,{r0-r12,r14}
     *
     * to save all registers before the aif preamble code trashes
     * them.  The rtl main entry restores the registers once the
     * preamble is complete.
     */
#define CALL_FIXUP 0xE92D5FFF

    assert(0xe1a00000 == *(UINT32 *)pMem);
    *(volatile UINT32 *)pMem = CALL_FIXUP;

    EntryPoint = (MODULEENTRY) pMem;
    DBGME(3,printf("Image %hs is x%x bytes at x%x\n", pImage, MemorySize, pMem));

    FlushCache();

    /* Create a descriptor for this module we just loaded.
     */
    sc = ModuleCreate( pImage,
                       pArgs,
                       pMem,
                       MemorySize,
                       EntryPoint,
                       NULL,  /* Export Data  */
                       0,     /* Export Size  */
                       0, 0,
                       0, 0,
                       Flags,
                       &pMod);

    if (FAILED(sc)) {
        DBGME(3,printf("ModuleCreate failed %x\n",sc));
        goto Exit;
    }

    /* Hand caller one ref iff
     */
    if (ppIModule) {
        pMod->v->AddRef(pMod);
        *ppIModule = pMod;
    }

    if (Flags & LOAD_IMAGE_CALL_ENTRY_POINT) {
        if (Flags & LOAD_IMAGE_CREATE_THREAD) {

            /* Find an IProcess descriptor for the module
             */
            (void) pMod->v->QueryInterface(pMod, &IID_IProcess, (void**) &pPrc);

            /* Create a new thread to invoke image's entry point
             */
            sc = pPrc->v->CreateThread(pPrc, Flags,
                        (THREAD_FUNCTION) EntryPoint,
                        pTheBaseRtl,
                        0,/*default stacksize*/
                        NULL,
                        NULL);

            /* Release all our references (new thread holds one)
             */
            (void) pPrc->v->Release(pPrc);

        } else {
            /* Call the entry point directly
             */
            (EntryPoint)( pTheBaseRtl );
        }
    }
    pMod->v->Release(pMod);

    pMem = NULL;

 Exit:
    if (pMem)
        pIHeap->v->Free(pIHeap,0,pMem);
    if (f != INVALID_FILE_HANDLE)
      _Close(f);

    DBGME(0,printf("LdrLoadImage returns %x\n",sc));
    return sc;
}

#else

/*
 * Load an AIF formatted image in memory
 */
SCODE LdrLoadImage(
        PINAMESPACE pNameSpace,
        const _TCHAR * pImage,
        const _TCHAR * pArgs,
        UINT32 Flags,
        IModule **ppIModule)
{
    UINT ImageSize, MemorySize;
    UINT RelocationSize;
    UINT CntRead;
    UINT64 FileSize;
    BYTE *pMem = NULL;
    PIHEAP pIHeap = CurrentHeap();
    PIFILE pFile = NULL;
    MODULEENTRY EntryPoint;
    AIF_HEADER *pAif;
    PIUNKNOWN pUnk;
    PIMODULE pMod;
    PIPROCESS pPrc;
    SCODE sc;
    UINT64 zero = Int64Initializer(0,0);

    /* Check the file exists, see how big it is
     */
    if (pNameSpace == NULL) {
        DBGME(3,printf("Invalid NameSpace\n"));
        return E_INVALID_PARAMETER;
    }

    sc = pNameSpace->v->Bind(pNameSpace, pImage, NAME_SPACE_READ, &pUnk);
    if (FAILED(sc)) {
        DBGME(3,printf("File %s not found.\n",pImage));
        return E_FILE_NOT_FOUND;
    }

    sc = pUnk->v->QueryInterface(pUnk, &IID_IFile, (void**) &pFile);
    pUnk->v->Release(pUnk); /* Done with pUnk */
    if (FAILED(sc)) {
        DBGME(3,printf("%s not a file.\n",pImage));
        goto Error;
    }

    sc = pFile->v->GetSize(pFile, &FileSize);
    if (FAILED(sc)) {
        DBGME(3,printf("GetSize failed %x\n",sc));
        goto Error;
    }

    ImageSize = Int64ToInt32(FileSize);
    if (Int64ToInt32(Uint64RShift(FileSize, 32)) != 0 ||
        ImageSize < sizeof(AIF_HEADER)) {
        DBGME(3,printf("Image %s too small (%d).\n",pImage,ImageSize));
        sc = E_BAD_EXE_FORMAT;
        goto Error;
    }

    /* Read in the Header
     */
    pAif = (AIF_HEADER*) pIHeap->v->Alloc(pIHeap, 0, sizeof(AIF_HEADER), 0);
    if (pAif == NULL) {
        DBGME(3,printf("Outtamem %d.\n",sizeof(AIF_HEADER)));
        sc = E_NOT_ENOUGH_MEMORY;
        goto Exit;
    }

    sc = pFile->v->ReadAt(pFile, zero, (UINT8 *) pAif, sizeof(AIF_HEADER), &CntRead);
    if (FAILED(sc) || (CntRead < sizeof(AIF_HEADER))) {
        pIHeap->v->Free(pIHeap,0,pAif);
        DBGME(3,printf("Short read %x %d\n",sc,CntRead));
        sc = E_INVALID_PARAMETER;
        goto Error;
    }

#if defined(MMLITE_VERSION)
    /* Is the version compatible with our build ?
     */
    if (pAif->VersionNumber > MMLITE_VERSION) {
        DBGME(3,printf("Error -- Incompatible version!\n"));
        pIHeap->v->Free(pIHeap,0,pAif);
        sc = E_INCOMPATIBLE_VERSION;
        goto Error;
    }
#endif

    /* Compute how much memory we need
     */
    MemorySize = pAif->TextSize + pAif->DataSize;
    RelocationSize = ImageSize - MemorySize - pAif->DebugSize;
    MemorySize += (RelocationSize > pAif->BssSize) ?
        RelocationSize : pAif->BssSize;

#if SHRINK_DEBUGS
    /* If we drop DEBUGs, drop BSS as well
     * Read relocs separately.
     */
    if (pAif->DebugSize)
        ImageSize = pAif->TextSize + pAif->DataSize;
#else
    /* Read all of it in, DEBUG included
     */
    if (ImageSize > MemorySize) 
        MemorySize = ImageSize;
#endif

    /* Done with header
     */
    pIHeap->v->Free(pIHeap,0,pAif);

    /* Allocate a descriptor for this image */

    /* Allocate enough memory for image
     */
    pMem = (PBYTE) pIHeap->v->Alloc(pIHeap, 0, MemorySize, 0);
    if (pMem == NULL) {
        DBGME(3,printf("Outtamem %d.\n",MemorySize));
        sc = E_NOT_ENOUGH_MEMORY;
        goto Exit;
    }

    DBGME(0,printf("mem=%x img=%x rel=%x\n",
                 MemorySize,ImageSize,RelocationSize));

    if (RelocationSize == 0) {
        DBGME(3,printf("Image has no relocations !?!?\n"));
        /* Continue anyways and hope for the best
         * BUGBUG Should check if linkbase is ok and reject ifnot
         */
    }

    /* This is the header, again
     */
    pAif = (AIF_HEADER *)pMem;

    /* Read Image in
     */
    sc = pFile->v->ReadAt(pFile, zero, pMem, ImageSize, &CntRead);
    if (FAILED(sc) || (CntRead < ImageSize)) {
        DBGME(3,printf("Short read %x %d\n",sc,CntRead));
        sc = E_INVALID_PARAMETER;
        goto Exit;
    }
#if SHRINK_DEBUGS
    /* Drop DEBUG info, if any
     * If there are no symbols, the relocations have been
     * already read in above.  Else we gotta do it now.
     */
    if (pAif->DebugSize) {
        UINT64 FileOffset;
        UINT RelocOffset;

        /* Distance from file header of reloc info:
         * Offset in Words, 2 words of PC pipeline, 1 word into header
         */
        Int64FromHighAndLow(FileOffset, 0, ((pAif->BlSelfReloc & 0x00ffffff) << 2) + 8 + 4);

        /* Where in memory it goes. We put it where the BSS *will be*,
         * knowing that the BSS is zeroed after relocation.
         */
        RelocOffset = pAif->TextSize + pAif->DataSize;
        CntRead = 0;
        sc = pFile->v->ReadAt(pFile, FileOffset, pMem + RelocOffset, RelocationSize, &CntRead);
        if (CntRead < RelocationSize)
            sc = E_INVALID_PARAMETER;
        if (FAILED(sc)) {
            DBGME(3,printf("Short reloc read %x %d\n",sc,CntRead));
            goto Exit;
        }

        /* One more twist.  The relocation list contains relocations
         * for the DEBUG info too (sic). Make sure the self-relocation
         * code will only patch within the image's boundaries.
         */
        FixRelocations(pMem + RelocOffset, pAif);

        /* Now patch branch inside header
         */
        pAif->BlSelfReloc = 0xeb000000 |  ((RelocOffset - 8 - 4) >> 2);
    }
#endif

    /* No need to zero BSS, program will do it itself
     */

    /* Done with file handle
     */
    pFile->v->Release(pFile);
    pFile = NULL;

    /* The first word of an AIF image is the entry point and is
     * a nop instruction.  Replace it with:
     *
     *   stmfd   r13!,{r0-r12,r14}
     *
     * to save all registers before the aif preamble code trashes
     * them.  The rtl main entry restores the registers once the
     * preamble is complete.
     */
#define CALL_FIXUP 0xE92D5FFF

    assert(0xe1a00000 == *(UINT32 *)pMem);
    *(volatile UINT32 *)pMem = CALL_FIXUP;

    EntryPoint = (MODULEENTRY) pMem;
    DBGME(3,printf("Image %hs is x%x bytes at x%x\n", pImage, MemorySize, pMem));

    FlushCache();

    /* Create a descriptor for this module we just loaded.
     */
    sc = ModuleCreate( pImage,
                       pArgs,
                       pMem,
                       MemorySize,
                       EntryPoint,
                       NULL,    /* ExportData */
                       0,       /* ExportSize  */
                       NULL,    /* ImportedModules */
                       0,       /* nImportedModules */
                       0,       /* Metadata */
                       0,       /* com+ entry point token */
                       Flags,
                       &pMod);

    if (FAILED(sc)) {
        DBGME(3,printf("ModuleCreate failed %x\n",sc));
        goto Exit;
    }

    /* Hand caller one ref iff
     */
    if (ppIModule) {
        pMod->v->AddRef(pMod);
        *ppIModule = pMod;
    }

    if (Flags & LOAD_IMAGE_CALL_ENTRY_POINT) {
        if (Flags & LOAD_IMAGE_CREATE_THREAD) {
            /* Find an IProcess descriptor for the module
             */
            (void) pMod->v->QueryInterface(pMod, &IID_IProcess, (void**) &pPrc);

            /* Create a new thread to invoke image's entry point
             */
            sc = pPrc->v->CreateThread(pPrc, Flags,
                        (THREAD_FUNCTION) EntryPoint,
                        pTheBaseRtl,
                        0,/*default stacksize*/
                        NULL,
                        NULL);

            /* Release all our references (new thread holds one)
             */
            (void) pPrc->v->Release(pPrc);

        } else {
            /* Call the entry point directly
             */
            (EntryPoint)( pTheBaseRtl );
        }
    }
    pMod->v->Release(pMod);

    pMem = NULL;

 Error:
 Exit:
    if (pMem)
        pIHeap->v->Free(pIHeap,0,pMem);
    if (pFile)
        pFile->v->Release(pFile);

    DBGME(0,printf("LdrLoadImage returns %x\n",sc));
    return sc;
}

#endif