//
// Copyright (c) Microsoft Corporation.  All rights reserved.
//
//
// Use of this source code is subject to the terms of the Microsoft shared
// source or premium shared source license agreement under which you licensed
// this source code. If you did not accept the terms of the license agreement,
// you are not authorized to use this source code. For the terms of the license,
// please see the license agreement between you and Microsoft or, if applicable,
// see the SOURCE.RTF on your install media or the root of your tools installation.
// THE SOURCE CODE IS PROVIDED "AS IS", WITH NO WARRANTIES.
//
// reloc.cpp

#include "diffbin.h"

BOOL
HeaderMagicOk(const e32_exe& e32)
/*---------------------------------------------------------------------------*\
 * REVIEW: should this check for the two zero bytes also?
\*---------------------------------------------------------------------------*/
{
	return strcmp(reinterpret_cast<const char*>(&e32.e32_magic[0]),"PE") == 0;

} /* HeaderMagicOk()
   */

HRESULT
DoFixups(CModuleData *pOldMod, CModuleData *pNewMod, LPBYTE pbFixupBuffer, 
         DWORD dwLen)
/*---------------------------------------------------------------------------*\
 *
\*---------------------------------------------------------------------------*/
{
    HRESULT             hr              = NOERROR;
    o32_rom *           dataptr;
	struct info *       blockptr;
    struct info *       blockstart;
	ULONG               blocksize       = dwLen;
	LPDWORD             FixupAddress;
    LPWORD              currptr;
	WORD                curroff;
    DWORD               offset;
	BOOL                MatchedReflo    =FALSE;
	int                 iLoop;
    LPBYTE              pbO32Addr;
    CSectionData *      pSection        = NULL;
    
	RETAILMSG(ZONE_FIXUP_VERBOSE, (TEXT("DoFixups: Old e32_vbase=0x%08lx, New e32_vbase=0x%08lx\n"),
		pOldMod->GetE32()->e32_vbase, pNewMod->GetE32()->e32_vbase));

	blockstart = blockptr = (struct info *)pbFixupBuffer;
	while (((ULONG)blockptr < (ULONG)blockstart + blocksize) && blockptr->size) {
		currptr = (LPWORD)(((ULONG)blockptr)+sizeof(struct info));
		if ((DWORD)currptr >= ((DWORD)blockptr+blockptr->size)) {
			blockptr = (struct info *)(((DWORD)blockptr)+blockptr->size);
			continue;
		}
		dataptr = 0;
		while ((ULONG)currptr < ((ULONG)blockptr+blockptr->size)) {
			curroff = *currptr&0xfff;
			if (!curroff && !blockptr->rva) {
				currptr++;
				continue;
			}
			if (!dataptr || (dataptr->o32_rva > blockptr->rva + curroff) ||
				(dataptr->o32_rva + dataptr->o32_vsize <= blockptr->rva + curroff)) 
        {
				for (iLoop = 0; iLoop < pNewMod->GetE32()->e32_objcnt; iLoop++) {
					dataptr = pNewMod->GetO32(iLoop);
					if ((dataptr->o32_rva <= blockptr->rva + curroff) &&
              (dataptr->o32_rva+dataptr->o32_vsize > blockptr->rva + curroff)) 
            {
                // Got new O32 section... Got to go and find it.
                offset = pOldMod->GetO32(iLoop)->o32_realaddr - pNewMod->GetO32(iLoop)->o32_realaddr;

                RETAILMSG(ZONE_FIXUP_VERBOSE, (TEXT("DoFixups: Offset is %8.8lx\n"), offset));

                pSection = NULL;

                CHR(pNewMod->GetParentImage()->FindDataPtr(dataptr->o32_dataptr, 
                                                           min(dataptr->o32_psize, dataptr->o32_vsize), 
                                                           &pSection, &pbO32Addr));

                ASSERT(pSection);
                ASSERT(pbO32Addr);

          RETAILMSG(ZONE_FIXUP_VERBOSE, (TEXT("DoFixups: SectionAddr=0x%08lx  SectionPtr=0x%08lx  O32Addr=0x%08lx\n"),
            pSection->GetSectionHeader()->Address, pSection->GetUncompressedDataPtr(), pbO32Addr));

                CHR(pSection->BeginReloc(offset));
						break;
          }
				}
				CBRA(iLoop != pNewMod->GetE32()->e32_objcnt);
			}
			FixupAddress = (LPDWORD)(blockptr->rva - dataptr->o32_rva +
				(*currptr&0xfff) + pbO32Addr);

			RETAILMSG(ZONE_FIXUP_VERBOSE, (TEXT("DoFixups: type %d, low %8.8lx, page %8.8lx, addr %8.8lx, op %8.8lx\n"),
				(*currptr>>12)&0xf, (*currptr)&0x0fff,blockptr->rva,FixupAddress,*FixupAddress));

            switch ((*currptr >> 12) & 0xF)  {
										// No location is necessary; reference is to an absolute
										// value
				case IMAGE_REL_BASED_ABSOLUTE:
					break;

                                        // Word - (32-bits) relocate the entire address.
				case IMAGE_REL_BASED_HIGHLOW:
                    CHR(pSection->AddReloc((DWORD)FixupAddress - (DWORD)pSection->GetUncompressedDataPtr()));
					*FixupAddress += offset;
					break;

				case IMAGE_REL_BASED_HIGHADJ:
				case IMAGE_REL_BASED_MIPS_JMPADDR:
	            case IMAGE_REL_BASED_MIPS_JMPADDR16: 
				case IMAGE_REL_BASED_HIGH:
				case IMAGE_REL_BASED_LOW:
				default :
                    RETAILMSG(1, (TEXT("DoFixups : error : Unsupported fixup type!!! type=%d\n"), (*currptr >> 12) & 0xF));
                    CHRA(E_FAIL);
					break;
			}
			RETAILMSG(ZONE_FIXUP_VERBOSE, (TEXT("DoFixups: reloc complete, new op %8.8lx\n"), *FixupAddress));
			currptr++;
		}
		blockptr = (struct info *)(((ULONG)blockptr)+blockptr->size);
	}
Error:
    return hr;

} /* DoFixups()
   */

