Problems to redirect the TObject.AfterConstruction to other procedure

delphi, delphi-xe

Solution

EDITED - Now working to increase and decrease the number of items. To make it work is just to put the unit as the first unit of your dpr. Now, I'll just optimize some methods and put here the outputs I want. (I'll not reedit the post, not necessary) But if you'd like to use, fell free to test and report bugs. I put a simple out put if you'd like to test, the procedure SaveInstancesToFile, it creates a test.txt file in your application path with the output of counters.

unit ObjectCounter;

  {  Develop by rodrigofrezino@gmail.com
     Stackoverflow: http://stackoverflow.com/users/225010/saci
     Please, any bug let me know}

interface

  procedure SaveInstancesToFile;

implementation

uses
  Windows, SysUtils, Classes, TypInfo;

type

  PClassVars = ^TClassVars;
  TClassVars = class(TObject)
  private
    class var ListClassVars: TList;
  public
    InstanceCount: integer;
    BaseClassName: string;
    constructor Create;

    class procedure SaveToDisk;
  end;

  PJump = ^TJump;
  TJump = packed record
    OpCode: Byte;
    Distance: Pointer;
  end;

  TObjectHack = class(TObject)
  private
    class procedure SetClassVars(AClassVars: TClassVars);
    class function GetClassVars: TClassVars;

    procedure IncCounter;
    procedure DecCounter;
    procedure OldFreeInstace;
  public
    class function InitInstance(Instance: Pointer): TObject;
  end;

var
  FOldFreeInstance: Pointer;

procedure SaveInstancesToFile;
begin
  TClassVars.SaveToDisk;
end;

function GetMethodAddress(AStub: Pointer): Pointer;
const
  CALL_OPCODE = $E8;
begin
  if PBYTE(AStub)^ = CALL_OPCODE then
  begin
    Inc(Integer(AStub));
    Result := Pointer(Integer(AStub) + SizeOf(Pointer) + PInteger(AStub)^);
  end
  else
    Result := nil;
end;

procedure AddressPatch(const ASource, ADestination: Pointer);
const
  JMP_OPCODE = $E9;
  SIZE = SizeOf(TJump);
var
  NewJump: PJump;
  OldProtect: Cardinal;
begin
  if VirtualProtect(ASource, SIZE, PAGE_EXECUTE_READWRITE, OldProtect) then
  begin
    NewJump := PJump(ASource);
    NewJump.OpCode := JMP_OPCODE;
    NewJump.Distance := Pointer(Integer(ADestination) - Integer(ASource) - 5);

    FlushInstructionCache(GetCurrentProcess, ASource, SizeOf(TJump));
    VirtualProtect(ASource, SIZE, OldProtect, @OldProtect);
  end;
end;

procedure PatchCodeDWORD(ACode: PDWORD; AValue: DWORD);
var
  LRestoreProtection, LIgnore: DWORD;
begin
  if VirtualProtect(ACode, SizeOf(ACode^), PAGE_EXECUTE_READWRITE, LRestoreProtection) then
  begin
    ACode^ := AValue;
    VirtualProtect(ACode, SizeOf(ACode^), LRestoreProtection, LIgnore);
    FlushInstructionCache(GetCurrentProcess, ACode, SizeOf(ACode^));
  end;
end;

procedure OldAfterConstruction;
asm
  call TObject.InitInstance;
end;

{ TCriticalSectionHack }
procedure TObjectHack.DecCounter;
begin
  if (Self.ClassType <> TClassVars) then
    Dec(GetClassVars.InstanceCount);
  OldFreeInstace;
end;

class function TObjectHack.GetClassVars: TClassVars;
begin
  Result := PClassVars(Integer(Self) + vmtAutoTable)^;
end;

class procedure TObjectHack.SetClassVars(AClassVars: TClassVars);
begin
  AClassVars.BaseClassName := Self.ClassName;
  PatchCodeDWORD(PDWORD(Integer(Self) + vmtAutoTable), DWORD(AClassVars));
end;

