drivers/kmtest/reghelper.c

   1 /*
   2  * Driver Regression Tests
   3  *
   4  * Copyright 2009 Michael Martin <martinmnet@hotmail.com>
   5  *
   6  * This library is free software; you can redistribute it and/or
   7  * modify it under the terms of the GNU Library General Public
   8  * License as published by the Free Software Foundation; either
   9  * version 2 of the License, or (at your option) any later version.
  10  *
  11  * This library is distributed in the hope that it will be useful,
  12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14  * Library General Public License for more details.
  15  *
  16  * You should have received a copy of the GNU Library General Public
  17  * License along with this library; see the file COPYING.LIB.
  18  * If not, write to the Free Software Foundation,
  19  * 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  20  */
  21
  22 /* INCLUDES *******************************************************************/
  23
  24 #include "kmtest.h"
  25
  26 /*
  27     Adds a service registry entry for a driver
  28     The driver must reside in the same path as this loaded driver
  29     The caller is resposible for releasing memory
  30 */
  31 PWCHAR CreateLowerDeviceRegistryKey(PUNICODE_STRING RegistryPath, PWCHAR NewDriver)
  32 {
  33     OBJECT_ATTRIBUTES ObjectAttributes;
  34     UNICODE_STRING Name;
  35     UNICODE_STRING Value;
  36     UNICODE_STRING NewDriverRegPath;
  37     PKEY_VALUE_PARTIAL_INFORMATION ValuePartialInfo = NULL;
  38     HANDLE ServiceKey;
  39     NTSTATUS Status;
  40     ULONG Disposition;
  41     ULONG ServiceDWordValue;
  42     ULONG ResultLength = 0;
  43     ULONG Length = 0;
  44     PWCHAR ReturnPath = NULL;
  45     /* Now lets find out where we were loaded from by using registry */
  46     InitializeObjectAttributes(&ObjectAttributes, RegistryPath, OBJ_CASE_INSENSITIVE, NULL, NULL);
  47     Status = ZwOpenKey(&ServiceKey, KEY_READ, &ObjectAttributes);
  48     if (!NT_SUCCESS(Status))
  49     {
  50         DbgPrint("ZwOpenKey () failed (Status %x)\n", Status);
  51         return NULL;
  52     }
  53
  54     RtlInitUnicodeString(&Name, L"ImagePath");
  55
  56     /* First query how much memory we need */
  57     Status = ZwQueryValueKey(ServiceKey, &Name, KeyValuePartialInformation, 0, 0, &ResultLength);
  58
  59     ResultLength += sizeof(KEY_VALUE_PARTIAL_INFORMATION);
  60     ValuePartialInfo = ExAllocatePool(PagedPool, ResultLength);
  61     if (!ValuePartialInfo)
  62     {
  63         DbgPrint("Out of memory!\n");
  64         goto cleanup;
  65     }
  66
  67     Length = ResultLength;
  68     Status = ZwQueryValueKey(ServiceKey, &Name, KeyValuePartialInformation, (PVOID)ValuePartialInfo, Length, &ResultLength);
  69     if (!NT_SUCCESS(Status))
  70     {
  71         DbgPrint("ZwQueryValueKey() failed (Status %lx)\n", Status);
  72         goto cleanup;
  73     }
  74
  75     /* Remove the current driver name from the string */
  76     /* FIXME: Dont use hard coded driver name, determine it from the string returned from the above Query */
  77     Length = (wcslen((PWCHAR)ValuePartialInfo->Data) * 2) - (wcslen(L"kmtest.sys") * 2);
  78     RtlZeroMemory((PVOID)((ULONG_PTR)ValuePartialInfo->Data + Length),
  79     wcslen(L"drvtests.sys") * 2);
  80     ZwClose(ServiceKey);
  81
  82     /* Now add a registry entry for the driver */
  83
  84     NewDriverRegPath.Length = 0;
  85     NewDriverRegPath.MaximumLength = (wcslen(L"\\Registry\\Machine\\System\\CurrentControlSet\\Services\\") +
  86                                      wcslen(NewDriver) + 1) * sizeof(WCHAR);
  87     NewDriverRegPath.Buffer = ExAllocatePool(PagedPool, NewDriverRegPath.MaximumLength);
  88     if (!NewDriverRegPath.Buffer)
  89     {
  90         DbgPrint("Out of memory!\n");
  91         ExFreePool(NewDriverRegPath.Buffer);
  92         goto cleanup;
  93     }
  94
  95     RtlAppendUnicodeToString(&NewDriverRegPath, L"\\Registry\\Machine\\System\\CurrentControlSet\\Services\\");
  96     RtlAppendUnicodeToString(&NewDriverRegPath, NewDriver);
  97
  98     InitializeObjectAttributes(&ObjectAttributes,
  99                                &NewDriverRegPath,
 100                                OBJ_CASE_INSENSITIVE | OBJ_OPENIF,
 101                                0,
 102                                NULL);
 103
 104     Status = ZwCreateKey(&ServiceKey,
 105                          KEY_ALL_ACCESS,
 106                          &ObjectAttributes,
 107                          0,
 108                          NULL,
 109                          REG_OPTION_VOLATILE,
 110                          &Disposition);
 111
 112     if (!NT_SUCCESS(Status))
 113     {
 114         DbgPrint("ZwCreateKey() failed (Status %lx)\n", Status);
 115         ExFreePool(NewDriverRegPath.Buffer);
 116         goto cleanup;
 117     }
 118
 119     ReturnPath = NewDriverRegPath.Buffer;
 120     RtlInitUnicodeString(&Name, L"ImagePath");
 121
 122     Value.Length = 0;
 123     Value.MaximumLength = (wcslen((PWCHAR)ValuePartialInfo->Data) +
 124                           wcslen(NewDriver) + 5) * sizeof(WCHAR);
 125     Value.Buffer = ExAllocatePool(PagedPool, Value.MaximumLength);
 126
 127     if (!Value.Buffer)
 128     {
 129         DbgPrint("Out of memory!\n");
 130         ExFreePool(Value.Buffer);
 131         goto cleanup;
 132     }
 133
 134     RtlAppendUnicodeToString(&Value, (PWCHAR)ValuePartialInfo->Data);
 135     RtlAppendUnicodeToString(&Value, NewDriver);
 136     RtlAppendUnicodeToString(&Value, L".sys");
 137
 138     Status = ZwSetValueKey(ServiceKey,
 139                            &Name,
 140                            0,
 141                            REG_SZ,
 142                            Value.Buffer,
 143                            (wcslen(Value.Buffer)+1) * sizeof(WCHAR));
 144     ExFreePool(Value.Buffer);
 145     if (!NT_SUCCESS(Status))
 146     {
 147         DbgPrint("ZwCreateKey() failed (Status %lx)\n", Status);
 148         goto cleanup;
 149     }
 150
 151     RtlInitUnicodeString(&Name, L"DisplayName");
 152     RtlInitUnicodeString(&Value, NewDriver);
 153
 154     Status = ZwSetValueKey(ServiceKey,
 155                            &Name,
 156                            0,
 157                            REG_SZ,
 158                            Value.Buffer,
 159                            (wcslen(Value.Buffer)+1) * sizeof(WCHAR));
 160
 161     if (!NT_SUCCESS(Status))
 162     {
 163         DbgPrint("ZwCreateKey() failed (Status %lx)\n", Status);
 164         goto cleanup;
 165     }
 166
 167     RtlInitUnicodeString(&Name, L"ErrorControl");
 168     ServiceDWordValue = 0;
 169
 170     Status = ZwSetValueKey(ServiceKey,
 171                            &Name,
 172                            0,
 173                            REG_DWORD,
 174                            &ServiceDWordValue,
 175                            sizeof(ULONG));
 176
 177     if (!NT_SUCCESS(Status))
 178     {
 179         DbgPrint("ZwCreateKey() failed (Status %lx)\n", Status);
 180         goto cleanup;
 181     }
 182
 183     RtlInitUnicodeString(&Name, L"Start");
 184     ServiceDWordValue = 3;
 185     Status = ZwSetValueKey(ServiceKey,
 186                            &Name,
 187                            0,
 188                            REG_DWORD,
 189                            &ServiceDWordValue,
 190                            sizeof(ULONG));
 191
 192     if (!NT_SUCCESS(Status))
 193     {
 194         DbgPrint("ZwCreateKey() failed (Status %lx)\n", Status);
 195         goto cleanup;
 196     }
 197
 198     RtlInitUnicodeString(&Name, L"Type");
 199     ServiceDWordValue = 0;
 200     Status = ZwSetValueKey(ServiceKey,
 201                            &Name,
 202                            0,
 203                            REG_DWORD,
 204                            &ServiceDWordValue,
 205                            sizeof(ULONG));
 206
 207     if (!NT_SUCCESS(Status))
 208     {
 209         DbgPrint("ZwCreateKey() failed (Status %lx)\n", Status);
 210         goto cleanup;
 211     }
 212
 213 cleanup:
 214     ZwClose(ServiceKey);
 215     if (ValuePartialInfo) ExFreePool(ValuePartialInfo);
 216
 217     return ReturnPath;
 218
 219 }