
/**
 * @file vad-lib.c
 *
 * Tento modul implementuje nekolik rutin umoznujicich pohodlnejsi praci se 
 * strukturami virtualnich deskriptoru (VAD). Zatim umoznuje tyto struktury
 * pouze prochazet a cist nektere jejich polozky.
 */

#include <ntddk.h>
#include "version.h"
#include "driver.h"
#include "vad-def.h"
#include "vad-lib.h"


// Offset polozky VadRoot relativni k pocatku struktury EPROCESS
static ULONG VadRootOffset = 0;
// Offset polozky PhysicalVadRoot relativni k pocatku struktury EPROCESS
static ULONG PhysicalVadRootOffset = 0;
// Cislo verze operacniho systemu, na ktere ovladac prave bezi
static ULONG VersionCode = VERSION_UNKNOWN;


/** Zjisti adresu korenoveho VADu pro dany proces. Tento postup je mirne odlisny
 *  pro ruzne verze Windows.
 *
 *  @param Process Adresa struktury EPROCESS reprezentujici cilovy proces.
 *
 *  @return Funkce vraci adresu korenoveho VADu daneho procesu. V pripade neuspechu
 *  je vracena hodnota NULL.
 */
static PVOID GetProcessVadRootNode(IN PEPROCESS Process)
{
  PVOID Ret = NULL;
  PMM_AVL_TABLE avl = NULL;
  KdPrint(("vad-lib.c: GetProcessVadRootNode(Process=0x%p)\n", Process));

  avl = (PMM_AVL_TABLE)((PCHAR)Process + VadRootOffset);
  switch (VersionCode) {
    case VERSION_WXP:
      Ret = *(PVOID *)avl;
      break;
    case VERSION_W2K3:
    case VERSION_VISTA:
    case VERSION_W7:
      Ret = avl->RootNode.RightChild;
      break;
    default:
      KdPrint(("vad-lib.c: GetProcessVadRootNode(): ERROR: OS not supported\n"));
      ASSERT(FALSE);
      break;
  }

  KdPrint(("vad-lib.c: GetProcessVadRootNode(-):0x%p\n", Ret));
  return Ret;
}


/** Vrati adresu sekundarniho korene virtualnich deskriptoru pro dany proces. Tento
 *  koren je zrejme pouzivan pro nektere neobvykle deskriptory souvisejici s mapovanim
 *  fyzicke pameti ci velkych stranek.
 *
 *  @param Process Adresa struktury EPROCESS reprezentujici cilovy proces.
 *
 *  @return Vrati adresu sekundarniho korene deskriptoru virtualnich adres. V pripade
 *  neuspechu je vracena hodnota NULL.
 */
static PVOID GetProcessPhysicalVadRootNode(IN PEPROCESS Process)
{
  PVOID Ret = NULL;
  PMM_AVL_TABLE avl = NULL;
  KdPrint(("vad-lib.c: GetProcessPhysicalVadRootNode(Process=0x%p)\n", Process));

  avl = *((PMM_AVL_TABLE *)((PCHAR)Process + PhysicalVadRootOffset));
  switch (VersionCode) {
    case VERSION_W2K3:
    case VERSION_VISTA:
    case VERSION_W7:
      if (avl != NULL)
        Ret = avl->RootNode.RightChild;
      break;
    default:
      KdPrint(("vad-lib.c: GetProcessVadRootNode(): ERROR: OS not supported\n"));
      ASSERT(FALSE);
      break;
  }

  KdPrint(("vad-lib.c: GetProcessPhysicalVadRootNode(-):0x%p\n", Ret));
  return Ret;
}


