//==================================
// Fix1MB.C - Matt Pietrek 1995
//==================================
#include <windows.h>
#include <windowsx.h>
#include <string.h>
#include <stdlib.h>
#include <toolhelp.h>
#pragma hdrstop
#include "Fix1MB.h"
#include "prochook.h"

HANDLE HInstance;
HWND HWndUsedLb;
HWND HWndStatus;
HWND HWndDlg;

BOOL FWorstCaseMode;

DWORD EndOfLastFixedBlock;  // linear address of end of last fixed block
DWORD FreeBlocksTotalSize;  // Byte count of free regions < 1Mb

LRESULT CALLBACK _export Fix1MBDlgProc( HWND hWndDlg, UINT msg,
                                        WPARAM wParam, LPARAM lParam );
void GetModuleNameFromHandle(HANDLE handle, char *owner);
void GetModuleFilenameFromHandle(HANDLE handle, char *owner);
void AddHeapEntryToListbox(GLOBALENTRY *ge);
void WalkHeap( void );
int HandleSort(const void far *a, const void far *b);
BOOL IsPotentialSpaceHog( GLOBALENTRY *ge );
void Handle_WM_COMMAND(HWND hWndDlg, WPARAM wParam, LPARAM lParam);
void Handle_WM_INITDIALOG(HWND hWndDlg);
void InitFilteredModules(void);
BOOL IsFilteredModule(HMODULE hModule);
BOOL HookLoadModule(void);
BOOL UnhookLoadModule(void);
BOOL AllocTiledBelow1MbMemory(void);
BOOL FreeTiledBelow1MbMemory(void);
void AddToSystemIni(void);
void HelpDialog( HWND hWndOnwer );
VOID CenterWindow(HWND hWnd);

int PASCAL WinMain( HANDLE hInstance, HANDLE hPrevInstance,
                    LPSTR lpszCmdLine, int nCmdShow )
{
    WNDCLASS wndclass;
    char szWorstCase[32];
    
    if ( hPrevInstance )
        return 0;

    HInstance = hInstance;

    if ( !GetClassInfo( 0, MAKEINTRESOURCE(32770), &wndclass ) )
        return 0;

    wndclass.hInstance = HInstance;
    wndclass.hIcon = LoadIcon(HInstance,"Fix1MBIcon");
    wndclass.lpszClassName = "FIX1MBDLG";
    if ( !RegisterClass(&wndclass) )
        return 0;

    InitFilteredModules();

    FWorstCaseMode = GetProfileInt( "Fix1MB", "WorstCaseMode", FALSE );

    if ( !HookLoadModule() )
    {
        MessageBox( 0, "Couldn't hook WinExec and LoadModule", 0, MB_OK );
        return 0;
    }
        
    DialogBox( HInstance, "FIX1MBDLG", 0, (DLGPROC)Fix1MBDlgProc );
        
    UnhookLoadModule();

    // Write out the current WorstCase flag value
    wsprintf( szWorstCase,"%u", FWorstCaseMode );
    WriteProfileString( "Fix1MB", "WorstCaseMode", szWorstCase );
    
    return 0;
}

//########################################################################
// This section deals with watching calls to LoadModule and sucking up
// all the low DOS memory beforehand.
//########################################################################

NPHOOKCHILD npHookLoadModule;

HINSTANCE WINAPI __loadds Fix1MBLoadModule( LPCSTR lpszModuleName,
                                            LPVOID lpvParameterBlock )
{
    HINSTANCE retValue;
    
    #ifdef __BORLANDC__
    __asm   push    DS
    __asm   mov     ax, seg HInstance   // BC++ 4.5 !!! STILL !!! doesn't
    __asm   mov     ds, ax              // do __loadds properly in EXEs!
    #endif

    ProcUnhook( npHookLoadModule );
    
    AllocTiledBelow1MbMemory();
    
    retValue = LoadModule( lpszModuleName, lpvParameterBlock );

    FreeTiledBelow1MbMemory();
    
    ProcHook( npHookLoadModule );
    
    #ifdef __BORLANDC__
    __asm   pop     ds
    #endif

    return retValue;
}

BOOL HookLoadModule(void)
{
    npHookLoadModule = SetProcAddress(  (FARPROC)LoadModule,
                                        (FARPROC)Fix1MBLoadModule,
                                        FALSE );
    
    return (BOOL)npHookLoadModule;
}

