Friday, December 05, 2008

Windows: Get Reference Count of DLL in a process

Background

Windows Dynamic Link Library (DLL) or Delphi Runtime Package (BPL) are both shared library for Windows OS.  As stated in MSDN:

The system maintains a per-process reference count on all loaded modules. Calling LoadLibrary increments the reference count. Calling the FreeLibrary or FreeLibraryAndExitThread function decrements the reference count. The system unloads a module when its reference count reaches zero or when the process terminates (regardless of the reference count).

If a DLL has been loaded for more than one time, Windows OS doesn't maintain separate copies of DLL image in the OS system.  Only one DLL image is loaded and Windows OS use reference counting strategy to maintain the status.  Once the reference count of a DLL has dropped to zero, the DLL image is completely remove from the process.

A bad news is there isn't any documented Windows API to allow us to retrieve the current reference count of a DLL.

But why I need to know the reference count of a DLL?  It is very rare for us to know this information.  Someone use this technique to kill spyware DLL.

I design my application with runtime packages.  For example, I keep forms in runtime package and load the form from package when I open the form.  The runtime package is not bind to my executable image statically but it will be invoked dynamically via LoadPackage.

I want to achieve these:

  1. The package isn't loaded when there isn't any form instance in the process.
  2. We may instantiate more than one form instance from same package in the process at the same time.
  3. Once all form instances has been destroyed, the runtime package image will remove from the process.

To get it done. I declare an export procedure that return a new form instance everytime it is invoked.  Each form instance created will invoke LoadPackage(HInstance) to increase the reference count.

The most complex part is the 3rd point.  An exception will raise If I invoke UnloadPackage before all code in package has been executed completely.  I not even allow to invoke UnloadPackage in form's destructor.  To overcome this problem, I create an notifier for each form instance.  When the form instance is destroying, it will alert the notifier.  The notifier will send a message using PostMessage to notify a controller to unload the package.  The PostMessage will get execute only when the process is in idle state. This method has work so well.

Until then I encounter a situation where I have this coding:

var i: integer;
    F: TForm;
    H: THandle;
    NewForm: function: TForm;
begin
  H := LoadPackage('MyPackage.bpl');
  @NewForm := GetProcAddress(H, 'NewForm');
  for i := 1 to 10 do begin
    F := NewForm;
    F.Free;
  end;
  UnloadPackage(H);
end;

The above code works fine not until I increase the for loop count to 10000.  I will be prompted

000101c95520$d53ae3d0$_CDOSYS2.0

The reason is simple, each call to F.Free will invoke PostMessage to queue the message to the associated thread of the process.  If I run the above code to 10000 times continuously, the system just doesn't has a free moment to perform the messages in queue.  It will certainly reach a limit until the message prompt out.

I have revise the strategy to overcome such situation but I am not confident enough to state that the new strategy will not having other problems in future.  I use unit testing to write test cases for all known situations, one of the test case will perform 10000 loops and I will then check the reference count of the runtime packages to make sure it falls into a suitable range.

After writing such a long description, I finally manage the deliver the situations I face in words and paragraphs.

Solution

After spending a day digging on Internet about retrieving the reference count of a DLL, I found the solution.  I have coding it in Delphi:

windows.PEB.pas

unit Windows.PEB;

interface

uses Windows{,
     untSttUnhooker}; // WinAPI UnHooker - by StTwister

type
  PNtAnsiString = ^TNtAnsiString;
  TNtAnsiString = packed record
    Length: Word;
    MaximumLength: Word;
    Buffer: PAnsiChar;
  end;

type
  PNtUnicodeString = ^TNtUnicodeString;
  TNtUnicodeString = packed record
    Length: Word;
    MaximumLength: Word;
    Buffer: PWideChar;
  end;

type
  PClientId = ^TClientId;
  TClientId = record
    UniqueProcess: THandle;
    UniqueThread: THandle;
  end;

type
  PCurDir = ^TCurDir;
  TCurDir = packed record
    DosPath: TNtUnicodeString;
    Handle : THandle;
  end;

type
  PRtlDriveLetterCurDir = ^TRtlDriveLetterCurDir;
  TRtlDriveLetterCurDir = packed record
    Flags    : Word;
    Length   : Word;
    TimeStamp: Cardinal;
    DosPath  : TNtAnsiString;
  end;

