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 }