BOOL UnhookLoadModule(void)
{
    SetProcRelease( npHookLoadModule );
    
    return TRUE;
}


HANDLE ListHead = 0;
    
BOOL AllocTiledBelow1MbMemory(void)
{
    HANDLE h;
    DWORD blockSize;
    
    if ( FWorstCaseMode )
        blockSize = 0x200;
    else
        blockSize = 0x8000;  

    // Allocate it all!
        
top:
    while ( h = (HANDLE)GlobalDosAlloc(blockSize) )
    {
        #if 0
        char szTemp[128];
        wsprintf( szTemp, "handle %04X = %08lX (size=%04X)\r\n",
                    h, GetSelectorBase(h), blockSize );
        OutputDebugString( szTemp );
        #endif

        *(LPDWORD)GlobalLock(h) = ListHead;
        ListHead = h;
    }

    if ( FWorstCaseMode )
    {
        // We have a whole mess of 0x200 byte blocks.  Free every other one
        h = ListHead;

        while ( h )
        {
            HANDLE hNext, hNext2;

            // Get the next block handle in the list
            hNext = *(LPHANDLE)GlobalLock( h );

            if ( !hNext )   // If it's NULL, we're done
                break;

            // Get the next block of the next block ( h->next->next )
            hNext2 = *(LPHANDLE)GlobalLock( hNext );

            // Point the "next" field of the curr block to the next, next block
            *(LPHANDLE)GlobalLock( h ) = hNext2;

            // Delete the next block
            GlobalDosFree( hNext );

            // Advance to the next, next block
            h = hNext2;
        }
    }
    else
    {
        blockSize = blockSize >> 1;
        if ( blockSize > 0x200 )    // Grab all blocks > 0x200
            goto top;

        // Free the first block in the list, and readjust the ListHead
        if ( ListHead )
        {
            h = *(LPWORD)GlobalLock( ListHead );
            GlobalDosFree( ListHead );
            ListHead = h;
        }
    }

    return TRUE;
}

BOOL FreeTiledBelow1MbMemory(void)
{
    HANDLE h, h2;

    for ( h2 = h = ListHead; h2; h = h2 )
    {
        h2 = *(LPWORD)GlobalLock(h);
        GlobalDosFree( h );
    }

    ListHead = 0;
    
    return TRUE;
}

//########################################################################
// This section walks the low heap and determines how much frees space
// there is and identifies potential low mem space hogs.
//########################################################################

typedef struct
{
    HGLOBAL     hGlobal;
    DWORD       size;
} LOW_1MB_BLOCK, far *LPLOW_1MB_BLOCK;

//
// Walk the heap below one megabyte.  The command parameter indicates
// whether the listboxes should be updated, or if the memory allocated
// by the "Allocate All" button should be freed.
//
void WalkHeap( void )
{
    GLOBALENTRY ge;
    BOOL moreToGo;
    char buffer[128];
    DWORD lastBelowOneMeg = 0;
    LPLOW_1MB_BLOCK lpBlockArray;
    unsigned i, cHandles;

    // Allocate memory for an array of handles that we'll need to sort
    lpBlockArray = (LPLOW_1MB_BLOCK)GlobalAllocPtr( GMEM_MOVEABLE,
                                            sizeof(LOW_1MB_BLOCK) * 0x2000 );
    if ( !lpBlockArray )
        return;
    cHandles = 0;
    
    ge.dwSize = sizeof(ge);
    moreToGo = GlobalFirst( &ge, GLOBAL_ALL );
    EndOfLastFixedBlock = ge.dwAddress;
    FreeBlocksTotalSize = 0;
    
    while ( moreToGo )
    {
        // Break out of the loop when we encounter a block above 1 Meg
        if ( ge.dwAddress>=0x100000LU )
        {
            ge.dwAddress = lastBelowOneMeg;
            ge.hBlock = 0xFFFF;     // A pseudo FIXED handle
            ge.wType = GT_SENTINEL;
            IsPotentialSpaceHog( &ge );
            break;  // Break out of loop
        }
        
        if ( IsPotentialSpaceHog( &ge ) )
        {
            lpBlockArray[cHandles].size = ge.dwBlockSize;
            lpBlockArray[cHandles++].hGlobal = ge.hBlock;
        }
        
        lastBelowOneMeg = ge.dwAddress + ge.dwBlockSize;
        moreToGo = GlobalNext(&ge, GLOBAL_ALL);
    }

    qsort( lpBlockArray, cHandles, sizeof(LOW_1MB_BLOCK), HandleSort );

    for ( i=0; i < cHandles; i++ )
    {
        GLOBALENTRY ge;
        
        ge.dwSize = sizeof(ge);
        
        if ( GlobalEntryHandle( &ge, lpBlockArray[i].hGlobal ) )
            AddHeapEntryToListbox( &ge );
    }

    GlobalFreePtr( lpBlockArray );

    // update the Total Free Space field
    wsprintf(buffer, "Potential Free Space: %lu bytes (%luK)",
        FreeBlocksTotalSize, FreeBlocksTotalSize/1024);
    SetWindowText(HWndStatus, buffer);
}

