/**
 * \file
 * \author Egor Tensin <Egor.Tensin@gmail.com>
 * \date 2015
 * \copyright This file is licensed under the terms of the MIT License.
 *            See LICENSE.txt for details.
 */

#include "device.h"

#include <ntddk.h>

static NTSTATUS device_open(DEVICE_OBJECT *device_object, IRP *irp)
{
    NTSTATUS status = STATUS_SUCCESS;

    irp->IoStatus.Status = status;
    irp->IoStatus.Information = 0;
    IoCompleteRequest(irp, IO_NO_INCREMENT);
    return status;
}

typedef NTSTATUS (*ioctl_handler)(void *, unsigned long,
                                  void *, unsigned long,
                                  ULONG_PTR *);

static NTSTATUS handle_say_hello(void *in_buf,
                                 unsigned long in_buf_size,
                                 void *out_buf,
                                 unsigned long out_buf_size,
                                 ULONG_PTR *nbwritten)
{
    DbgPrint("Hello, world!\n");
    return STATUS_SUCCESS;
}

static NTSTATUS handle_exchange_ints(void *in_buf,
                                     unsigned long in_buf_size,
                                     void *out_buf,
                                     unsigned long out_buf_size,
                                     ULONG_PTR *nbwritten)
{
    unsigned int read;
    unsigned int written = 0xdeadbeef;

    if (in_buf_size != sizeof(read))
        return STATUS_INVALID_BUFFER_SIZE;

    RtlCopyMemory(&read, in_buf, in_buf_size);
    DbgPrint("%08x\n", read);

    if (out_buf_size < sizeof(written))
        return STATUS_BUFFER_TOO_SMALL;

    RtlCopyMemory(out_buf, &written, sizeof(written));
    *nbwritten += sizeof(written);

    return STATUS_SUCCESS;
}

#define SAY_HELLO CTL_CODE(0x8000, 0x800, METHOD_BUFFERED, FILE_ANY_ACCESS)
#define EXCHANGE_INTS CTL_CODE(0x8001, 0x800, METHOD_BUFFERED, FILE_ANY_ACCESS)

static NTSTATUS device_ioctl(DEVICE_OBJECT *device_object, IRP *irp)
{
    IO_STACK_LOCATION *io_stack_loc;
    void* in_buf, *out_buf;
    unsigned long in_buf_size, out_buf_size;
    ioctl_handler handler;
    NTSTATUS status = STATUS_UNSUCCESSFUL;

    irp->IoStatus.Status = status;
    irp->IoStatus.Information = 0;
    io_stack_loc = IoGetCurrentIrpStackLocation(irp);

    in_buf = out_buf = irp->AssociatedIrp.SystemBuffer;
    in_buf_size = io_stack_loc->Parameters.DeviceIoControl.InputBufferLength;
    out_buf_size = io_stack_loc->Parameters.DeviceIoControl.OutputBufferLength;

    switch (io_stack_loc->Parameters.DeviceIoControl.IoControlCode)
    {
        case SAY_HELLO:
            handler = handle_say_hello;
            break;
        case EXCHANGE_INTS:
            handler = handle_exchange_ints;
            break;
        default:
            status = irp->IoStatus.Status = STATUS_NOT_SUPPORTED;
            goto complete_request;
    }

    status = irp->IoStatus.Status = handler(in_buf, in_buf_size,
                                            out_buf, out_buf_size,
                                            &irp->IoStatus.Information);

complete_request:
    IoCompleteRequest(irp, IO_NO_INCREMENT);

    return status;
}

typedef struct
{
    const wchar_t *path;
    const wchar_t *symlink;
} device_info;

typedef struct
{
    DEVICE_OBJECT *object;
    UNICODE_STRING path;
    UNICODE_STRING symlink;
} device;

#define NUMOF_DEVICES 2

static device_info devices_info[NUMOF_DEVICES] =
{
    {
        L"\\Device\\test_device1",
        L"\\DosDevices\\test_device1",
    },
    {
        L"\\Device\\test_device2",
        L"\\DosDevices\\test_device2",
    },
};

static device devices[NUMOF_DEVICES];

static void destroy_device(int i)
{
    IoDeleteSymbolicLink(&devices[i].symlink);
    IoDeleteDevice(devices[i].object);
}

void destroy_devices()
{
    int i;
    for (i = 0; i < NUMOF_DEVICES; ++i)
        destroy_device(i);
}

static NTSTATUS set_up_device(DRIVER_OBJECT *driver_object, int i)
{
    NTSTATUS status = STATUS_SUCCESS;

    DbgPrint("Setting up device...\n");
    DbgPrint("\tPath: %ws\n", devices_info[i].path);
    DbgPrint("\tSymlink: %ws\n", devices_info[i].symlink);

    RtlInitUnicodeString(&devices[i].path, devices_info[i].path);
    RtlInitUnicodeString(&devices[i].symlink, devices_info[i].symlink);

    status = IoCreateDevice(driver_object,
                            0,
                            &devices[i].path,
                            FILE_DEVICE_UNKNOWN,
                            FILE_DEVICE_SECURE_OPEN,
                            FALSE,
                            &devices[i].object);

    if (!NT_SUCCESS(status))
        return status;

    devices[i].object->Flags |= DO_BUFFERED_IO;
    devices[i].object->Flags &= ~DO_DEVICE_INITIALIZING;

    if (!NT_SUCCESS(status = IoCreateSymbolicLink(&devices[i].symlink,
                                                  &devices[i].path)))
        goto delete_device;

    return status;

delete_device:
    IoDeleteDevice(devices[i].object);

    return status;
}

NTSTATUS set_up_devices(DRIVER_OBJECT *driver_object)
{
    int i, j;
    NTSTATUS status = STATUS_SUCCESS;

    for (i = 0; i < NUMOF_DEVICES; ++i)
        if (!NT_SUCCESS(status = set_up_device(driver_object, i)))
            goto destroy_devices;

    driver_object->MajorFunction[IRP_MJ_CREATE] = device_open;
    driver_object->MajorFunction[IRP_MJ_DEVICE_CONTROL] = device_ioctl;

    return status;

destroy_devices:
    for (j = 0; j < i; ++j)
        destroy_device(j);

    return status;
}