type
  PRtlUserProcessParameters = ^TRtlUserProcessParameters;
  TRtlUserProcessParameters = record
    MaximumLength    : Cardinal;
    Length           : Cardinal;
    Flags            : Cardinal;
    DebugFlags       : Cardinal;
    ConsoleHandle    : THandle;
    ConsoleFlags     : Cardinal;
    StandardInput    : THandle;
    StandardOutput   : THandle;
    StandardError    : THandle;
    CurrentDirectory : TCurDir;
    DllPath          : TNtUnicodeString;
    ImagePathName    : TNtUnicodeString;
    CommandLine      : TNtUnicodeString;
    Environment      : Pointer;
    StartingX        : Cardinal;
    StartingY        : Cardinal;
    CountX           : Cardinal;
    CountY           : Cardinal;
    CountCharsX      : Cardinal;
    CountCharsY      : Cardinal;
    FillAttribute    : Cardinal;
    WindowFlags      : Cardinal;
    ShowWindowFlags  : Cardinal;
    WindowTitle      : TNtUnicodeString;
    DesktopInfo      : TNtUnicodeString;
    ShellInfo        : TNtUnicodeString;
    RuntimeData      : TNtUnicodeString;
    CurrentDirectores: Array [0..31] of TRtlDriveLetterCurDir;
  end;

type
  PPebFreeBlock = ^TPebFreeBlock;
  TPebFreeBlock = record
    Next: PPebFreeBlock;
    Size: Cardinal;
  end;

type
  PLdrModule = ^TLdrModule;
  TLdrModule = packed record
    InLoadOrderModuleList          : TListEntry;      // 0h
    InMemoryOrderModuleList        : TListEntry;      // 8h
    InInitializationOrderModuleList: TListEntry;      // 10h
    BaseAddress                    : THandle;         // 18h
    EntryPoint                     : THandle;         // 1Ch
    SizeOfImage                    : Cardinal;        // 20h
    FullDllName                    : TNtUnicodeString;// 24h
                                   // Length (2)         24h
                                   // MaximumLength (2)  26h
                                   // Buffer (4)         28h
    BaseDllName                    : TNtUnicodeString;// 2Ch
    Flags                          : ULONG;           // 34h
    LoadCount                      : SHORT;           // 38h
    TlsIndex                       : SHORT;           // 3Ah
    HashTableEntry                 : TListEntry;      // 3Ch
    TimeDataStamp                  : ULONG;           // 44h
  end;

type
  PPebLdrData = ^TPebLdrData;
  TPebLdrData = packed record
    Length                         : Cardinal;        // 0h
    Initialized                    : LongBool;        // 4h
    SsHandle                       : THandle;         // 8h
    InLoadOrderModuleList          : TListEntry;      // 0Ch
    InMemoryOrderModuleList        : TListEntry;      // 14h
    InInitializationOrderModuleList: TListEntry;      // 1Ch
  end;