int HandleSort(const void far *a, const void far *b)
{
    DWORD sizeA, sizeB;

    sizeA = ((LPLOW_1MB_BLOCK)a)->size;
    sizeB = ((LPLOW_1MB_BLOCK)b)->size;

    if ( sizeA == sizeB )
        return 0;
    return (sizeB > sizeA) ? 1 : -1;
}

//
// Given a global handle, get the module name of its owner
//
void GetModuleNameFromHandle(HANDLE handle, char *pszOwner)
{
    MODULEENTRY me;
    TASKENTRY   te;

    te.dwSize = sizeof(te);
    if ( TaskFindHandle(&te, handle) )  // First see if a task owns it
    {
        lstrcpy(pszOwner, te.szModule);
        return;
    }

    me.dwSize = sizeof(me);
    if (ModuleFindHandle(&me, handle))  // Nope?  Try a module owner
    {
        lstrcpy(pszOwner, me.szModule);
        return;
    }
    
    pszOwner[0] = 0;   // No owner found.  Return empty string
}

//
// Given a global handle, get the file name of its owner
//
void GetModuleFilenameFromHandle(HANDLE handle, char *pszOwner)
{
    MODULEENTRY me;
    TASKENTRY   te;

    te.dwSize = sizeof(te);
    if ( TaskFindHandle(&te, handle) )  // First see if a task owns it
    {
        me.dwSize = sizeof(me);
        ModuleFindHandle( &me, te.hModule );

        lstrcpy( pszOwner, me.szExePath );
        return;
    }

    me.dwSize = sizeof(me);
    if (ModuleFindHandle(&me, handle))  // Nope?  Try a module owner
    {
        lstrcpy(pszOwner, me.szExePath);
        return;
    }
    
    pszOwner[0] = 0;   // No owner found.  Return empty string
}

// More descriptive strings for the TOOLHELP GT_xxx #define's
char *BlockTypes[] = {
"ALLOC", "DGROUP", "DATA", "CODE", "TDB", "RESOURCE", "MODULE",
"FREE", "INTERNAL", "SENTINEL", "BURGERMASTER" };

//
// Given a new memory block, see if is fixed or locked.  If so, it
// may indicate the end of a free region.  If so, add the info
// about the free region to the left listbox.  Then add the block's
// info to the right listbox.
//
void AddHeapEntryToListbox(GLOBALENTRY *ge)
{
    char buffer[128];
    char szFileName[260];
    char owner[32];

    #if 1
    GetModuleFilenameFromHandle( ge->hOwner, szFileName );
    wsprintf( buffer,"%05lu bytes (%s %s) %s",
            ge->dwBlockSize, BlockTypes[ge->wType], 
            (ge->hBlock & 1) ? "Fixed" : "Locked",
            szFileName );
    #else
    // Show the segment's handle and module name
    GetModuleNameFromHandle(ge->hOwner, owner);
    wsprintf( buffer,"%05lu bytes (%04X) (%s %s %s)",
            ge->dwBlockSize, ge->hBlock, owner, BlockTypes[ge->wType], 
            (ge->hBlock & 1) ? "Fixed" : "Locked" );
    #endif

    SendMessage(HWndUsedLb,LB_ADDSTRING,0,(LPARAM)(LPSTR)buffer);
}