/** Zajistuje rekurzivni prochazeni stromu virtualnich deskriptoru. Na kazdy nalezeny
 *  deskriptor zavola zadanou zpetne volanou funkci.
 *
 *  @param Vad Adresa virtualniho deskriptoru reprezentujiciho koren podstromu, ktery
 *  ma rutina projit.
 *  @param WalkRoutine Zpetne volana funkce, ktera bude vykonana pro kazdy nalezeny
 *  virtualni deskriptor, ktery patri do daneho podstromu (vyjma korene).
 *  @param Context Hodnota, ktera bude spolecne s adresou daneho deskriptoru predana
 *  do zpetne volane funkce.
 *
 *  @return Rutina vraci hodnotu NTSTATUS. Pokud navratova hodnota indikuje chybu
 *  (NT_SUCCESS se vyhodnoti na FALSE), prochazeni podstromu okamzite konci a chyba
 *  se propaguje vyse, takze zastavi prochazeni celeho stromu virtualnich deskriptoru.
 *  Pokud vracena hodnota znamena uspech (NT_SUCCESS se vyhodnoti na TRUE), v 
 *  podstromu se pokracuje.
 */
static NTSTATUS _VadTreeWalk(PVOID Vad, VAD_WALK_ROUTINE WalkRoutine, PVOID Context)
{
  MMADDRESS_NODE Node;
  MMADDRESS_NODE ChildNode;
  NTSTATUS Status = STATUS_UNSUCCESSFUL;
  KdPrint(("vad-lib.c: _VadTreeWalk(Vad=0x%p; WalkRoutine=0x%p; Context=0x%p)\n", Vad, WalkRoutine, Context));

  VDFillAddressNode(Vad, &Node);
  Status = STATUS_SUCCESS;
  if (Node.LeftChild != NULL)
    Status = _VadTreeWalk(Node.LeftChild, WalkRoutine, Context);

  if (NT_SUCCESS(Status)) {
    Status = WalkRoutine(Vad, Context);
    if (NT_SUCCESS(Status)) {
      if (Node.RightChild != NULL) 
        Status = _VadTreeWalk(Node.RightChild, WalkRoutine, Context);
	}
  }

  KdPrint(("vad-lib.c: _VadTreeWalk(-):0x%x\n", Status));
  return Status;
}


/** Projde hlavni a vedlejsi strom virtualnich deskriptoru zadaneho procesu.
 *
 *  Na kazdy nalezeny deskriptor je volana zadana zpetne volana funkce.
 *
 *  @param Process Adresa struktury EPROCESS reprezentujici cilovy proces.
 *  @param WalkRoutine Adresa zpetne volane funkce, kterou rutina vykona pro kazdy
 *  virtualni deskriptor.
 *  @param Context Hodnota, kterou bude rutina predavat zpetne volane funkci spolecne
 *  s adresou kazdeho virtualniho deskriptoru.
 *
 *  @return Funkce vraci hodnotu NTSTATUS indikujici uspech ci neuspech operace.
 *  Tato hodnota je primo ovlivnena navratovymi hodnotami z volani zpetne volane
 *  funkce predane v parametru WalkRoutine.
 */
NTSTATUS VadTreeWalk(PEPROCESS Process, VAD_WALK_ROUTINE WalkRoutine, PVOID Context)
{
  PVOID Vad = NULL;
  NTSTATUS Status = STATUS_UNSUCCESSFUL;
  KdPrint(("vad-lib.c: VadTreeWalk(Process=0x%p; WalkRoutine=0x%p; Context=0x%p)\n", Process, WalkRoutine, Context));

  Status = STATUS_NOT_FOUND;
  Vad = GetProcessVadRootNode(Process);
  if (Vad != NULL)
    Status = _VadTreeWalk(Vad, WalkRoutine, Context);

  if (VersionCode != VERSION_WXP && NT_SUCCESS(Status)) {
    Vad = GetProcessPhysicalVadRootNode(Process);
    if (Vad != NULL)
      Status = _VadTreeWalk(Vad, WalkRoutine, Context);
  }

  KdPrint(("vad-lib.c: VadTreeWalk(-):0x%x\n", Status));
  return Status;
}


/** Zpetne volana funkce pro knihovni rutinu, ktera prochazi stromy deskriptoru
 *  a zapisuje si o kazde polozce informace.
 *
 *  Rutina precte nektere udaje z virtualniho deskriptoru zadaneho v prvnim parametru
 *  a ulozi je do seznamu zadaneho v parametru druhem.
 *
 *  @param Vad Adresa VADu.
 *  @param Context Adresa hlavy spojoveho seznamu, kam rutina zapisuje informace
 *  o jednotlivych nalezenych virtualnich deskriptorech.
 *
 *  @return Funkce vraci hodnotu NTSTATUS indikujici uspech ci neupsech operace.
 */
