
/**
 * @file: intcount.c
 *
 * Tento soubor implementuje modul pro pocitani vyskytu jednotlivych
 * preruseni. Metoda pocitani je zalozena na modifikaci adres obsluznych
 * rutin v deskriptorech vsech IDT.
 */

#include <ntddk.h>
#include "sync.h"
#include "idt.h"
#include "ioctls.h"
#include "driver.h"
#include "intcount.h"


// Pocet aktivnich logickych procesoru
static LONG CPUCount = 0;
// Puvodni adresy obsluznych rutin preruseni
static PVOID *IsrAddresses;
// Citace jednotlivych preruseni
static ULONG Counters [0x100];
// Instance novych obsluznych rutin pro preruseni
static PVOID IsrHandlers [0x100 * MAX_PROCESSORS];



/** Jedna se o metarutinu. Pri inicializaci bude rozkoporovana na mnoho
 *  instanci (jejich pocet odpovida 0x100 * pocet procesoru, coz je maxi-
 *  malni pocet ruznych preruseni) a offsety v jejim tele budou nahrazeny
 *  nasledovne:
 *  + 0xBBBBBBBB za adresu citace prislusneho preruseni
 *  + 0xAAAAAAAA za relativni adresu puvodni obsluzne rutiny vzhledem
 *    k prislusne instanci nove obsluzne rutiny.
 *
 *  Metarutina pouze zvysi citac odpovidajici danemu preruseni o jednicku
 *  a zavola originalni obsluznou rutinu.
 */
static __declspec(naked) VOID DefaultIsrHandler()
{
  __asm
  {
    lock inc dword ptr ds:[0xBBBBBBBB]
    __emit 0xe9
    __emit 0xAA
    __emit 0xAA
    __emit 0xAA
    __emit 0xAA
  }
}



/** Obnovi puvodni stav tabulky vektoru preruseni aktualniho procesoru.
 *
 * Rutina zapise do deskriptoru IDT aktualniho procesoru adresy obsluznych
 * rutin, ktere tam byly pred zapnutim pocitani preruseni. Hlavicka rutiny
 * je kompatibilni s odlozene volanou procedurou.
 *
 * @param Dpc Tetno parametr je pritomen pouze z duvodu kompatibility s
 * odlozene volanou procedurou.
 * @param IDTsRestored Adresa promenne, ktera uchovava pocet tabulek vektoru
 * preruseni, ktere se jiz podarilo obnovit do puvodniho stavu.
 * @param Arg1 Tetno parametr je pritomen pouze z duvodu kompatibility s
 * odlozene volanou procedurou.
 * @param Arg2 Tetno parametr je pritomen pouze z duvodu kompatibility s
 * odlozene volanou procedurou.
 */
static VOID _RestoreIDT(IN PKDPC Dpc, IN PVOID IDTsRestored, IN PVOID Arg1, IN PVOID Arg2)
{
  ULONG i;
  ULONG CurrentCPU = KeGetCurrentProcessorNumber();
  KdPrint(("intcount.c: _RestoreIDT(Dpc=0x%p; IDTsRestored=0x%p; Arg1=0x%p; Arg2=0x%p)\n", Dpc, IDTsRestored, Arg1, Arg2));

  for (i = 0; i < 256; i++)
    IdtModifyIsr(i, IsrAddresses[CurrentCPU * 0x100 + i]);
  
  InterlockedIncrement((PLONG)IDTsRestored);

  KdPrint(("intcount.c: _RestoreIDT(-)\n"));
  return;
}


/** Obnovi vsechny tabulky vektoru preruseni do puvodniho stavu.
 */
static VOID RestoreAllIDTs(VOID)
{
  LONG IDTsRestored = 0;
  PKDPC Dpc = NULL;
  NTSTATUS Status = STATUS_UNSUCCESSFUL;
  KdPrint(("intcount.c: RestoreAllIDTs()\n"));

  Status = ExecuteAtEveryProcessor(_RestoreIDT, &IDTsRestored, NULL, NULL, &Dpc);
  if (NT_SUCCESS(Status)) {
    // Cekani, dokud nejsou obnoveny vsechny IDT
    while (InterlockedCompareExchange(&IDTsRestored, IDTsRestored, IDTsRestored) < CPUCount)
      __asm nop
    
    ExFreePoolWithTag(Dpc, POOL_TAG);
  }

  KdPrint(("intcount.c: RestoreAllIDTs(-)\n"));
  return;
}