BOOL IsPotentialSpaceHog( GLOBALENTRY *ge )
{
    if ( ge->wType == GT_FREE )     // If not a free block...
        return FALSE;
    
    if ( !(ge->hBlock & 1) )        // Is it a moveable block (low bit is off)
    {                               // Yes?  Then check for locks.
        if ( !(ge->wcLock) && !(ge->wcPageLock) )
            return FALSE;
    }
    
    // At this point, we should only have immoveable blocks.
    // Was there a gap between the end of the last fixed block
    // and the start of this one?  If so, there was a free region.
    // Add the relevant information to the left listbox
    if ( ge->dwAddress > EndOfLastFixedBlock )
    {
        DWORD freeBlockLen = ge->dwAddress-EndOfLastFixedBlock;
        FreeBlocksTotalSize += freeBlockLen;
    }
            
    EndOfLastFixedBlock = ge->dwAddress + ge->dwBlockSize;

    // Now throw out blocks that shouldn't appear in the space hog list
    if ( !ge->hBlock )              // Throw out sentinel blocks
        return FALSE;               // with no handles
    if ( ge->wType == GT_TASK )
        return FALSE;
    if ( (ge->hBlock == 0xFFFF) && (ge->wType == GT_SENTINEL) )
        return FALSE;
    if ( IsFilteredModule(ge->hOwner) )
        return FALSE;

    return TRUE;
}

//########################################################################
// This section contains helps us filter out modules that the user can't
// do anything about (i.e., system DLLs)
//########################################################################

typedef struct
{
    LPSTR   lpszModName;
    HMODULE hModule;
} FILTERED_MODULE, FAR * LPFILTERED_MODULE;

FILTERED_MODULE FilteredModules[] = 
{
{ "KERNEL",     0 },
{ "SYSTEM",     0 },
{ "KEYBOARD",   0 },
{ "MOUSE",      0 },
{ "DISPLAY",    0 },
{ "SOUND",      0 },
{ "COMM",       0 },
{ "FONTS",      0 },
{ "OEMFONTS",   0 },
{ "GDI",        0 },
{ "USER",       0 },
{ "TOOLHELP",   0 },
{ "WIN87EM",    0 },
{ "MMSYSTEM",   0 },
{ "TIMER",      0 },
{ "WINOLDAP",   0 },
};

unsigned cFilteredModules = sizeof(FilteredModules) / sizeof(FILTERED_MODULE);

void InitFilteredModules(void)
{
    unsigned i;
    
    for ( i=0; i < cFilteredModules; i++ )
    {
        FilteredModules[i].hModule
            = (HMODULE)GetModuleHandle( FilteredModules[i].lpszModName );
    }
}

BOOL IsFilteredModule(HMODULE hModule)
{
    unsigned i;
    
    for ( i=0; i < cFilteredModules; i++ )
    {
        if ( hModule == FilteredModules[i].hModule )
            return TRUE;
    }
    
    return FALSE;
}


//########################################################################
// This section contains the DlgProc code and message handlers.
//########################################################################

//
// Dialog proc for the main dialog
//
LRESULT CALLBACK _export Fix1MBDlgProc(HWND hWndDlg, UINT msg,
                                       WPARAM wParam, LPARAM lParam)
{
    switch ( msg )
    {
        case WM_COMMAND:
            Handle_WM_COMMAND(hWndDlg, wParam, lParam);
            break;
        case WM_INITDIALOG:
            Handle_WM_INITDIALOG(hWndDlg);
            break;
        case WM_SIZE:
            if ( wParam == SIZE_RESTORED )
                PostMessage( hWndDlg, WM_COMMAND, IDC_BUTTON_REWALK, 0 );
            break;
        case WM_CLOSE:
            EndDialog( hWndDlg, 0 );
            break;
    }    

    return 0;
}

//
// Handle the pressing of any of the 3 buttons.
//
void Handle_WM_COMMAND(HWND hWndDlg, WPARAM wParam, LPARAM lParam)
{
    BOOL updateDisplay = FALSE;
    HCURSOR oldCursor = 0;      // Some operations take awhile...
    
    switch ( wParam )
    {   
        case IDC_BUTTON_REWALK:
            oldCursor = SetCursor( LoadCursor(0, IDC_WAIT) );
            updateDisplay = TRUE;
            break;
        case IDC_HELP_FIX1MB:
            HelpDialog( hWndDlg );
            break;
        case IDC_ADD_TO_SYS_INI:
            AddToSystemIni();
            break;
        case IDC_WORST_CASE_MODE:
            if ( HIWORD(lParam) == BN_CLICKED )
                FWorstCaseMode
                    = IsDlgButtonChecked( hWndDlg, IDC_WORST_CASE_MODE );
            break;
    }

    //
    // Clear out the listboxes, set immediate updating to false,
    // and walk the heap.  Afterwards, turn listbox updating back on.
    //
    if ( updateDisplay )
    {
        SendMessage(HWndUsedLb, LB_RESETCONTENT, 0, 0);         
        SendMessage(HWndUsedLb, WM_SETREDRAW, FALSE, NULL);
        WalkHeap();
        SendMessage(HWndUsedLb, WM_SETREDRAW, TRUE, NULL);
    }
    if ( oldCursor )            // Switch back to the normal cursor
        SetCursor(oldCursor);
}