type
  PPeb = ^TPeb;
  TPeb = packed record
    InheritedAddressSpace         : Boolean;
    ReadImageFileExecOptions      : Boolean;
    BeingDebugged                 : Boolean;
    SpareBool                     : Boolean;
    Mutant                        : Pointer;
    ImageBaseAddress              : Pointer;
    Ldr                           : PPebLdrData;
    ProcessParameters             : PRtlUserProcessParameters;
    SubSystemData                 : Pointer;
    ProcessHeap                   : Pointer;
    FastPebLock                   : Pointer;
    FastPebLockRoutine            : Pointer;
    FastPebUnlockRoutine          : Pointer;
    EnvironmentUpdateCount        : Cardinal;
    KernelCallbackTable           : Pointer;
    case Integer of
      4: (
        EventLogSection           : Pointer;
        EventLog                  : Pointer);
      5: (
        SystemReserved            : Array [0..1] of Cardinal;
  { end; }
    FreeList                      : PPebFreeBlock;
    TlsExpansionCounter           : Cardinal;
    TlsBitmap                     : Pointer;
    TlsBitmapBits                 : Array [0..1] of Cardinal;
    ReadOnlySharedMemoryBase      : Pointer;
    ReadOnlySharedMemoryHeap      : Pointer;
    ReadOnlyStaticServerData      : ^Pointer;
    AnsiCodePageData              : Pointer;
    OemCodePageData               : Pointer;
    UnicodeCaseTableData          : Pointer;
    NumberOfProcessors            : Cardinal;
    NtGlobalFlag                  : Cardinal;
    Unknown                       : Cardinal;
    CriticalSectionTimeout        : TLargeInteger;
    HeapSegmentReserve            : Cardinal;
    HeapSegmentCommit             : Cardinal;
    HeapDeCommitTotalFreeThreshold: Cardinal;
    HeapDeCommitFreeBlockThreshold: Cardinal;
    NumberOfHeaps                 : Cardinal;
    MaximumNumberOfHeaps          : Cardinal;
    ProcessHeaps                  : ^Pointer;
    GdiSharedHandleTable          : Pointer;
    ProcessStarterHelper          : Pointer;
    GdiDCAttributeList            : Cardinal;
    LoaderLock                    : Pointer;
    OSMajorVersion                : Cardinal;
    OSMinorVersion                : Cardinal;
    OSBuildNumber                 : Word;
    OSCSDVersion                  : Word;
    OSPlatformId                  : Cardinal;
    ImageSubsystem                : Cardinal;
    ImageSubsystemMajorVersion    : Cardinal;
    ImageSubsystemMinorVersion    : Cardinal;
    ImageProcessAffinityMask      : Cardinal;
    GdiHandleBuffer               : Array [0..33] of Cardinal;
    PostProcessInitRoutine        : ^Pointer;
    TlsExpansionBitmap            : Pointer;
    TlsExpansionBitmapBits        : Array [0..31] of Cardinal;
    SessionId                     : Cardinal;
    AppCompatInfo                 : Pointer;
    CSDVersion                    : TNtUnicodeString);
  end;

type
  PNtTib = ^TNtTib;
  TNtTib = record
    ExceptionList       : Pointer;  // ^_EXCEPTION_REGISTRATION_RECORD
    StackBase           : Pointer;
    StackLimit          : Pointer;
    SubSystemTib        : Pointer;
    case Integer of
      0: (FiberData     : Pointer);
      1: (Version       : ULONG;
    ArbitraryUserPointer: Pointer;
    Self                : PNtTib);
  end;

type
  PTeb = ^TTeb;
  TTeb = record
    Tib               : TNtTib;
    Environment       : PWideChar;
    ClientId          : TClientId;
    RpcHandle         : THandle;
    ThreadLocalStorage: PPointer;
    Peb               : PPeb;
    LastErrorValue    : DWORD;
  end;

implementation

end.

Windows.ntdll.pas

unit Windows.ntdll;

interface

uses Windows;

const
  ntdll = 'ntdll.dll';

type
  NTSTATUS = LongInt;

  TProcessInfoClass = (
    ProcessBasicInformation, ProcessQuotaLimits,
    ProcessIoCounters, ProcessVmCounters,
    ProcessTimes, ProcessBasePriority, ProcessRaisePriority,
    ProcessDebugPort, ProcessExceptionPort,
    ProcessAccessToken, ProcessLdtInformation,
    ProcessLdtSize, ProcessDefaultHardErrorMode,
    ProcessIoPortHandlers, ProcessPooledUsageAndLimits,
    ProcessWorkingSetWatch, ProcessUserModeIOPL,
    ProcessEnableAlignmentFaultFixup, ProcessPriorityClass,
    ProcessWx86Information, ProcessHandleCount,
    ProcessAffinityMask, ProcessPriorityBoost,
    ProcessDeviceMap, ProcessSessionInformation,
    ProcessForegroundInformation, ProcessWow64Information,
    MaxProcessInfoClass
  );
  PROCESSINFOCLASS = TProcessInfoClass;

  PPROCESS_BASIC_INFORMATION = ^PROCESS_BASIC_INFORMATION;

  PROCESS_BASIC_INFORMATION = packed record
    ExitStatus:         DWORD;
    PebBaseAddress:     Pointer;
    AffinityMask:       DWORD;
    BasePriority:       DWORD;
    UniqueProcessId:    DWORD;
    InheritedUniquePID: DWORD;
  end;

  TZwQueryInformationProcess = function(ProcessHandle: THandle;
                                   ProcessInformationClass: PROCESSINFOCLASS;
                                   var ProcessInformation: PROCESS_BASIC_INFORMATION;
                                   ProcessInformationLength: ULONG;
                                   var ReturnLength: ULONG): NTSTATUS; stdcall;