static NTSTATUS _GetVadInfo(PVOID Vad, PVOID Context)
{
  MMADDRESS_NODE Node;
  MI_VAD_TYPE VadType = VadNone;
  VAD_STRUCTURE_TYPE VadStructureType = VadStructureUnknown;
  MMVAD_FLAGS VadFlags;
  MMVAD_FLAGS2 VadFlags2;
  MMVAD_FLAGS3 VadFlags3;
  PVAD_RECORD VadRecord = NULL;
  NTSTATUS Status = STATUS_UNSUCCESSFUL;
  PLIST_ENTRY VadInfoList = (PLIST_ENTRY)Context;
  KdPrint(("vad-lib.c: GetVadInfo(Vad=0x%p; Context=0x%p)\n", Vad, Context));
  
  VadRecord = (PVAD_RECORD)ExAllocatePoolWithTag(PagedPool, sizeof(VAD_RECORD), POOL_TAG);
  if (VadRecord != NULL) {
    RtlZeroMemory(VadRecord, sizeof(VAD_RECORD));
    InitializeListHead(&VadRecord->Entry);
    VDFillAddressNode(Vad, &Node);
    VadType = VDGetVadType(Vad);
    VadStructureType = VDGetVadStructureType(Vad);
    VadRecord->Record.StartAddress = Node.StartingVPN << 12;
    VadRecord->Record.EndAddress = (Node.EdningVPN << 12) | 0xFFF;
    VadRecord->Record.VadType = VadType;
    VadRecord->Record.StructureType = VadStructureType;
    switch (VadStructureType) {
      case VadStructureShort:
        VadFlags = VDGetVadFlags(Vad);
        VadRecord->Record.u1.Flags1 = VadFlags;
        if (VersionCode == VERSION_VISTA || VersionCode == VERSION_W7) {
          VadFlags3 = VDGetVadFlags3(Vad);
          VadRecord->Record.u3.Flags3 = VadFlags3;
        }
        break;
      case VadStructureNormal:
        VadFlags = VDGetVadFlags(Vad);
        VadRecord->Record.u1.Flags1 = VadFlags;
        VadFlags2 = VDGetVadFlags2(Vad);
        VadRecord->Record.u2.Flags2 = VadFlags2;
        if (VersionCode == VERSION_VISTA || VersionCode == VERSION_W7) {
          VadFlags3 = VDGetVadFlags3(Vad);
          VadRecord->Record.u3.Flags3 = VadFlags3;
        }
        break;
      case VadStructureLong:
        VadFlags = VDGetVadFlags(Vad);
        VadRecord->Record.u1.Flags1 = VadFlags;
        VadFlags2 = VDGetVadFlags2(Vad);
        VadRecord->Record.u2.Flags2 = VadFlags2;
        VadFlags3 = VDGetVadFlags3(Vad);
        VadRecord->Record.u3.Flags3 = VadFlags3;
        break;
      default:
        KdPrint(("vad-lib.c: _GetVadInfo(): ERROR: Unknown type of VAD\n"));
        ASSERT(FALSE);
        break;
    }

    InsertTailList(VadInfoList, &VadRecord->Entry);
    Status = STATUS_SUCCESS;
  } else Status = STATUS_UNSUCCESSFUL;

  KdPrint(("vad-lib.c: GetVadInfo(-):0x%x\n", Status));
  return Status;
}


/** Zjisti udaje o vsech virtualnich deskriptorech pro dany proces.
 *
 *  @param Process Adresa struktury EPROCESS reprezentujici cilovy proces.
 *  @param ListHead Hlava spojoveho seznamu, kam rutina v pripade uspechu
 *  ulozi informace o jednotlivych virtualnich deskriptorech ve forme zaznamu
 *  typu VAD_RECORD.
 *
 *  @return Funkce vraci hodnotu NTSTATUS indikujici uspech ci neuspech operace.
 */