#if 0
// Old code to support ALL reloc types
// Current code ONLY supports IMAGE_REL_BASE_HIGHLOW

            switch ((*currptr >> 12) & 0xF)  {
										// MIPS relocation Types:

										// No location is necessary; reference is to an absolute
										// value
				case IMAGE_REL_BASED_ABSOLUTE:
					break;

										// Reference to the upper 16 bits of a 32-bit address.
										// Save the address and go to get REF_LO paired with this one
				case IMAGE_REL_BASED_HIGH:
					FixupAddressHi = (LPWORD)FixupAddress;
					MatchedReflo = TRUE;
					break;

										// Direct reference to a 32-bit address. With RISC
										// architecture, such an address cannot fit into a single
										// instruction, so this reference is likely pointer data
										// Low - (16-bit) relocate high part too.
				case IMAGE_REL_BASED_LOW:
					if (MatchedReflo) {
						FixupValue = (DWORD)(long)((*FixupAddressHi) << 16) +
							*(LPWORD)FixupAddress + offset;
						*FixupAddressHi = (short)((FixupValue + 0x8000) >> 16);
						MatchedReflo = FALSE;
					} else {
						FixupValue = *(short *)FixupAddress + offset;
					}
					*(LPWORD)FixupAddress = (WORD)(FixupValue & 0xffff);
					break;

										// Word - (32-bits) relocate the entire address.
				case IMAGE_REL_BASED_HIGHLOW:
					*FixupAddress += offset;
					break;

										// 32 bit relocation of 16 bit high word, sign extended
				case IMAGE_REL_BASED_HIGHADJ:
					RETAILMSG(ZONE_FIXUP, (TEXT("DoFixups: Grabbing extra data %8.8lx\n"),*(currptr+1)));
					*(LPWORD)FixupAddress += (WORD)((  SignExtendShortToLong(*(++currptr)) +offset+0x00008000)>>16);
					break;

										// Reference to the low portion of a 32-bit address.
										// jump to 26 bit target (shifted left 2)
				case IMAGE_REL_BASED_MIPS_JMPADDR:
					FixupValue = ((*FixupAddress)&0x03ffffff) + (offset>>2);
					*FixupAddress = (*FixupAddress & 0xfc000000) | (FixupValue & 0x03ffffff);
					break;
	            case IMAGE_REL_BASED_MIPS_JMPADDR16: // MIPS16 jal/jalx to 26 bit target (shifted left 2)
    	            FixupValue = (*(LPWORD)FixupAddress) & 0x03ff;
					FixupValue = (((FixupValue >> 5) | ((FixupValue & 0x1f) << 5)) << 16) | *((LPWORD)FixupAddress+1);
	                FixupValue += offset >> 2;
	                *((LPWORD)FixupAddress+1) = (WORD)(FixupValue & 0xffff);
	                FixupValue = (FixupValue >> 16) & 0x3ff;
	                *(LPWORD)FixupAddress = (WORD)((*(LPWORD)FixupAddress & 0x1c00) | (FixupValue >> 5) | ((FixupValue & 0x1f) << 5));
	                break;
				default :
					RETAILMSG(ZONE_FIXUP, (TEXT("romimage.c(DoFixups) : error : Not doing fixup type %d\n"),*currptr>>12));
					break;
			}
