[install] Add reg_set_sz() function
[sanbootconf.git] / driver / registry.c
1 /*
2  * Copyright (C) 2008 Michael Brown <mbrown@fensystems.co.uk>.
3  *
4  * This program is free software; you can redistribute it and/or
5  * modify it under the terms of the GNU General Public License as
6  * published by the Free Software Foundation; either version 2 of the
7  * License, or any later version.
8  *
9  * This program is distributed in the hope that it will be useful, but
10  * WITHOUT ANY WARRANTY; without even the implied warranty of
11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
12  * General Public License for more details.
13  *
14  * You should have received a copy of the GNU General Public License
15  * along with this program; if not, write to the Free Software
16  * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
17  */
18
19 #include <ntddk.h>
20 #include <ntstrsafe.h>
21 #include "sanbootconf.h"
22 #include "registry.h"
23
24 /**
25  * Open registry key
26  *
27  * @v reg_key_name      Registry key name
28  * @v reg_key           Registry key to fill in
29  * @ret ntstatus        NT status
30  */
31 NTSTATUS reg_open ( LPCWSTR reg_key_name, PHANDLE reg_key ) {
32         UNICODE_STRING unicode_string;
33         OBJECT_ATTRIBUTES object_attrs;
34         NTSTATUS status;
35
36         RtlInitUnicodeString ( &unicode_string, reg_key_name );
37         InitializeObjectAttributes ( &object_attrs, &unicode_string,
38                                      OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE,
39                                      NULL, NULL );
40         status = ZwOpenKey ( reg_key, KEY_ALL_ACCESS, &object_attrs );
41         if ( ! NT_SUCCESS ( status ) ) {
42                 DbgPrint ( "Could not open %S: %x\n", reg_key_name, status );
43                 return status;
44         }
45
46         return STATUS_SUCCESS;
47 }
48
49 /**
50  * Close registry key
51  *
52  * @v reg_key           Registry key
53  */
54 VOID reg_close ( HANDLE reg_key ) {
55         ZwClose ( reg_key );
56 }
57
58 /**
59  * Fetch registry key value information
60  *
61  * @v reg_key           Registry key
62  * @v value_name        Registry value name
63  * @v kvi               Key value information block to allocate and fill in
64  * @ret ntstatus        NT status
65  *
66  * The caller must eventually free the allocated key value information
67  * block.
68  */
69 NTSTATUS fetch_reg_kvi ( HANDLE reg_key, LPCWSTR value_name,
70                          PKEY_VALUE_PARTIAL_INFORMATION *kvi ) {
71         UNICODE_STRING u_value_name;
72         ULONG kvi_len;
73         NTSTATUS status;
74
75         /* Get value length */
76         RtlInitUnicodeString ( &u_value_name, value_name );
77         status = ZwQueryValueKey ( reg_key, &u_value_name,
78                                    KeyValuePartialInformation, NULL, 0,
79                                    &kvi_len );
80         if ( ! ( ( status == STATUS_SUCCESS ) ||
81                  ( status == STATUS_BUFFER_OVERFLOW ) ||
82                  ( status == STATUS_BUFFER_TOO_SMALL ) ) ) {
83                 DbgPrint ( "Could not get KVI length for \"%S\": %x\n",
84                            value_name, status );
85                 goto err_zwqueryvaluekey_len;
86         }
87
88         /* Allocate value buffer */
89         *kvi = ExAllocatePoolWithTag ( NonPagedPool, kvi_len,
90                                        SANBOOTCONF_POOL_TAG );
91         if ( ! *kvi ) {
92                 DbgPrint ( "Could not allocate KVI for \"%S\": %x\n",
93                            value_name, status );
94                 goto err_exallocatepoolwithtag_kvi;
95         }
96
97         /* Fetch value */
98         status = ZwQueryValueKey ( reg_key, &u_value_name,
99                                    KeyValuePartialInformation, *kvi,
100                                    kvi_len, &kvi_len );
101         if ( ! NT_SUCCESS ( status ) ) {
102                 DbgPrint ( "Could not get KVI for \"%S\": %x\n",
103                            value_name, status );
104                 goto err_zwqueryvaluekey;
105         }
106
107         return STATUS_SUCCESS;
108
109  err_zwqueryvaluekey:
110         ExFreePool ( kvi );
111  err_exallocatepoolwithtag_kvi:
112  err_zwqueryvaluekey_len:
113         return status;
114 }
115
116 /**
117  * Fetch registry string value
118  *
119  * @v reg_key           Registry key
120  * @v value_name        Registry value name
121  * @v value             String value to allocate and fill in
122  * @ret ntstatus        NT status
123  *
124  * The caller must eventually free the allocated value.
125  */
126 NTSTATUS fetch_reg_sz ( HANDLE reg_key, LPCWSTR value_name, LPWSTR *value ) {
127         PKEY_VALUE_PARTIAL_INFORMATION kvi;
128         ULONG value_len;
129         NTSTATUS status;
130
131         /* Fetch key value information */
132         status = fetch_reg_kvi ( reg_key, value_name, &kvi );
133         if ( ! NT_SUCCESS ( status ) )
134                 goto err_fetch_reg_kvi;
135
136         /* Allocate and populate string */
137         value_len = ( kvi->DataLength + sizeof ( value[0] ) );
138         *value = ExAllocatePoolWithTag ( NonPagedPool, value_len,
139                                          SANBOOTCONF_POOL_TAG );
140         if ( ! *value ) {
141                 DbgPrint ( "Could not allocate value for \"%S\"\n",
142                            value_name );
143                 status = STATUS_UNSUCCESSFUL;
144                 goto err_exallocatepoolwithtag_value;
145         }
146         RtlZeroMemory ( *value, value_len );
147         RtlCopyMemory ( *value, kvi->Data, kvi->DataLength );
148
149  err_exallocatepoolwithtag_value:
150         ExFreePool ( kvi );
151  err_fetch_reg_kvi:
152         return status;
153 }
154
155 /**
156  * Fetch registry multiple-string value
157  *
158  * @v reg_key           Registry key
159  * @v value_name        Registry value name
160  * @v values            Array of string values to allocate and fill in
161  * @ret ntstatus        NT status
162  *
163  * The caller must eventually free the allocated values.
164  */
165 NTSTATUS fetch_reg_multi_sz ( HANDLE reg_key, LPCWSTR value_name,
166                               LPWSTR **values ) {
167         PKEY_VALUE_PARTIAL_INFORMATION kvi;
168         LPWSTR string;
169         ULONG num_strings;
170         ULONG values_len;
171         ULONG i;
172         NTSTATUS status;
173
174         /* Fetch key value information */
175         status = fetch_reg_kvi ( reg_key, value_name, &kvi );
176         if ( ! NT_SUCCESS ( status ) )
177                 goto err_fetch_reg_kvi;
178
179         /* Count number of strings in the array.  This is a
180          * potential(ly harmless) overestimate.
181          */
182         num_strings = 0;
183         for ( string = ( ( LPWSTR ) kvi->Data ) ;
184               string < ( ( LPWSTR ) ( kvi->Data + kvi->DataLength ) ) ;
185               string++ ) {
186                 if ( ! *string )
187                         num_strings++;
188         }
189         
190         /* Allocate and populate string array */
191         values_len = ( ( ( num_strings + 1 ) * sizeof ( values[0] ) ) +
192                        kvi->DataLength + sizeof ( values[0][0] ) );
193         *values = ExAllocatePoolWithTag ( NonPagedPool, values_len,
194                                           SANBOOTCONF_POOL_TAG );
195         if ( ! *values ) {
196                 DbgPrint ( "Could not allocate value array for \"%S\"\n",
197                            value_name );
198                 status = STATUS_UNSUCCESSFUL;
199                 goto err_exallocatepoolwithtag_value;
200         }
201         RtlZeroMemory ( *values, values_len );
202         string = ( ( LPWSTR ) ( *values + num_strings + 1 ) );
203         RtlCopyMemory ( string, kvi->Data, kvi->DataLength );
204         for ( i = 0 ; i < num_strings ; i++ ) {
205                 (*values)[i] = string;
206                 while ( *string )
207                         string++;
208                 while ( ! *string )
209                         string++;
210         }
211
212  err_exallocatepoolwithtag_value:
213         ExFreePool ( kvi );
214  err_fetch_reg_kvi:
215         return status;
216 }
217
218 /**
219  * Store registry string value
220  *
221  * @v reg_key           Registry key
222  * @v value_name        Registry value name
223  * @v value             String value to store
224  * @ret ntstatus        NT status
225  */
226 NTSTATUS reg_store_sz ( HANDLE reg_key, LPCWSTR value_name, LPWSTR value ) {
227         UNICODE_STRING u_value_name;
228         SIZE_T value_len;
229         NTSTATUS status;
230
231         RtlInitUnicodeString ( &u_value_name, value_name );
232         value_len = ( ( wcslen ( value ) + 1 ) * sizeof ( value[0] ) );
233         status = ZwSetValueKey ( reg_key, &u_value_name, 0, REG_SZ,
234                                  value, ( ( ULONG ) value_len ) );
235         if ( ! NT_SUCCESS ( status ) ) {
236                 DbgPrint ( "Could not store value \"%S\": %x\n",
237                            value_name, status );
238                 return status;
239         }
240
241         return STATUS_SUCCESS;
242 }
243
244 /**
245  * Store registry multiple-string value
246  *
247  * @v reg_key           Registry key
248  * @v value_name        Registry value name
249  * @v ...               String values to store (NULL terminated)
250  * @ret ntstatus        NT status
251  */
252 NTSTATUS reg_store_multi_sz ( HANDLE reg_key, LPCWSTR value_name, ... ) {
253         UNICODE_STRING u_value_name;
254         va_list args;
255         LPCWSTR string;
256         SIZE_T values_len;
257         LPWSTR values;
258         LPWSTR value;
259         SIZE_T values_remaining;
260         SIZE_T value_len;
261         NTSTATUS status;
262
263         /* Calculate total buffer length */
264         values_len = sizeof ( string[0] );
265         va_start ( args, value_name );
266         while ( ( string = va_arg ( args, LPCWSTR ) ) != NULL ) {
267                 values_len += ( ( wcslen ( string ) + 1 ) *
268                                 sizeof ( string[0] ) );
269         }
270         va_end ( args );
271
272         /* Allocate buffer */
273         values = ExAllocatePoolWithTag ( NonPagedPool, values_len,
274                                          SANBOOTCONF_POOL_TAG );
275         if ( ! values ) {
276                 DbgPrint ( "Could not allocate value buffer for \"%S\"\n" );
277                 status = STATUS_UNSUCCESSFUL;
278                 goto err_exallocatepoolwithtag;
279         }
280
281         /* Copy strings into buffer */
282         RtlZeroMemory ( values, values_len );
283         value = values;
284         values_remaining = values_len;
285         va_start ( args, value_name );
286         while ( ( string = va_arg ( args, LPCWSTR ) ) != NULL ) {
287                 RtlStringCbCatW ( value, values_remaining, string );
288                 value_len = ( ( wcslen ( value ) + 1 ) * sizeof ( value[0] ) );
289                 value += ( value_len / sizeof ( value[0] ) );
290                 values_remaining -= value_len;
291         }
292         va_end ( args );
293
294         /* Store value */
295         RtlInitUnicodeString ( &u_value_name, value_name );
296         status = ZwSetValueKey ( reg_key, &u_value_name, 0, REG_MULTI_SZ,
297                                  values, ( ( ULONG ) values_len ) );
298         if ( ! NT_SUCCESS ( status ) ) {
299                 DbgPrint ( "Could not store value \"%S\": %x\n",
300                            value_name, status );
301                 goto err_zwsetvaluekey;
302         }
303
304  err_zwsetvaluekey:
305         ExFreePool ( values );
306  err_exallocatepoolwithtag:
307         return STATUS_SUCCESS;
308 }
309
310 /**
311  * Store registry dword value
312  *
313  * @v reg_key           Registry key
314  * @v value_name        Registry value name
315  * @v value             String value to store, or NULL
316  * @ret ntstatus        NT status
317  */
318 NTSTATUS reg_store_dword ( HANDLE reg_key, LPCWSTR value_name, ULONG value ) {
319         UNICODE_STRING u_value_name;
320         NTSTATUS status;
321
322         RtlInitUnicodeString ( &u_value_name, value_name );
323         status = ZwSetValueKey ( reg_key, &u_value_name, 0, REG_DWORD,
324                                  &value, sizeof ( value ) );
325         if ( ! NT_SUCCESS ( status ) ) {
326                 DbgPrint ( "Could not store value \"%S\": %x\n",
327                            value_name, status );
328                 return status;
329         }
330
331         return STATUS_SUCCESS;
332 }