procedure RegisterClassVarsSupport(const Classes: array of TObjectHack);
var
  LClass: TObjectHack;
  LRestoreProtection: DWORD;
  LIgnore: DWORD;
  LVMT: Pointer;
begin
  for LClass in Classes do
    if LClass.GetClassVars = nil then
    begin
      LClass.SetClassVars(TClassVars.Create);

      //Change de mvt to object mvt
      LVMT := PPointer(Integer(TObject) + vmtFreeInstance)^;
      if VirtualProtect(LVMT, SizeOf(LVMT^), PAGE_EXECUTE_READWRITE, LRestoreProtection) then
      begin
        LVMT :=  @TObjectHack.DecCounter;
        VirtualProtect(LVMT, SizeOf(LVMT^), LRestoreProtection, LIgnore);
        FlushInstructionCache(GetCurrentProcess, LVMT, SizeOf(LVMT^));
      end;
    end
    else
      raise Exception.CreateFmt('Class %s has automated section or duplicated registration.', [LClass.ClassName]);
end;

procedure TObjectHack.IncCounter;
begin
  if (Self.ClassType = TClassVars) then
    Exit;

  if GetClassVars = nil then
    RegisterClassVarsSupport(Self);

  Inc(GetClassVars.InstanceCount);
end;

class function TObjectHack.InitInstance(Instance: Pointer): TObject;
asm
        PUSH    EBX
        PUSH    ESI
        PUSH    EDI
        MOV     EBX,EAX
        MOV     EDI,EDX
        STOSD
        MOV     ECX,[EBX].vmtInstanceSize
        XOR     EAX,EAX
        PUSH    ECX
        SHR     ECX,2
        DEC     ECX
        REP     STOSD
        POP     ECX
        AND     ECX,3
        REP     STOSB
        MOV     EAX,EDX
        MOV     EDX,ESP
@@0:    MOV     ECX,[EBX].vmtIntfTable
        TEST    ECX,ECX
        JE      @@1
        PUSH    ECX
@@1:    MOV     EBX,[EBX].vmtParent
        TEST    EBX,EBX
        JE      @@2
        MOV     EBX,[EBX]
        JMP     @@0
@@2:    CMP     ESP,EDX
        JE      @@5
@@3:    POP     EBX
        MOV     ECX,[EBX].TInterfaceTable.EntryCount
        ADD     EBX,4
@@4:    MOV     ESI,[EBX].TInterfaceEntry.VTable
        TEST    ESI,ESI
        JE      @@4a
        MOV     EDI,[EBX].TInterfaceEntry.IOffset
        MOV     [EAX+EDI],ESI
@@4a:   ADD     EBX,TYPE TInterfaceEntry
        DEC     ECX
        JNE     @@4
        CMP     ESP,EDX
        JNE     @@3
@@5:    MOV     EBX,EAX
        CALL    TObjectHack.IncCounter
        MOV     EAX,EBX
        POP     EDI
        POP     ESI
        POP     EBX
end;

procedure TObjectHack.OldFreeInstace;
asm
  call FOldFreeInstance;
end;

procedure InitFreeInstance;
begin
  FOldFreeInstance := PPointer(Integer(TObject) + vmtFreeInstance)^;
end;

{ TClassVars }

constructor TClassVars.Create;
begin
  ListClassVars.Add(Self);
end;

class procedure TClassVars.SaveToDisk;
var
  LStringList: TStringList;
  i: Integer;
begin                               
  LStringList := TStringList.Create;
  try
    LStringList.Add('CLASS | NUMBER OF INSTANCES');
    for i := 0 to ListClassVars.Count -1 do
      LStringList.Add(TClassVars(ListClassVars.Items[I]).BaseClassName + '|' + IntToStr(TClassVars(ListClassVars.Items[I]).InstanceCount));

    LStringList.SaveToFile(ExtractFilePath(ParamStr(0)) + 'test.txt');
  finally
    FreeAndNil(LStringList);
  end;
end;

initialization
  TClassVars.ListClassVars := TList.Create;
  InitFreeInstance;
  AddressPatch(GetMethodAddress(@OldAfterConstruction), @TObjectHack.InitInstance);

end.

Problem

Original source