NTSTATUS ListProcessVads(PEPROCESS Process, PLIST_ENTRY ListHead)
{
  NTSTATUS Status = STATUS_UNSUCCESSFUL;
  KdPrint(("vad-lib.c: ListProcessVads(Process=0x%p; ListHead=0x%p)\n", Process, ListHead));

  InitializeListHead(ListHead);
  Status = VadTreeWalk(Process, _GetVadInfo, ListHead);
  if (!NT_SUCCESS(Status)) {
    PLIST_ENTRY Tmp = ListHead->Flink;
    while (Tmp != ListHead) {
      PLIST_ENTRY Item = Tmp;
	  Tmp = Tmp->Flink;
      ExFreePoolWithTag(Item, POOL_TAG);
	}

	InitializeListHead(ListHead);
  }

  KdPrint(("vad-list.c: ListProcessVads(-):0x%x\n", Status));
  return Status;
}



/** Inicializuje knihovnu pro praci s virtualnimi deskriptory procesu.
 *
 *  Rutina nejprve inicializuje modul vad-def.c.
 *  Pote na zaklade verze operacniho systemu vyplni promenne obsahujici offset
 *  polozek VadRoot a PhysicalVadRoot ve strukturach EPROCESS.
 *
 *  @return Funkce vraci hodnotu NTSTATUS indikujici uspech ci neupsech operace.
 */
NTSTATUS VadLibInit(VOID)
{
  NTSTATUS Status = STATUS_UNSUCCESSFUL;
  KdPrint(("vad-lib.c: VadLibInit()\n"));

  Status = VDInit();
  if (NT_SUCCESS(Status)) {
    VersionCode = GetVersionCode();
    if (sizeof(PVOID) == 4) {
      switch (VersionCode) {
        case VERSION_WXP:
          VadRootOffset = VAD_ROOT_OFFSET_WXP;
          break;
	    case VERSION_W2K3:
          VadRootOffset = VAD_ROOT_OFFSET_W2K3;
          PhysicalVadRootOffset = PHYSICAL_VAD_ROOT_OFFSET_W2K3;
          break;
        case VERSION_VISTA:
          VadRootOffset = VAD_ROOT_OFFSET_VISTA;
          PhysicalVadRootOffset = PHYSICAL_VAD_ROOT_OFFSET_VISTA;
          break;
        case VERSION_W7:
          VadRootOffset = VAD_ROOT_OFFSET_W7;
          PhysicalVadRootOffset = PHYSICAL_VAD_ROOT_OFFSET_W7;
          break;
	    default:
          Status = STATUS_NOT_SUPPORTED;
          break;
      }
    } else if (sizeof(PVOID) == 8) {
      switch (VersionCode) {
        case VERSION_W2K3:
          VadRootOffset = VAD_ROOT_OFFSET_W2K3_X64;
          PhysicalVadRootOffset = PHYSICAL_VAD_ROOT_OFFSET_W2K3_X64;
          break;
        case VERSION_VISTA:
          VadRootOffset = VAD_ROOT_OFFSET_VISTA_X64;
          PhysicalVadRootOffset = PHYSICAL_VAD_ROOT_OFFSET_VISTA_X64;
          break;
        case VERSION_W7:
          VadRootOffset = VAD_ROOT_OFFSET_W7_X64;
          PhysicalVadRootOffset = PHYSICAL_VAD_ROOT_OFFSET_W7_X64;
          break;
        default:
          Status = STATUS_NOT_SUPPORTED;
          break;
      }
    } else Status = STATUS_NOT_SUPPORTED;
  }

  KdPrint(("vad-lib.c: VadLibInit(-):0x%x\n", Status));
  return Status;
}


/** Provede uklid pro knihovne pro praci s virtualnimi deskriptory.
 *
 *  Protoze inicializacni rutina neprovadela zadne zasahy do datovych struktur
 *  jadra a ani s nim jinak neinteragovala, tato funkce pouze provede uklid
 *  modulu vad-def.c.
 */
VOID VadLibFinit(VOID)
{
  KdPrint(("vad-lib.c: VadLibFinit()\n"));

  VDFinit();

  KdPrint(("vad-lib.c: VadLibFinit(-)\n"));
  return;
}