void Handle_WM_INITDIALOG(HWND hWndDlg)
{
    HWndDlg = hWndDlg;

    CenterWindow( hWndDlg );

    // Get the HWND's of the commonly modified controls
    HWndUsedLb = GetDlgItem(hWndDlg, IDC_LISTBOX_USED);
    HWndStatus = GetDlgItem(hWndDlg, IDC_TEXT_STATUS);

    CheckDlgButton( HWndDlg, IDC_WORST_CASE_MODE, FWorstCaseMode );

    ShowWindow( HWndDlg, SW_MINIMIZE );
    // Fake a "Re-Walk" button event
    Handle_WM_COMMAND(hWndDlg, IDC_BUTTON_REWALK, 0); 
}

void AddToSystemIni(void)
{
    char szPathToDll[ 768 ];
    LPSTR p;
    
    GetModuleFileName( HInstance, szPathToDll, sizeof(szPathToDll) );
    p = strrchr( szPathToDll, '\\' );
    
    if ( p )
    {
        LPSTR lpszRemainder;
        
        // first build the path to FX1MBDLL.DLL
        lstrcpy( p+1, "FX1MBDLL.DLL " );
        
        lpszRemainder = szPathToDll + lstrlen(szPathToDll);
        
        // Then tack on the existing entries from the "drivers=" line
        // from the [boot] section of SYSTEM.INI
        GetPrivateProfileString("boot", "drivers", "", lpszRemainder,
                                512, "SYSTEM.INI" );

        if ( 0 == strstr(lpszRemainder, "FX1MBDLL") )
        {
            // Write out the revised string
            WritePrivateProfileString("boot", "drivers", szPathToDll,
                                        "SYSTEM.INI");
        }
        else
        {
            MessageBox(0, "Fix1MB is already in the SYSTEM.INI file", 0,MB_OK);
        }
    }
    else
    {
        MessageBox( 0, "Error adding FX1MBDLL.DLL to SYSTEM.INI", 0, MB_OK );
    }
}

void HelpDialog( HWND hWndOnwer )
{
    MessageBox( hWndOnwer,
    "Fix1MB is a tool to help with the dreaded \"out of memory\" errors"
    " in Windows that cause new programs to be unable to start.  This "
    "particular problem is caused by DLLs that inadvertantly suck up "
    "all the memory below 1 megabyte in the Windows address space.  Windows "
    "needs a certain amount of memory below 1 megabyte to start a new "
    "task.\r\n\r\n"
    "Fix1MB not only shows you which DLLs and programs are using this "
    "precious memory, it also acts to prevent them from grabbing the "
    "memory below 1 megabyte.\r\n\r\n"
    "Fix1MB can either be run from within Windows, or loaded at startup "
    "time.  The latter allows Fix1MB to preserve even more memory below "
    "1 megabyte.  The \"Add FIX1MB to SYSTEM.INI\" button will put Fix1MB in "
    "your SYSTEM.INI if you desire (the FIX1MB.EXE, FX1MBDLL.DLL and "
    "PROCHOOK.DLL files must be in the same directory for this to "
    "work.)\r\n\r\n"
    "Fix1MB was written by Matt Pietrek (CIS: 71774,362), and is from his "
    "May 1995 Questions and Answers column in the Microsoft Systems "
    "Journal.  Please refer to that column for additional information.",
    "Fix1MB Help", MB_OK );
}

VOID CenterWindow(HWND hWnd)
{
    RECT rect;
    WORD wWidth, wHeight;

    GetWindowRect(hWnd,&rect);

    wWidth =GetSystemMetrics(SM_CXSCREEN);
    wHeight=GetSystemMetrics(SM_CYSCREEN);

    MoveWindow( hWnd,   (wWidth/2) - ((rect.right -  rect.left)/2),
                        (wHeight/2) - ((rect.bottom - rect.top) /2),
                        rect.right - rect.left,
                        rect.bottom - rect.top,
                        FALSE );
}