#endif //0





BOOL
RelocateImageFile(CModuleData *pOldMod, CModuleData *pNewMod, LPCSTR szFilename)
/*---------------------------------------------------------------------------*\
 *
\*---------------------------------------------------------------------------*/
{
    HRESULT                 hr              = NOERROR;
    HANDLE					hf              = INVALID_HANDLE_VALUE;
	ULONG					pOffset;
    e32_exe	                e32;			// exe header buffer
	DWORD					dwBytesRead;
	LONG					lPosHigh = 0;
    DWORD					dwLoop;
    DWORD                   dwFixupSize     = 0;
    LPBYTE                  pbFixupBuffer   = NULL;
    o32_obj	                o32;	

	// attempt to open image file
	CBR((hf = CreateFile(
			szFilename,
			GENERIC_READ,
			FILE_SHARE_READ | FILE_SHARE_WRITE,
			NULL,					// security
			OPEN_EXISTING,
			FILE_ATTRIBUTE_NORMAL,
			INVALID_HANDLE_VALUE
			)) != INVALID_HANDLE_VALUE);

	// seek to the offset of the e32 structure
	CBR(SetFilePointer(hf, 0x3c, &lPosHigh, FILE_BEGIN) != 0xffffffff);

    // read in the e32 offset
	CBR(ReadFile(
			hf,
			&pOffset,
			sizeof(ULONG),
			&dwBytesRead,
			NULL
			));

    // seek to the e32 structure
	CBR(SetFilePointer(hf, pOffset, &lPosHigh, FILE_BEGIN) != 0xffffffff);

    // read in the e32 structure
	CBR(ReadFile(
			hf,
			&e32,
			sizeof(e32),
			&dwBytesRead,
			NULL));
	
    CBR(HeaderMagicOk(e32));

    dwFixupSize = e32.e32_unit[FIX].size;

    if(dwFixupSize == 0) {
        // No fixups means no work to do!
        hr = NOERROR;
        goto Error;
    }

    CPR(pbFixupBuffer = (LPBYTE)LocalAlloc(LMEM_FIXED, dwFixupSize));

    //
    // Load the fixups
    //

    // Find the O32 section which contains the fixups, and load them into pbFixupBuffer
	for (dwLoop = 0; dwLoop < e32.e32_objcnt; dwLoop++) {

        // read in o32 structure
        CBR(SetFilePointer(hf, pOffset+sizeof(e32_exe)+dwLoop*sizeof(o32_obj), &lPosHigh, FILE_BEGIN) != 0xffffffff);

		CBR(ReadFile(
				hf,
				(LPBYTE)&o32,
				sizeof(o32),
				&dwBytesRead,
				NULL));

		if (o32.o32_rva <= e32.e32_unit[FIX].rva && 
		    o32.o32_rva + o32.o32_vsize >= e32.e32_unit[FIX].rva &&
		    o32.o32_rva <= e32.e32_unit[FIX].rva + e32.e32_unit[FIX].size &&
		    o32.o32_rva + o32.o32_vsize >= e32.e32_unit[FIX].rva + e32.e32_unit[FIX].size) 
        {
            CBRA(o32.o32_psize && o32.o32_dataptr);
            CBRA(o32.o32_psize >= dwFixupSize);

            CBR(SetFilePointer(hf, o32.o32_dataptr + e32.e32_unit[FIX].rva - o32.o32_rva, &lPosHigh, FILE_BEGIN) != 0xffffffff);
            CBR(ReadFile(
				    hf,
				    pbFixupBuffer,
				    dwFixupSize,
				    &dwBytesRead,
				    NULL));

            break;
		}
	}

	CloseHandle(hf);					// close the image file

    //
    // Apply the fixups
    //
    DoFixups(pOldMod, pNewMod, pbFixupBuffer, dwFixupSize);

Error:
    return hr;

} /* RelocateImageFile()
   */

