
/**
 * @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-secure.h"
#include "vad-lib.h"


// Offset polozky VadRoot relativni k pocatku struktury EPROCESS
static ULONG VadRootOffset = 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;
}



/** 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);

  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;
  SIZE_T VadRecordLength = 0;
  KdPrint(("vad-lib.c: GetVadInfo(Vad=0x%p; Context=0x%p)\n", Vad, Context));
  
  VadRecordLength = sizeof(VAD_RECORD);
  if (VSIsVadSecured(Vad))
    VadRecordLength += VSGetVadSecuritySize(Vad);

  VadRecord = (PVAD_RECORD)ExAllocatePoolWithTag(PagedPool, VadRecordLength, POOL_TAG);
  if (VadRecord != NULL) {
    RtlZeroMemory(VadRecord, VadRecordLength);
    InitializeListHead(&VadRecord->Entry);
    VDFillAddressNode(Vad, &Node);
    VadType = VDGetVadType(Vad);
    VadStructureType = VDGetVadStructureType(Vad);
    VadRecord->Record.Length = VadRecordLength - (sizeof(VAD_RECORD) - sizeof(UM_VAD_RECORD));
    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;
        if (VSIsVadSecured(Vad))
          VSGetVadSecurity(Vad, &VadRecord->Record.SecureInfo);
        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)) {
    Status = VSModuleInit();
    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; break;
          case VERSION_VISTA: VadRootOffset = VAD_ROOT_OFFSET_VISTA; break;
          case VERSION_W7: VadRootOffset = 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; break;
          case VERSION_VISTA: VadRootOffset = VAD_ROOT_OFFSET_VISTA_X64; break;
          case VERSION_W7: VadRootOffset = VAD_ROOT_OFFSET_W7_X64; break;
          default:
            Status = STATUS_NOT_SUPPORTED;
            break;
        }
      } else {
        Status = STATUS_NOT_SUPPORTED;
        VSModuleFinit();
        VDFinit();
      }
    } else VDFinit();
  }

  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"));

  VSModuleFinit();
  VDFinit();

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