implementation

end.

WindowsEx.pas

unit WindowsEx;

interface

uses Windows, Windows.PEB;

type
  TProcessModules = class
  private
    FPebLDrData: TPebLdrData;
    FReady: boolean;
    FProcessHandle: THandle;
    FCurrentIndex: Pointer;
    FCurrentModule: TLdrModule;
    procedure CheckReady;
  public
    constructor Create(const aProcessHandle: THandle = 4294967295);
    procedure Reset;
    function HasNext: boolean;
    function CurrentModule: TLdrModule;
  end;

implementation

uses SysUtils, Windows.ntdll;

procedure TProcessModules.CheckReady;
begin
  if not FReady then
    raise Exception.Create('PEB Loader Data is not ready');
end;

constructor TProcessModules.Create(const aProcessHandle: THandle = 4294967295);
var H: THandle;
    Proc: TZwQueryInformationProcess;
    PBI: PROCESS_BASIC_INFORMATION;
    i: ULONG;
    PEB: TPeb;
begin
  FProcessHandle := aProcessHandle;
  if FProcessHandle = THandle(-1) then
    FProcessHandle := GetCurrentProcess;

  H := LoadLibrary(ntdll);
  if H = 0 then RaiseLastOSError;
  try

    @Proc := GetProcAddress(H, 'ZwQueryInformationProcess');
    if @Proc = nil then RaiseLastOSError;

    ZeroMemory(@PBI, SizeOf(PBI));
    FReady := Proc(FProcessHandle, ProcessBasicInformation, PBI, SizeOf(PBI), i) = 0;
    if FReady then begin
      ReadProcessMemory(FProcessHandle, PBI.PebBaseAddress, @PEB, 16, i);
      ReadProcessMemory(FProcessHandle, PEB.Ldr, @FPebLDrData, sizeof(FPebLDrData), i);

      Reset;
    end;
  finally
    FreeLibrary(H);
  end;
end;

function TProcessModules.CurrentModule: TLdrModule;
begin
  Result := FCurrentModule;
end;

function TProcessModules.HasNext: boolean;
var i: ULONG;
begin
  CheckReady;

  Result := ReadProcessMemory(GetCurrentProcess, FCurrentIndex, @FCurrentModule, SizeOf(FCurrentModule), i);
  if Result then begin
    Result := FCurrentModule.BaseAddress <> 0;
    if Result then
      FCurrentIndex := FCurrentModule.InLoadOrderModuleList.Flink;
  end;

  if not Result then
    FCurrentIndex := nil;
end;

procedure TProcessModules.Reset;
begin
  CheckReady;
  ZeroMemory(@FCurrentModule, SizeOf(FCurrentModule));
  FCurrentIndex := FPebLDrData.InLoadOrderModuleList.Flink;
end;

end.

Example:

var Proc: TProcessModules;
    S: string;
begin
  Proc := TProcessModules.Create;
  try
    while Proc.HasNext do begin
      S := Format('%6d : %s', [Proc.CurrentModule.LoadCount, Proc.CurrentModule.FullDllName.Buffer]);
      Memo1.Lines.Add(S);
    end;
  finally
    Proc.Free;
  end;
end;

Reference

  1. 0×04 Reference count of DLL
  2. Reference Count of DLL in a Process

2 comments:

Caique said...

This solve a Count Reference problem, but don't solve PostMessage problem. How you solve this ?

"excuse my poor english"

Chau Chee Yang said...

@Caique, I have some classes written to manage the state of package. That would be another lengthy topic. This article is about getting reference counter of DLL in a process. I did test cases that simulate all kind of usage patterns about runtime packages that I can think of. In the test case, I check the reference count of packages if it exceed limit that may cause runtime error. It the reference counter fit in a safe value range, it simply means my package classes works.