/** Zmeni obslune rutiny vsech preruseni v tabulce vektoru preruseni
 *  aktualniho procesoru.
 *
 * @param Dpc Tetno parametr se nevyzuva a je pritomen pouze v dusledku
 * dodrzeni kompatiblity s odlozene volanou procedurou.
 * @param IDTsHooked Adresa promenne, ktera udava, kolik tabulek vektoru
 * preruseni jiz bylo zmeneno.
 * @param OffsetAAAAAAAA Offset cisla 0xAAAAAAAA v ramci obsluze rutiny preruseni.
 * Toto cislo musi byt pred nasmerovanim obsluzne rutiny upraveno.
 * @param OffsetBBBBBBBB Offset cisla 0xBBBBBBBB v ramci obsluze rutiny preruseni.
 * Toto cislo musi byt pred nasmerovanim obsluzne rutiny upraveno.
 */
static VOID _HookIDT(IN PKDPC Dpc, IN PVOID IDTsHooked, IN PVOID OffsetAAAAAAAA, IN PVOID OffsetBBBBBBBB)
{
  ULONG i;
  PULONG Mem = NULL;
  ULONG CurrentCPU = KeGetCurrentProcessorNumber();
  ULONG OffsetA = (ULONG)OffsetAAAAAAAA;
  ULONG OffsetB = (ULONG)OffsetBBBBBBBB;
  KdPrint(("intcount.c: _HookIDT(Dpc=0x%p; IDTsHooked=0x%p; OffsetAAAAAAAA=%u; OffsetBBBBBBBB=%u)\n", Dpc, IDTsHooked, OffsetAAAAAAAA, OffsetBBBBBBBB));

  for (i = 0; i < 256; ++i) {
    if (i == 2 || i == 8)
      continue;

    IsrAddresses[CurrentCPU * 0x100 + i] = IdtGetIsrAddress(i);
    RtlCopyMemory(IsrHandlers[CurrentCPU * 0x100 + i], DefaultIsrHandler, HANDLER_SIZE);
    Mem = (PULONG)((PCHAR)(IsrHandlers[CurrentCPU * 0x100 + i]) + OffsetA);
    *Mem = (ULONG)IsrAddresses[CurrentCPU * 0x100 + i] - (ULONG)Mem + 1 - 5;
    Mem = (PULONG)((PCHAR)(IsrHandlers[CurrentCPU * 0x100 + i]) + OffsetB);
    *Mem = (ULONG)&Counters[i];
    IdtModifyIsr(i, IsrHandlers[i + CurrentCPU * 0x100]);
  }

  InterlockedIncrement((PLONG)IDTsHooked);

  KdPrint(("intcount.c: _HookIDT(-)\n"));
  return;
}



/** Najde offsety cisel 0xAAAAAAAA a 0xBBBBBBBB v ramci obsluzne rutiny
 *  pro vsechna preruseni.
 *
 * @param OffsetA Adresa promenne, kam funkce v pripade uspechu zapise
 * offset cisla 0xAAAAAAAA relativni k pocatku obsluzne rutiny.
 * @param OffsetB Adresa promenne, kam funkce v pripade uspechu zapise
 * offset cisla 0xBBBBBBBB relativni k pocatku obsluzne rutiny.
 *
 * @return Funkce vraci hodnotu NTSTATUS indikujici uspech ci neuspech
 * operace.
 */
static NTSTATUS FindOffsets(OUT PULONG OffsetA, OUT PULONG OffsetB)
{
  PULONG Mem = NULL;
  LONG i = 0;
  NTSTATUS Status = STATUS_UNSUCCESSFUL;
  KdPrint(("intcount.c: FindOffsets(OffsetA=0x%p; OffsetB=0x%p)\n", OffsetA, OffsetB));

  for (i = 0; i < HANDLER_SIZE - sizeof(ULONG) + 1; i++) {
    Mem = (PULONG)((PCHAR)&DefaultIsrHandler + i);
    if (*Mem == 0xAAAAAAAA)
      *OffsetA = i;
    if (*Mem == 0xBBBBBBBB)
      *OffsetB = i;
    if (*OffsetA != 0 && *OffsetB != 0) {
      Status = STATUS_SUCCESS;
      break;
    }
  }

  KdPrint(("intcount.c: FindOffsets(-):0x%x, *OffsetA=%u, *OffsetB=%u\n", Status, OffsetA, OffsetB));
  return Status;
}


/** Modifikuje tabulky vektoru preruseni na vsech procesorech.
 *
 * Rutina nejprve zjisti offsety cisel 0xAAAAAAAA a 0xBBBBBBBB v ramci
 * nove obsluzne rutiny. Tuto rutinu pak rozkopiruje na tolik instanci,
 * aby kazde preruseni mohlo ziskat prave jednu instanci. Nasledne pomoci
 * odlozeneho volani procedur donuti kazdy procesor modifikovat svoji
 * IDT.
 *
 * @return Funkce vraci hodnotu NTSTATUS indikujici uspech ci neuspech
 * operace.
 */
