windows-tests-nm-20090130
[openafs.git] / src / WINNT / tests / nmtest / nmtest.c
1 /* 
2  * Copyright (c) 2009 - Secure Endpoints Inc.
3  *
4  * Author: Asanka Herath <asanka@secure-endpoints.com>
5  */
6
7 #include<windows.h>
8 #include<wincrypt.h>
9 #include<stdlib.h>
10 #include<stdio.h>
11
12 #define BLOCKSIZE 1024
13 #define MAX_PARAM 1
14
15 typedef __int64 offset_t;
16
17 HCRYPTPROV h_prov = 0;
18
19 offset_t N = 0;
20 offset_t M = 0;
21 const char * filename = "";
22 BOOL  verify = FALSE;
23 BOOL  show_offsets = FALSE;
24
25 typedef struct hash_data {
26     offset_t offset;
27     DWORD    param;
28 } hash_data;
29
30 #define ROUNDUP(x, align) (((((x) - 1) / (align)) + 1) * (align))
31
32 void dump_hex(BYTE * buffer, int cb)
33 {
34     static const char *htable[] = {
35         "0","1","2","3","4","5","6","7","8","9","A","B","C","D","E","F"
36     };
37     int i;
38
39     for (i=0; i < cb; i++) {
40         if ((i % 16) == 0) {
41             fprintf(stderr, "\n%08X : ", i);
42         }
43
44         fprintf(stderr, "%s%s",
45                 htable[(buffer[i] >> 4) & 0xf],
46                 htable[(buffer[i] & 0xf)]);
47
48         if ((i % 16) == 8)
49             fprintf(stderr, "-");
50         else
51             fprintf(stderr, " ");
52     }
53
54     fprintf(stderr, "\n");
55 }
56
57 #define CCall(f) if (!f) goto done;
58
59 BOOL init_test_data(void)
60 {
61     return CryptAcquireContext(&h_prov, NULL, MS_DEF_PROV, PROV_RSA_FULL,
62                                CRYPT_VERIFYCONTEXT | CRYPT_SILENT);
63 }
64
65 void exit_test_data(void)
66 {
67     if (h_prov)
68         CryptReleaseContext(h_prov, 0);
69     h_prov = 0;
70 }
71
72 BOOL generate_test_data_block(offset_t offset, DWORD param, BYTE buffer[BLOCKSIZE])
73 {
74     HCRYPTHASH h_hash = 0;
75     HCRYPTKEY  h_key = 0;
76     hash_data  d;
77     DWORD      cb_data;
78     BOOL rv = FALSE;
79
80     ZeroMemory(&d, sizeof(d));
81
82     d.offset = offset;
83     d.param = param;
84
85     CCall(CryptCreateHash(h_prov, CALG_SHA1, 0, 0, &h_hash));
86     CCall(CryptHashData(h_hash, (BYTE *) &d, sizeof(d), 0));
87     CCall(CryptDeriveKey(h_prov, CALG_RC4, h_hash, CRYPT_NO_SALT, &h_key));
88
89     ZeroMemory(buffer, BLOCKSIZE);
90     cb_data = BLOCKSIZE;
91     CCall(CryptEncrypt(h_key, 0, TRUE, 0, buffer, &cb_data, cb_data));
92
93     //dump_hex(buffer, 32);
94
95     rv = TRUE;
96  done:
97
98     if (h_hash)
99         CryptDestroyHash(h_hash);
100     if (h_key)
101         CryptDestroyKey(h_key);
102
103     return rv;
104 }
105
106 BOOL verify_test_data_block(offset_t offset, DWORD param, BYTE buffer[BLOCKSIZE])
107 {
108     BYTE expected[BLOCKSIZE];
109
110     if (!generate_test_data_block(offset, param, expected)) {
111         return FALSE;
112     }
113     if (!memcmp(expected, buffer, BLOCKSIZE)) {
114         if (param > 0 && show_offsets)
115             printf("... [%I64d]\n", offset, param);
116         return TRUE;
117     }
118
119     return FALSE;
120 }
121
122 typedef struct bitmap {
123     DWORD    magic;
124     offset_t length;
125     offset_t o_data;
126     BYTE     bits[1];
127 } bitmap;
128
129 #define BMAGIC 0xbeefface
130
131 #define NBITMAPBYTES(n) (ROUNDUP((n)/BLOCKSIZE, 8)/8)
132 #define BITMAPSIZE(n) (sizeof(bitmap) + NBITMAPBYTES(n) - 1)
133 #define FILESIZE(n)   (DATAOFFSET(n) + (n))
134 #define DATAOFFSET(n) ROUNDUP(BITMAPSIZE(n), BLOCKSIZE)
135
136 bitmap * allocbits(offset_t N)
137 {
138     bitmap * b;
139
140     b = malloc(BITMAPSIZE(N));
141     if (b == NULL)
142         return NULL;
143
144     memset(b, 0, BITMAPSIZE(N));
145
146     b->magic = BMAGIC;
147     b->length = N;
148     b->o_data = DATAOFFSET(N);
149
150     return b;
151 }
152
153 void freebits(bitmap * b)
154 {
155     if (b)
156         free(b);
157 }
158
159 void setbit(bitmap * b, offset_t o)
160 {
161     if (o < b->o_data || o >= b->o_data + b->length || (o % BLOCKSIZE) != 0) {
162         fprintf(stderr, "Internal error. Invalid offset\n");
163         exit(1);
164     }
165
166     o = (o - b->o_data) / BLOCKSIZE;
167
168     b->bits[o / 8] |= (1 << (o & 7));
169 }
170
171 BOOL getbit(bitmap * b, offset_t o)
172 {
173     if (o < b->o_data || o >= b->o_data + b->length || (o % BLOCKSIZE) != 0) {
174         fprintf(stderr, "Internal error. Invalid offset\n");
175         exit(1);
176     }
177
178     o = (o - b->o_data) / BLOCKSIZE;
179
180     return !!(b->bits[o / 8] & (1 << (o & 7)));
181 }
182
183 BOOL write_bitmap(HANDLE h, bitmap * b)
184 {
185     LARGE_INTEGER li;
186     OVERLAPPED ov;
187     offset_t b_size, o;
188     BOOL success = FALSE;
189
190     ZeroMemory(&ov, sizeof(ov));
191
192     b_size = BITMAPSIZE(b->length);
193
194     li.QuadPart = 0;
195
196     if (!SetFilePointerEx(h, li, &li, FILE_BEGIN)) {
197         fprintf(stderr, "Can't set file pointer. GLE=%d\n", GetLastError());
198         goto done;
199     }
200
201     for (o = 0; o < b_size; o += BLOCKSIZE) {
202         BYTE buffer[BLOCKSIZE];
203         DWORD n_written = 0;
204
205         if (o + BLOCKSIZE <= b_size) {
206             memcpy(buffer, ((BYTE *) b) + o, BLOCKSIZE);
207         } else {
208             memcpy(buffer, ((BYTE *) b) + o, (b_size - o));
209             memset(buffer + (b_size - o), 0, BLOCKSIZE - (b_size - o));
210         }
211
212         li.QuadPart = o;
213         ov.Offset = li.LowPart;
214         ov.OffsetHigh = li.HighPart;
215
216         if (!LockFileEx(h, LOCKFILE_EXCLUSIVE_LOCK, 0, BLOCKSIZE, 0, &ov)) {
217             fprintf(stderr, "Can't lock file. GLE=%d\n", GetLastError());
218             goto done;
219         }
220
221         if (!WriteFile(h, buffer, BLOCKSIZE, &n_written, NULL) || n_written != BLOCKSIZE) {
222             fprintf(stderr, "Can't write data. GLE=%d\n", GetLastError());
223             goto done;
224         }
225
226         if (!UnlockFileEx(h, 0, BLOCKSIZE, 0, &ov)) {
227             fprintf(stderr, "Can't unlock file. GLE=%d\n", GetLastError());
228             goto done;
229         }
230     }
231
232     li.QuadPart = 0;
233     if (!SetFilePointerEx(h, li, &li, FILE_CURRENT)) {
234         fprintf(stderr, "Can't set file pointer. GLE=%d\n", GetLastError());
235         goto done;
236     }
237
238     if (li.QuadPart != b->o_data) {
239         fprintf(stderr, "Current file pointer is not at start of data.\n");
240         goto done;
241     }
242
243     success = TRUE;
244
245  done:
246
247     return success;    
248 }
249
250 bitmap* read_bitmap(HANDLE h)
251 {
252     LARGE_INTEGER li;
253     OVERLAPPED ov;
254     bitmap * bfile = NULL;
255     bitmap * bprep = NULL;
256     bitmap * brv = NULL;
257     BYTE buffer[BLOCKSIZE];
258     DWORD n_read = 0;
259     offset_t o;
260
261     ZeroMemory(&ov, sizeof(ov));
262
263     li.QuadPart = 0;
264
265     if (!SetFilePointerEx(h, li, &li, FILE_BEGIN)) {
266         fprintf(stderr, "Can't set file pointer. GLE=%d\n", GetLastError());
267         goto done;
268     }
269
270     ov.Offset = 0;
271     ov.OffsetHigh = 0;
272
273     if (!LockFileEx(h, 0, 0, BLOCKSIZE, 0, &ov)) {
274         fprintf(stderr, "Can't lock file. GLE=%d\n", GetLastError());
275         goto done;
276     }
277
278     if (!ReadFile(h, buffer, BLOCKSIZE, &n_read, NULL) || n_read != BLOCKSIZE) {
279         fprintf(stderr, "Can't read data. GLE=%d\n", GetLastError());
280         goto done;
281     }
282
283     if (!UnlockFileEx(h, 0, BLOCKSIZE, 0, &ov)) {
284         fprintf(stderr, "Can't unlock file. GLE=%d\n", GetLastError());
285         goto done;
286     }
287
288     bfile = (bitmap *)buffer;
289
290     if (bfile->magic != BMAGIC) {
291         fprintf(stderr, "Corrupt data. Magic number is invalid\n");
292         goto done;
293     }
294
295     if (!GetFileSizeEx(h, &li)) {
296         fprintf(stderr, "Can't get file size. GLE=%d\n", GetLastError());
297         goto done;
298     }
299
300     if (li.QuadPart != FILESIZE(bfile->length)) {
301         fprintf(stderr, "Corrupt data. Invalid file size.\n");
302         goto done;
303     }
304
305     if (bfile->o_data != DATAOFFSET(bfile->length)) {
306         fprintf(stderr, "Corrupt data. Invalid data offset.\n");
307         goto done;
308     }
309
310     bprep = allocbits(bfile->length);
311
312     if (bprep == NULL) {
313         fprintf(stderr, "Can't allocate memory for bitmap\n");
314         goto done;
315     }
316
317     memcpy(bprep, bfile, __min(BITMAPSIZE(bfile->length), BLOCKSIZE));
318
319     for (o = BLOCKSIZE; o < BITMAPSIZE(bprep->length); o += BLOCKSIZE) {
320         li.QuadPart = o;
321         ov.Offset = li.LowPart;
322         ov.OffsetHigh = li.HighPart;
323
324         if (!LockFileEx(h, 0, 0, BLOCKSIZE, 0, &ov)) {
325             fprintf(stderr, "Can't lock file. GLE=%d\n", GetLastError());
326             goto done;
327         }
328
329         if (!ReadFile(h, buffer, BLOCKSIZE, &n_read, NULL) || n_read != BLOCKSIZE) {
330             fprintf(stderr, "Can't read data. GLE=%d\n", GetLastError());
331             goto done;
332         }
333
334         if (!UnlockFileEx(h, 0, BLOCKSIZE, 0, &ov)) {
335             fprintf(stderr, "Can't unlock file. GLE=%d\n", GetLastError());
336             goto done;
337         }
338
339         memcpy(((BYTE *) bprep) + o, buffer, __min(BITMAPSIZE(bprep->length) - o, BLOCKSIZE));
340     }
341
342     li.QuadPart = 0;
343     if (!SetFilePointerEx(h, li, &li, FILE_CURRENT)) {
344         fprintf(stderr, "Can't set file pointer. GLE=%d\n", GetLastError());
345         goto done;
346     }
347
348     if (li.QuadPart != bprep->o_data) {
349         fprintf(stderr, "Current file pointer not at start of data.\n");
350         goto done;
351     }
352
353     brv = bprep;
354
355  done:
356     if (brv == NULL && bprep != NULL) {
357         freebits(bprep);
358     }
359     return brv;
360 }
361
362 int do_verify_test(void)
363 {
364     HANDLE h_file = NULL;
365     int rv = 1;
366     offset_t offset;
367     LARGE_INTEGER li;
368     OVERLAPPED ov;
369     bitmap * b = NULL;
370
371     printf("Verifying test data file [%s]\n", filename);
372
373     h_file = CreateFile(filename, GENERIC_READ, FILE_SHARE_WRITE, NULL,
374                         OPEN_EXISTING, FILE_FLAG_SEQUENTIAL_SCAN, NULL);
375     if (h_file == INVALID_HANDLE_VALUE) {
376         fprintf(stderr, "Can't open file [%s] GLE=%d\n", filename, GetLastError());
377         goto done;
378     }
379
380     b = read_bitmap(h_file);
381
382     if (b == NULL) {
383         goto done;
384     }
385
386     N = b->length;
387
388     printf("File size N = %I64d\n", N);
389
390     if (!init_test_data()) {
391         fprintf(stderr, "Initialize crypto. GLE=%d\n", GetLastError());
392         goto done;
393     }
394
395     printf("Verifying data ... \n");
396
397     ZeroMemory(&ov, sizeof(ov));
398
399     for (offset = b->o_data; offset < b->o_data+N; offset += BLOCKSIZE) {
400         BYTE buffer[BLOCKSIZE];
401         DWORD n_read = 0;
402
403         li.QuadPart = offset;
404         ov.Offset = li.LowPart;
405         ov.OffsetHigh = li.HighPart;
406
407         if (!LockFileEx(h_file, 0, 0, BLOCKSIZE, 0, &ov)) {
408             printf("ERROR\n");
409             fprintf(stderr, "Can't lock file. GLE=%d\n", GetLastError());
410             goto done;
411         }
412
413         if (!ReadFile(h_file, buffer, BLOCKSIZE, &n_read, NULL) || n_read != BLOCKSIZE) {
414             printf("ERROR\n");
415             fprintf(stderr, "Can't read data. GLE=%d\n", GetLastError());
416             goto done;
417         }
418
419         if (!verify_test_data_block(offset, ((getbit(b, offset))? 1 : 0), buffer)) {
420             printf("VERIFY FAILED\n");
421             fprintf(stderr, "Verification failed at offset %I64d\n", offset);
422             goto done;
423         }
424
425         if (!UnlockFileEx(h_file, 0, BLOCKSIZE, 0, &ov)) {
426             printf("ERROR\n");
427             fprintf(stderr, "Can't unlock file. GLE=%d\n", GetLastError());
428             goto done;
429         }
430     }
431
432     printf("Verify succeeded!\n");
433
434     rv = 0;
435
436  done:
437     if (b)
438         freebits(b);
439
440     if (h_file)
441         CloseHandle(h_file);
442
443     exit_test_data();
444
445     return rv;
446 }
447
448 int do_write_test(void)
449 {
450     HANDLE h_file = NULL;
451     int rv = 1;
452     offset_t offset;
453     OVERLAPPED ov;
454     bitmap * b = NULL;
455
456     printf("Generating test data file [%s]\n", filename);
457
458     N = (((N - 1)/BLOCKSIZE) + 1) * BLOCKSIZE;
459     M = (((M - 1)/BLOCKSIZE) + 1) * BLOCKSIZE;
460
461     printf("Using N = %I64d and M = %I64d\n", N, M);
462
463     h_file = CreateFile(filename, GENERIC_WRITE, FILE_SHARE_WRITE, NULL,
464                         CREATE_ALWAYS, FILE_FLAG_RANDOM_ACCESS, NULL);
465     if (h_file == INVALID_HANDLE_VALUE) {
466         fprintf(stderr, "Can't create file [%s] GLE=%d\n", filename, GetLastError());
467         goto done;
468     }
469
470     b = allocbits(N);
471     if (b == NULL) {
472         fprintf(stderr, "Can't allocate bitmap.\n");
473         goto done;
474     }
475
476     if (!init_test_data()) {
477         fprintf(stderr, "Initialize crypto. GLE=%d\n", GetLastError());
478         goto done;
479     }
480
481     printf("Phase 1: Generating test data ... ");
482
483     if (!write_bitmap(h_file, b)) {
484         goto done;
485     }
486
487     ZeroMemory(&ov, sizeof(ov));
488
489     for (offset = b->o_data; offset < b->o_data + N; offset += BLOCKSIZE) {
490         BYTE buffer[BLOCKSIZE];
491         DWORD n_written = 0;
492         LARGE_INTEGER li;
493
494         if (!generate_test_data_block(offset, 0, buffer)) {
495             printf("ERROR\n");
496             fprintf(stderr, "Can't generate test data. GLE=%d\n", GetLastError());
497             goto done;
498         }
499
500         li.QuadPart = offset;
501         ov.Offset = li.LowPart;
502         ov.OffsetHigh = li.HighPart;
503
504         if (!LockFileEx(h_file, LOCKFILE_EXCLUSIVE_LOCK, 0, BLOCKSIZE, 0, &ov)) {
505             printf("ERROR\n");
506             fprintf(stderr, "Can't lock file. GLE=%d\n", GetLastError());
507             goto done;
508         }
509
510         if (!WriteFile(h_file, buffer, BLOCKSIZE, &n_written, NULL) || n_written != BLOCKSIZE) {
511             printf("ERROR\n");
512             fprintf(stderr, "Can't write data. GLE=%d\n", GetLastError());
513             goto done;
514         }
515
516         if (!UnlockFileEx(h_file, 0, BLOCKSIZE, 0, &ov)) {
517             printf("ERROR\n");
518             fprintf(stderr, "Can't unlock file. GLE=%d\n", GetLastError());
519             goto done;
520         }
521     }
522
523     printf("Done\n");
524
525     if (!FlushFileBuffers(h_file)) {
526         fprintf(stderr, "Can't flush file. GLE=%d\n", GetLastError());
527         goto done;
528     }
529
530     printf("Phase 2: Overwriting test data ... \n");
531
532     for (offset = 0; offset < M; offset += BLOCKSIZE) {
533         offset_t orandom;
534         BYTE buffer[BLOCKSIZE];
535         LARGE_INTEGER li;
536         DWORD n_written = 0;
537
538         if (!CryptGenRandom(h_prov, sizeof(orandom), (BYTE *) &orandom)) {
539             printf("ERROR\n");
540             fprintf(stderr, "Can't generate random number. GLE=%d\n", GetLastError());
541             goto done;
542         }
543
544         if (orandom < 0)
545             orandom = -orandom;
546         orandom = ((orandom % N) / BLOCKSIZE) * BLOCKSIZE;
547         orandom += b->o_data;
548
549         if (show_offsets)
550             printf("... [%I64d]\n", orandom);
551
552         if (!generate_test_data_block(orandom, 1, buffer)) {
553             printf("ERROR\n");
554             fprintf(stderr, "Can't generate test data. GLE=%d\n", GetLastError());
555             goto done;
556         }
557
558         li.QuadPart = orandom;
559         ov.Offset = li.LowPart;
560         ov.OffsetHigh = li.HighPart;
561
562         if (!LockFileEx(h_file, LOCKFILE_EXCLUSIVE_LOCK, 0, BLOCKSIZE, 0, &ov)) {
563             printf("ERROR\n");
564             fprintf(stderr, "Can't lock file. GLE=%d\n", GetLastError());
565             goto done;
566         }
567
568         if (!SetFilePointerEx(h_file, li, NULL, FILE_BEGIN)) {
569             printf("ERROR\n");
570             fprintf(stderr, "Can't set file pointer. GLE=%d\n", GetLastError());
571             goto done;
572         }
573
574         if (!WriteFile(h_file, buffer, BLOCKSIZE, &n_written, NULL) || n_written != BLOCKSIZE) {
575             printf("ERROR\n");
576             fprintf(stderr, "Can't write data. GLE=%d\n", GetLastError());
577             goto done;
578         }
579
580         if (!UnlockFileEx(h_file, 0, BLOCKSIZE, 0, &ov)) {
581             printf("ERROR\n");
582             fprintf(stderr, "Can't unlock file. GLE=%d\n", GetLastError());
583             goto done;
584         }
585
586         setbit(b, orandom);
587     }
588
589     if (!write_bitmap(h_file, b)) {
590         goto done;
591     }
592
593     printf("Done.\n");
594
595     rv = 0;
596
597  done:
598     if (b)
599         freebits(b);
600
601     if (h_file)
602         CloseHandle(h_file);
603
604     exit_test_data();
605
606     return rv;
607 }
608
609 BOOL show_usage(const char * fn)
610 {
611     fprintf(stderr,
612             "%s : Generate or verify test data file.\n"
613             "\n"
614             "Usage: %s [-r <filename> | -w <N> <M> <filename>]\n"
615             "\n"
616             "  -w <N> <M> <filename> :\n"
617             "     First writes N bytes of random data into <filename> and then\n"
618             "     overwrites M bytes with different random data.\n"
619             "\n"
620             "  -r <filename> : \n"
621             "     Verify the contents of <filename>.  Verification succeeds if\n"
622             "     the contents of <filename> was generated using the -w option\n",
623             fn, fn);
624     return FALSE;
625 }
626
627 BOOL parse_cmdline(int argc, char ** argv)
628 {
629     if (argc == 3) {
630         if (strcmp(argv[1], "-r"))
631             return show_usage(argv[0]);
632         verify = TRUE;
633         N = 0;
634         M = 0;
635         filename = argv[2];
636         return TRUE;
637     }
638
639     if (argc == 5) {
640         if (strcmp(argv[1], "-w"))
641             return show_usage(argv[0]);
642         verify = FALSE;
643         N = atol(argv[2]);
644         if (N == 0)
645             return show_usage(argv[0]);
646         M = atol(argv[3]);
647         if (M == 0)
648             return show_usage(argv[0]);
649         filename = argv[4];
650         return TRUE;
651     }
652
653     return show_usage(argv[0]);
654 }
655
656 int do_tests(void)
657 {
658     if (verify)
659         return do_verify_test();
660     else
661         return do_write_test();
662 }
663
664 int main(int argc, char ** argv)
665 {
666     if (!parse_cmdline(argc, argv))
667         return 1;
668
669     return do_tests();
670 }