HRESULT
ApplyReverseFixups(CImageData *pimgOld, CImageData *pimgNew)
/*---------------------------------------------------------------------------*\
 *
\*---------------------------------------------------------------------------*/
{
    HRESULT         hr          = NOERROR;
    CModuleData *   pOldMod     = NULL;
    CModuleData *   pNewMod     = NULL;
    int             iModOld;
    int             i;
    //
    // Loop through all modules in the new image applying reverse fixups
    //

    for(i = 0; i < pimgNew->GetModuleCount(); i++) {
        ASSERT(pNewMod == NULL);
        ASSERT(pOldMod == NULL);

        if(pNewMod = pimgNew->GetModuleData(i)) {
            LPCSTR szNewFilename = pNewMod->GetFilename();
            if(szNewFilename) {
                iModOld = pimgOld->FindModuleByName(szNewFilename);
                if(iModOld != -1) {
                    if(pOldMod = pimgOld->GetModuleData(iModOld)) {
                        e32_rom *pOldE32 = pOldMod->GetE32();
                        e32_rom *pNewE32 = pNewMod->GetE32();

                        if(pOldE32->e32_vbase != pNewE32->e32_vbase) {
                            RETAILMSG(ZONE_VERBOSE, (TEXT("Moved module: '%s'  OldVBase=0x%08lx  NewVBase=0x%08lx\n"), 
                                    szNewFilename, pOldE32->e32_vbase, pNewE32->e32_vbase));

                            RelocateImageFile(pOldMod, pNewMod, szNewFilename);
                        }

                        delete pOldMod;
                        pOldMod = NULL;
                    }
                } else {
                    RETAILMSG(ZONE_VERBOSE, (TEXT("Added module: '%s'  (No reverse fixups needed)\n"), szNewFilename));
                }
            } else {
                RETAILMSG(ZONE_VERBOSE, (TEXT("Error : Unable to get filename for new module index %d\n"), i));
                CHR(E_FAIL);
            }

            delete pNewMod;
            pNewMod = NULL;
        }
    }

    //
    // Loop through the module in the old image to detect deleted modules
    //
    for(i = 0; i < pimgOld->GetModuleCount(); i++) {
        ASSERT(pNewMod == NULL);
        ASSERT(pOldMod == NULL);

        if(pOldMod = pimgOld->GetModuleData(i)) {
            LPCSTR szOldFilename = pOldMod->GetFilename();
            if(szOldFilename) {
                if(pimgNew->FindModuleByName(szOldFilename) == -1) {
                    RETAILMSG(ZONE_VERBOSE, (TEXT("Deleted module: '%s'  (No reverse fixups needed)\n"), szOldFilename));
                }
            } else {
                RETAILMSG(ZONE_VERBOSE, (TEXT("Error : Unable to get filename for old module index %d\n"), i));
                CHR(E_FAIL);
            }

            delete pOldMod;
            pOldMod = NULL;
        }
    }

Error:
    delete pNewMod;
    delete pOldMod;

    return hr;

} /* ApplyReverseFixups()
   */