static NTSTATUS HookAllIDTs(VOID)
{
  LONG i = 0;
  LONG j = 0;
  LONG IDTsHooked = 0;
  PULONG Mem = NULL;
  PKDPC Dpc = NULL;
  ULONG OffsetAAAAAAAA = 0;
  ULONG OffsetBBBBBBBB = 0;
  NTSTATUS Status = STATUS_UNSUCCESSFUL;
  KdPrint(("intcount.c: HookAllIDTs()\n"));

  Status = FindOffsets(&OffsetAAAAAAAA, &OffsetBBBBBBBB);
  if (NT_SUCCESS(Status)) {
    IsrAddresses = ExAllocatePoolWithTag(NonPagedPool, 0x100 * sizeof(PVOID) * CPUCount, POOL_TAG);
    if (IsrAddresses != NULL) {
      for (i = 0; i < 0x100 * CPUCount; ++i) {
        IsrHandlers[i] = ExAllocatePoolWithTag(NonPagedPool, HANDLER_SIZE, POOL_TAG);
        if (IsrHandlers[i] == NULL) {
          Status = STATUS_INSUFFICIENT_RESOURCES;
          for (j = 0; j < i; ++j)
            ExFreePoolWithTag(IsrHandlers[j], POOL_TAG);

          break;
        }
      }

      if (NT_SUCCESS(Status)) {
        Status = ExecuteAtEveryProcessor(_HookIDT, &IDTsHooked, (PVOID)OffsetAAAAAAAA, (PVOID)OffsetBBBBBBBB, &Dpc);
        if (NT_SUCCESS(Status)) {
          while (InterlockedCompareExchange(&IDTsHooked, IDTsHooked, IDTsHooked) < CPUCount)
            __asm nop;
      
          ExFreePoolWithTag(Dpc, POOL_TAG);
        } else {
          for (i = 0; i < 0x100 * CPUCount; i++)
            ExFreePoolWithTag(IsrHandlers[i], POOL_TAG);
        }
      }

      if (!NT_SUCCESS(Status))
        ExFreePoolWithTag(IsrAddresses, POOL_TAG);
    } else Status = STATUS_INSUFFICIENT_RESOURCES;
  }

  KdPrint(("intcount.c: HookAllIDTs(-):0x%x\n", Status));
  return Status;
}


/** Vrati adresu pole citacu vskytu jednotlivych preruseni.
 *
 * @return Funkce vrati adresu pole citacu vyskytu jednotlivych preruseni.
 */
PULONG IntCountGetCounterArray(VOID)
{
  return (PULONG)&Counters;
}


/** Vrati velikost pole citacu v bajtech.
 *
 * @return Funkce vraci velikost pole citacu v bajtech.
 */
ULONG IntCountGetCounterArrayLength(VOID)
{
  return sizeof(Counters);
}


/** Vynuluje pole citacu vyskytu jednotlivych preruseni.
 */
VOID IntCountResetCounterArray(VOID)
{
  RtlZeroMemory(&Counters, sizeof(Counters));

  return;
}



/** Inicializuje mechanismus pro pocitani vyskytu preruseni.
 *
 * Rutina vynuluje pole citacu, zjisti pocet aktivnich logickych procesoru
 * a modifikuje vsechny tabulky vektoru preruseni. Tim zapne cely mechanismus.
 *
 * @return Funkce vraci hodnotu NTSTATUS indikujici uspech ci neuspech
 * operace.
 */
NTSTATUS IntCountInit(VOID)
{
  NTSTATUS Status = STATUS_UNSUCCESSFUL;
  KdPrint(("intcount.c: IntCountInit()\n"));

  IntCountResetCounterArray();
  CPUCount = KeNumberProcessors;
  Status = HookAllIDTs();

  KdPrint(("intcount.c: IntCountInit(-):0x%x\n", Status));
  return Status;
}


/** Vypne pocitani vyskytu preruseni a uvolni pamet, kterou tento modul
 *  pro tyto ucely alokoval.
 */
VOID IntCountFinit(VOID)
{
  LONG i = 0;
  KdPrint(("intcount.c: IntCountFinit()\n"));

  RestoreAllIDTs();
  ExFreePoolWithTag(IsrAddresses, POOL_TAG);
  for (i = 0; i < CPUCount * 0x100; ++i)
    ExFreePoolWithTag(IsrHandlers[i], POOL_TAG);

  KdPrint(("intcount.c: IntCountFinit(-)\n"));
  return;
}
