
/**
 * @file driver.c
 *
 * Hlavni soubor ovladace vad.sys. Zajistuje inicializaci vsech jeho soucasti
 * a uklid po nich. Dale obstarava komunikaci ovladace s aplikacemi ci pripadne
 * ostatnimi ovladaci.
 */

#include <ntifs.h>
#include "ioctls.h"
#include "version.h"
#include "vad-lib.h"
#include "driver.h"



/** Obsluha pozadavku typu IRP_MJ_CREATE a IRP_MJ_CLOSE. Jedna se pouze o stub,
 *  rutina vzdy vraci uspech.
 */
NTSTATUS DriverCreateClose(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
  NTSTATUS Status = STATUS_UNSUCCESSFUL;
  KdPrint(("driver.c: DriverCreateClose(DeviceObject=0x%p; Irp=0x%p)\n", DeviceObject, Irp));

  Status = STATUS_SUCCESS;
  Irp->IoStatus.Status = Status;
  IoCompleteRequest(Irp, IO_NO_INCREMENT);

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


/** Obsluha pozadavku IRP_MJ_DEVICE_CONTROL.
 *
 *  Ovladac odpovida pouze na pozadavek IOCTL_VAD_GET. V odpovedi vraci informace
 *  o virtualnich deskriptorech procesu. Struktura vystupniho bufferu je nasledujici:
 *    0x000      ULONG64       [Pocet virtualnich deskriptoru]
 *    0x008      UM_VAD_RECORD [Informace o prvnim deskriptoru]
 *    0xYYY      UM_VAD_RECORD [Informace o druhem deskriptoru]
 *    .....      
 *    0xXXX      UM_VAD_RECORD [Informace o poslednim deskriptoru]
 *
 *  Struktura vstupniho bufferu:
 *    0x000     HANDLE       [PID ciloveho procesu]
 *
 */
NTSTATUS DriverDeviceControl(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
  PIO_STACK_LOCATION IrpStack = NULL;
  PVOID OutBuffer = NULL;
  ULONG OutBufferLength = 0;
  PVOID InBuffer = NULL;
  ULONG InBufferLength = 0;
  ULONG ControlCode = 0;
  NTSTATUS Status = STATUS_UNSUCCESSFUL;
  KdPrint(("driver.c: DriverDeviceControl(DeviceObject=0x%p; Irp=0x%p)\n", DeviceObject, Irp));

  IrpStack = IoGetCurrentIrpStackLocation(Irp);
  OutBufferLength = IrpStack->Parameters.DeviceIoControl.OutputBufferLength;
  OutBuffer = Irp->AssociatedIrp.SystemBuffer;
  InBufferLength = IrpStack->Parameters.DeviceIoControl.InputBufferLength;
  InBuffer = Irp->AssociatedIrp.SystemBuffer;
  ControlCode = IrpStack->Parameters.DeviceIoControl.IoControlCode;
  Irp->IoStatus.Information = 0;

  switch (ControlCode)
  {
    case IOCTL_VAD_GET:
      if (InBufferLength >= sizeof(HANDLE) && OutBufferLength >= sizeof(ULONG64) + sizeof(UM_VAD_RECORD)) {
        PEPROCESS TargetProcess = NULL;
        PULONG64 VadCount = (PULONG64)OutBuffer;

        Status = PsLookupProcessByProcessId(*(PHANDLE)InBuffer, &TargetProcess);
        OutBuffer = (PCHAR)OutBuffer + sizeof(ULONG64);
        *VadCount = 0;
        if (NT_SUCCESS(Status)) {
          LIST_ENTRY VadListHead;

          Status = ListProcessVads(TargetProcess, &VadListHead);
          if (NT_SUCCESS(Status)) {
            ULONG RequiredLength = sizeof(ULONG64);
            PVAD_RECORD VadRecord = (PVAD_RECORD)VadListHead.Flink;

            while (&VadRecord->Entry != &VadListHead) {
              RequiredLength += (ULONG)VadRecord->Record.Length;
              if (RequiredLength <= OutBufferLength) {
                KdPrint(("VAD: 0x%p\n", VadRecord));
                RtlCopyMemory(OutBuffer, &VadRecord->Record, (SIZE_T)VadRecord->Record.Length);
                OutBuffer = (PCHAR)OutBuffer + VadRecord->Record.Length;
                Irp->IoStatus.Information = RequiredLength;
                (*VadCount)++;
              } else {
                Status = STATUS_BUFFER_TOO_SMALL;
                Irp->IoStatus.Information = 0;
                *VadCount = 0;
                break;
              }

              VadRecord = (PVAD_RECORD)VadRecord->Entry.Flink;
            }

			{
              PLIST_ENTRY Tmp = VadListHead.Blink;
              PLIST_ENTRY Item = NULL;

              while (Tmp != &VadListHead) {
                Item = Tmp;
                Tmp = Tmp->Blink;
	            ExFreePoolWithTag(Item, POOL_TAG);
              }
            }
          }

          ObDereferenceObject(TargetProcess);
        }
      } else Status = STATUS_BUFFER_TOO_SMALL;
      break;

    default:
      KdPrint(("driver.c: DriverDeviceControl(): Invalid device control request\n"));
      Status = STATUS_NOT_IMPLEMENTED;
      break;
  }  

  Irp->IoStatus.Status = Status;
  IoCompleteRequest (Irp, IO_NO_INCREMENT);

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



/** Odstrani zarizeni pro komunikaci s okolim a symbolicky odkaz na nej.
 *
 *  @param Adresa struktury DRIVER_OBJECT reprezentujici ovladac.
 */
VOID DriverFinit(IN PDRIVER_OBJECT DriverObject)
{
  UNICODE_STRING uSymLink;
  KdPrint(("driver.c: DriverFinit(Driverobject=0x%p)\n", DriverObject));

  RtlInitUnicodeString(&uSymLink, DRIVER_SYMLINK);
  IoDeleteSymbolicLink(&uSymLink);
  IoDeleteDevice(DriverObject->DeviceObject);

  KdPrint(("driver.c: DriverFinit(-)\n"));
  return;
}


/** Jadro vola rutina tesne pred ostranenim ovladace ze pameti.
 *  Rutina zajisti uklid po vsech soucastech ovladace.
 */
VOID DriverUnload(IN PDRIVER_OBJECT DriverObject)
{
  KdPrint(("driver.c: DriverUnload(DriverObject=0x%p)\n", DriverObject));

  DriverFinit(DriverObject);
  VadLibFinit();
  VersionFinit();

  KdPrint(("driver.c: DriverUnload(-)\n"));
  return;
}


/** Inicializuje primitiva potrebna pro komunikaci s okolim (vytvori objekt zarizeni
 *  a symbolicky odkaz). Rutina dale nastavi obsluzne funkce pro podporovane typy
 *  pozadavku (IRP_MJ_CREATE, IRP_MJ_CLOSE, IRP_MJ_DEVICE_CONTROL) a zajisti, ze
 *  ovladac lze dynamicky z pameti jadra uvolnit.
 *
 *  @param DriverObject Adresa struktury DRIVER_OBJECT reprezentujici ovladac vad.sys.
 *
 *  @return Funkce vraci hodnotu NTSTATUS indikujici uspech ci neuspech operace.
 */
NTSTATUS DriverInit(IN PDRIVER_OBJECT DriverObject)
{
  UNICODE_STRING uDevice;
  UNICODE_STRING uSymLink;
  NTSTATUS Status = STATUS_UNSUCCESSFUL;
  KdPrint(("driver.c: DriverInit(DriverObject=0x%p)\n", DriverObject));

  RtlInitUnicodeString(&uDevice, DRIVER_DEVICE);
  Status = IoCreateDevice(DriverObject, 0, &uDevice, FILE_DEVICE_UNKNOWN, 0, FALSE, &DriverObject->DeviceObject);
  if (NT_SUCCESS(Status)) {
    RtlInitUnicodeString(&uSymLink, DRIVER_SYMLINK);
    Status = IoCreateSymbolicLink(&uSymLink, &uDevice);
    if (NT_SUCCESS(Status)) {
      DriverObject->DriverUnload = DriverUnload;
      DriverObject->MajorFunction[IRP_MJ_CREATE] = DriverCreateClose;
      DriverObject->MajorFunction[IRP_MJ_CLOSE] = DriverCreateClose;
      DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = DriverDeviceControl;
      if (!NT_SUCCESS(Status)) {
        IoDeleteSymbolicLink(&uSymLink);
        IoDeleteDevice(DriverObject->DeviceObject);
      }
    } else IoDeleteDevice(DriverObject->DeviceObject);
  }

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


/** Jadro vola tuto rutinu tesne po namapovani ovladace do pameti. Rutina provede
 *  inicializaci vsech soucasti ovladace.
 */
NTSTATUS DriverEntry(IN PDRIVER_OBJECT DriverObject, IN PUNICODE_STRING RegistryPath)
{
  NTSTATUS Status = STATUS_UNSUCCESSFUL;
  KdPrint(("driver.c: DriverEntry(DriverObject=0x%p; RegistryPath=\"%S\")\n", DriverObject, RegistryPath->Buffer));

  Status = VersionInit();
  if (NT_SUCCESS(Status)) {
    Status = VadLibInit();
    if (NT_SUCCESS(Status)) {
      Status = DriverInit(DriverObject);
      if (!NT_SUCCESS(Status)) {
        VadLibFinit();
        VersionFinit();
      }
    } else VersionFinit();
  }

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