tee.c

   1 /*
   2  * CertViewer - tee for Windows
   3  * Copyright (c) 2023 "dEajL3kA" <Cumpoing79@web.de>
   4  *
   5  * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
   6  * associated documentation files (the "Software"), to deal in the Software without restriction,
   7  * including without limitation the rights to use, copy, modify, merge, publish, distribute,
   8  * sub license, and/or sell copies of the Software, and to permit persons to whom the Software is
   9  * furnished to do so, subject to the following conditions: The above copyright notice and this
  10  * permission notice shall be included in all copies or substantial portions of the Software.
  11  *
  12  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
  13  * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  14  * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
  15  * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT
  16  * OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
  17  */
  18 #define WIN32_LEAN_AND_MEAN 1
  19 #include <Windows.h>
  20 #include <ShellAPI.h>
  21 #include <stdarg.h>
  22
  23 #pragma warning(disable: 4706)
  24 #define BUFFSIZE 8192U
  25
  26  // --------------------------------------------------------------------------
  27  // Utilities
  28  // --------------------------------------------------------------------------
  29
  30 static wchar_t to_lower(const wchar_t c)
  31 {
  32     return ((c >= L'A') && (c <= L'Z')) ? (L'a' + (c - L'A')) : c;
  33 }
  34
  35 static BOOL is_terminal(const HANDLE handle)
  36 {
  37     DWORD mode;
  38     return GetConsoleMode(handle, &mode);
  39 }
  40
  41 static const wchar_t *get_filename(const wchar_t *filePath)
  42 {
  43     for (const wchar_t *ptr = filePath; *ptr != L'\0'; ++ptr)
  44     {
  45         if ((*ptr == L'\\') || (*ptr == L'/'))
  46         {
  47             filePath = ptr + 1U;
  48         }
  49     }
  50     return filePath;
  51 }
  52
  53 static BOOL is_null_device(const wchar_t *filePath)
  54 {
  55     filePath = get_filename(filePath);
  56     if ((to_lower(filePath[0U]) == L'n') && (to_lower(filePath[1U]) == L'u') || (to_lower(filePath[2U]) == L'l'))
  57     {
  58         return ((filePath[3U] == L'\0') || (filePath[3U] == L'.'));
  59     }
  60     return FALSE;
  61 }
  62
  63 static wchar_t *concat_va(const wchar_t *const first, ...)
  64 {
  65     const wchar_t *ptr;
  66     va_list ap;
  67
  68     va_start(ap, first);
  69     size_t len = 0U;
  70     for (ptr = first; ptr != NULL; ptr = va_arg(ap, const wchar_t*))
  71     {
  72         len = lstrlenW(ptr);
  73     }
  74     va_end(ap);
  75
  76     wchar_t *const buffer = (wchar_t*)LocalAlloc(LPTR, sizeof(wchar_t) * (len + 1U));
  77     if (buffer)
  78     {
  79         va_start(ap, first);
  80         for (ptr = first; ptr != NULL; ptr = va_arg(ap, const wchar_t*))
  81         {
  82             lstrcatW(buffer, ptr);
  83         }
  84         va_end(ap);
  85     }
  86
  87     return buffer;
  88 }
  89
  90 #define CLOSE_HANDLE(HANDLE) do \
  91 { \
  92     if (((HANDLE) != NULL) && ((HANDLE) != INVALID_HANDLE_VALUE)) \
  93     { \
  94         CloseHandle((HANDLE)); \
  95         (HANDLE) = NULL; \
  96     } \
  97 } \
  98 while (0)
  99
 100 #define CONCAT(...) concat_va(__VA_ARGS__, NULL)
 101
 102 // --------------------------------------------------------------------------
 103 // Console CTRL+C handler
 104 // --------------------------------------------------------------------------
 105
 106 static volatile BOOL g_stop = FALSE;
 107
 108 static BOOL WINAPI console_handler(const DWORD ctrlType)
 109 {
 110     switch (ctrlType)
 111     {
 112     case CTRL_C_EVENT:
 113     case CTRL_BREAK_EVENT:
 114     case CTRL_CLOSE_EVENT:
 115         g_stop = TRUE;
 116         return TRUE;
 117     default:
 118         return FALSE;
 119     }
 120 }
 121
 122 // --------------------------------------------------------------------------
 123 // Version
 124 // --------------------------------------------------------------------------
 125
 126 static ULONGLONG get_version(void)
 127 {
 128     const HRSRC hVersion = FindResourceW(NULL, MAKEINTRESOURCE(VS_VERSION_INFO), RT_VERSION);
 129     if (hVersion)
 130     {
 131         const HGLOBAL hResource = LoadResource(NULL, hVersion);
 132         if (hResource)
 133         {
 134             const DWORD sizeOfResource = SizeofResource(NULL, hResource);
 135             if (sizeOfResource >= sizeof(VS_FIXEDFILEINFO))
 136             {
 137                 const PVOID addrResourceBlock = LockResource(hResource);
 138                 if (addrResourceBlock)
 139                 {
 140                     VS_FIXEDFILEINFO *fileInfoData;
 141                     UINT fileInfoSize;
 142                     if (VerQueryValueW(addrResourceBlock, L"\\", &fileInfoData, &fileInfoSize))
 143                     {
 144                         ULARGE_INTEGER fileVersion;
 145                         fileVersion.LowPart  = fileInfoData->dwFileVersionLS;
 146                         fileVersion.HighPart = fileInfoData->dwFileVersionMS;
 147                         return fileVersion.QuadPart;
 148                     }
 149                 }
 150             }
 151         }
 152     }
 153
 154     return 0U;
 155 }
 156
 157 static const wchar_t *get_version_string(void)
 158 {
 159     static wchar_t text[64U] = { '\0' };
 160     lstrcpyW(text, L"tee for Windows v#.#.# [" TEXT(__DATE__) L"]\n");
 161
 162     const ULONGLONG version = get_version();
 163     if (version)
 164     {
 165         text[17U] = L'0' + ((version >> 48) & 0xFFFF);
 166         text[19U] = L'0' + ((version >> 32) & 0xFFFF);
 167         text[21U] = L'0' + ((version >> 16) & 0xFFFF);
 168     }
 169
 170     return text;
 171 }
 172
 173 // --------------------------------------------------------------------------
 174 // Text output
 175 // --------------------------------------------------------------------------
 176
 177 static char *utf16_to_utf8(const wchar_t *const input)
 178 {
 179     const int buff_size = WideCharToMultiByte(CP_UTF8, 0, input, -1, NULL, 0, NULL, NULL);
 180     if (buff_size > 0)
 181     {
 182         char *const buffer = (char*)LocalAlloc(LPTR, buff_size);
 183         if (buffer)
 184         {
 185             const int result = WideCharToMultiByte(CP_UTF8, 0, input, -1, buffer, buff_size, NULL, NULL);
 186             if ((result > 0) && (result <= buff_size))
 187             {
 188                 return buffer;
 189             }
 190             LocalFree(buffer);
 191         }
 192     }
 193     return NULL;
 194 }
 195
 196 static BOOL write_text(const HANDLE handle, const wchar_t *const text)
 197 {
 198     BOOL result = FALSE;
 199     DWORD written;
 200     if (GetConsoleMode(handle, &written))
 201     {
 202         result = WriteConsoleW(handle, text, lstrlenW(text), &written, NULL);
 203     }
 204     else
 205     {
 206         char *const utf8_text = utf16_to_utf8(text);
 207         if (utf8_text)
 208         {
 209             result = WriteFile(handle, utf8_text, lstrlenA(utf8_text), &written, NULL);
 210             LocalFree(utf8_text);
 211         }
 212     }
 213     return result;
 214 }
 215
 216 #define WRITE_TEXT(...) do \
 217 { \
 218     wchar_t* const _message = CONCAT(__VA_ARGS__); \
 219     if (_message) \
 220     { \
 221         write_text(hStdErr, _message); \
 222         LocalFree(_message); \
 223     } \
 224 } \
 225 while (0)
 226
 227 // --------------------------------------------------------------------------
 228 // Writer thread
 229 // --------------------------------------------------------------------------
 230
 231 static BYTE buffer[2U][BUFFSIZE];
 232 static DWORD bytesTotal[2U] = { 0U, 0U };
 233 static volatile ULONG_PTR index = 0U;
 234
 235 typedef struct
 236 {
 237     HANDLE hOutput,hError;
 238     BOOL flush;
 239     HANDLE hEventReady[2U], hEventCompleted;
 240 }
 241 thread_t;
 242
 243 static DWORD WINAPI writer_thread_start_routine(const LPVOID lpThreadParameter)
 244 {
 245     DWORD bytesWritten;
 246     const thread_t *const param = (thread_t*) lpThreadParameter;
 247
 248     for (;;)
 249     {
 250         switch (WaitForMultipleObjects(2U, param->hEventReady, FALSE, INFINITE))
 251         {
 252         case WAIT_OBJECT_0:
 253             break;
 254         case WAIT_OBJECT_0 + 1U:
 255             SetEvent(param->hEventCompleted);
 256             return 0U;
 257         default:
 258             write_text(param->hError, L"[tee] System error: Failed to wait for event!\n");
 259             return 1U;
 260         }
 261
 262         const ULONG_PTR myIndex = index;
 263
 264         for (DWORD offset = 0U; offset < bytesTotal[myIndex]; offset += bytesWritten)
 265         {
 266             const BOOL result = WriteFile(param->hOutput, buffer[myIndex] + offset, bytesTotal[myIndex] - offset, &bytesWritten, NULL);
 267             if ((!result) || (!bytesWritten))
 268             {
 269                 write_text(param->hError, L"[tee] Error: Not all data could be written!\n");
 270                 break;
 271             }
 272         }
 273
 274         SetEvent(param->hEventCompleted);
 275
 276         if (param->flush)
 277         {
 278             FlushFileBuffers(param->hOutput);
 279         }
 280     }
 281 }
 282
 283 // --------------------------------------------------------------------------
 284 // Options
 285 // --------------------------------------------------------------------------
 286
 287 typedef struct
 288 {
 289     BOOL append, flush, ignore, help, version;
 290 }
 291 options_t;
 292
 293 #define PARSE_OPTION(SHRT, NAME) do \
 294 { \
 295     if ((lc == L##SHRT) || (name && (lstrcmpiW(name, L#NAME) == 0))) \
 296     { \
 297         options->NAME = TRUE; \
 298         return TRUE; \
 299     } \
 300 } \
 301 while (0)
 302
 303 static BOOL parse_option(options_t *const options, const wchar_t c, const wchar_t *const name)
 304 {
 305     const wchar_t lc = to_lower(c);
 306
 307     PARSE_OPTION('a', append);
 308     PARSE_OPTION('f', flush);
 309     PARSE_OPTION('i', ignore);
 310     PARSE_OPTION('h', help);
 311     PARSE_OPTION('v', version);
 312
 313     return FALSE;
 314 }
 315
 316 static BOOL parse_argument(options_t *const options, const wchar_t *const argument)
 317 {
 318     if ((argument[0U] != L'-') || (argument[1U] == L'\0'))
 319     {
 320         return FALSE;
 321     }
 322
 323     if (argument[1U] == L'-')
 324     {
 325         return (argument[2U] != L'\0') && parse_option(options, L'\0', argument + 2U);
 326     }
 327     else
 328     {
 329         for (const wchar_t* ptr = argument + 1U; *ptr != L'\0'; ++ptr)
 330         {
 331             if (!parse_option(options, *ptr, NULL))
 332             {
 333                 return FALSE;
 334             }
 335         }
 336         return TRUE;
 337     }
 338 }
 339
 340 // --------------------------------------------------------------------------
 341 // MAIN
 342 // --------------------------------------------------------------------------
 343
 344 int wmain(const int argc, const wchar_t *const argv[])
 345 {
 346     HANDLE hThreads[2U] = { NULL, NULL };
 347     HANDLE hEventStop = NULL, hEventThrdReady[2U] = { NULL, NULL }, hEventCompleted[2U] = { NULL, NULL };
 348     HANDLE hMyFile = INVALID_HANDLE_VALUE;
 349     int exitCode = 1, argOff = 1;
 350     options_t options;
 351     thread_t threadData[2U];
 352
 353     /* Clear options */
 354     SecureZeroMemory(&options, sizeof(options_t));
 355
 356     /* Initialize standard streams */
 357     const HANDLE hStdIn = GetStdHandle(STD_INPUT_HANDLE), hStdOut = GetStdHandle(STD_OUTPUT_HANDLE), hStdErr = GetStdHandle(STD_ERROR_HANDLE);
 358     if ((hStdIn == INVALID_HANDLE_VALUE) || (hStdOut == INVALID_HANDLE_VALUE) || (hStdErr == INVALID_HANDLE_VALUE))
 359     {
 360         return -1;
 361     }
 362
 363     /* Set up CRTL+C handler */
 364     SetConsoleCtrlHandler(console_handler, TRUE);
 365
 366     /* Parse command-line options */
 367     while ((argOff < argc) && (argv[argOff][0U] == L'-') && (argv[argOff][1U] != L'\0'))
 368     {
 369         const wchar_t *const argValue= argv[argOff++];
 370         if ((argValue[1U] == L'-') && (argValue[2U] == L'\0'))
 371         {
 372             break; /*stop!*/
 373         }
 374         else if (!parse_argument(&options, argValue))
 375         {
 376             WRITE_TEXT(L"[tee] Error: Invalid option \"", argValue, L"\" encountered!\n");
 377             return 1;
 378         }
 379     }
 380
 381     /* Print version information */
 382     if (options.version)
 383     {
 384         write_text(hStdErr, get_version_string());
 385         return 0;
 386     }
 387
 388     /* Print manual page */
 389     if (options.help)
 390     {
 391         write_text(hStdErr, get_version_string());
 392         write_text(hStdErr, L"\n"
 393             L"Copy standard input to output file, and also to standard output.\n\n"
 394             L"Usage:\n"
 395             L"  gizmo.exe [...] | tee.exe [options] <output_file>\n\n"
 396             L"Options:\n"
 397             L"  -a --append  Append to the existing file, instead of truncating\n"
 398             L"  -f --flush   Flush output file after each write operation\n"
 399             L"  -i --ignore  Ignore the interrupt signal (SIGINT), e.g. CTRL+C\n\n");
 400         return 0;
 401     }
 402
 403     /* Check output file name */
 404     if (argOff >= argc)
 405     {
 406         write_text(hStdErr, L"[tee] Error: Output file name is missing. Type \"tee --help\" for details!\n");
 407         return 1;
 408     }
 409
 410     /* Check for excess arguments */
 411     if (argOff + 1 < argc)
 412     {
 413         write_text(hStdErr, L"[tee] Warning: Excess command line argument(s) ignored!\n");
 414     }
 415
 416     /* Open output file */
 417     if (!is_null_device(argv[argOff]))
 418     {
 419         if ((hMyFile = CreateFileW(argv[argOff], GENERIC_WRITE, FILE_SHARE_READ, NULL, options.append ? OPEN_ALWAYS : CREATE_ALWAYS, 0U, NULL)) == INVALID_HANDLE_VALUE)
 420         {
 421             WRITE_TEXT(L"[tee] Error: Failed to open the output file \"", argv[argOff], L"\" for writing!\n");
 422             goto cleanup;
 423         }
 424
 425         /* Seek to the end of the file */
 426         if (options.append)
 427         {
 428             LARGE_INTEGER offset = { .QuadPart = 0LL };
 429             if (!SetFilePointerEx(hMyFile, offset, NULL, FILE_END))
 430             {
 431                 write_text(hStdErr, L"[tee] Error: Failed to move the file pointer to the end of the file!\n");
 432                 goto cleanup;
 433             }
 434         }
 435     }
 436
 437     /* Determine thread count */
 438     const DWORD threadCount = (hMyFile != INVALID_HANDLE_VALUE) ? 2U : 1U;
 439
 440     /* Create events */
 441     if (!(hEventStop = CreateEventW(NULL, TRUE, FALSE, NULL)))
 442     {
 443         write_text(hStdErr, L"[tee] System error: Failed to create event object!\n\n");
 444         goto cleanup;
 445     }
 446     for (size_t threadId = 0U; threadId < threadCount; ++threadId)
 447     {
 448         if (!(hEventThrdReady[threadId] = CreateEventW(NULL, FALSE, FALSE, NULL)))
 449         {
 450             write_text(hStdErr, L"[tee] System error: Failed to create event object!\n\n");
 451             goto cleanup;
 452         }
 453         if (!(hEventCompleted[threadId] = CreateEventW(NULL, FALSE, FALSE, NULL)))
 454         {
 455             write_text(hStdErr, L"[tee] System error: Failed to create event object!\n\n");
 456             goto cleanup;
 457         }
 458     }
 459
 460     /* Set up thread data */
 461     for (size_t threadId = 0; threadId < threadCount; ++threadId)
 462     {
 463         threadData[threadId].hOutput = (threadId > 0U) ? hMyFile : hStdOut;
 464         threadData[threadId].hError = hStdErr;
 465         threadData[threadId].flush = options.flush && (!is_terminal(threadData[threadId].hOutput));
 466         threadData[threadId].hEventReady[0U] = hEventThrdReady[threadId];
 467         threadData[threadId].hEventReady[1U] = hEventStop;
 468         threadData[threadId].hEventCompleted = hEventCompleted[threadId];
 469     }
 470
 471     /* Start threads */
 472     for (DWORD threadId = 0; threadId < threadCount; ++threadId)
 473     {
 474         if (!(hThreads[threadId] = CreateThread(NULL, 0U, writer_thread_start_routine, &threadData[threadId], 0U, NULL)))
 475         {
 476             write_text(hStdErr, L"[tee] System error: Failed to create thread!\n");
 477             goto cleanup;
 478         }
 479     }
 480
 481     /* Are we reading from a pipe? */
 482     const BOOL isPipeInput = (GetFileType(hStdIn) == FILE_TYPE_PIPE);
 483
 484     /* Initialize index */
 485     ULONG_PTR myIndex = 1U - index;
 486
 487     /* Process all input from STDIN stream */
 488     do
 489     {
 490         for (DWORD threadId = 0U; threadId < threadCount; ++threadId)
 491         {
 492             if (!SetEvent(hEventThrdReady[threadId]))
 493             {
 494                 write_text(hStdErr, L"[tee] System error: Failed to signal event!\n");
 495                 goto cleanup;
 496             }
 497         }
 498
 499         if (!ReadFile(hStdIn, buffer[myIndex], BUFFSIZE, &bytesTotal[myIndex], NULL))
 500         {
 501             if (GetLastError() != ERROR_BROKEN_PIPE)
 502             {
 503                 write_text(hStdErr, L"[tee] Error: Failed to read input data!\n");
 504                 goto cleanup;
 505             }
 506             break;
 507         }
 508
 509         if ((!bytesTotal[myIndex]) && (!isPipeInput)) /*pipes may return zero bytes, even when more data can become available later!*/
 510         {
 511             break;
 512         }
 513
 514         const DWORD waitResult = WaitForMultipleObjects(threadCount, hEventCompleted, TRUE, INFINITE);
 515         if ((waitResult != WAIT_OBJECT_0) && (waitResult != WAIT_OBJECT_0 + 1U))
 516         {
 517             write_text(hStdErr, L"[tee] System error: Failed to wait for events!\n");
 518             goto cleanup;
 519         }
 520
 521         myIndex = (ULONG_PTR) InterlockedExchangePointer((PVOID*)&index, (PVOID)myIndex);
 522     }
 523     while ((!g_stop) || options.ignore);
 524
 525     exitCode = 0;
 526
 527 cleanup:
 528
 529     /* Stop the worker threads */
 530     if (hEventStop)
 531     {
 532         SetEvent(hEventStop);
 533     }
 534
 535     /* Wait for worker threads to exit */
 536     const DWORD _threadCount = hThreads[0U] ? (hThreads[1U] ? 2U : 1U) : 0U;
 537     if (_threadCount)
 538     {
 539         const DWORD waitResult = WaitForMultipleObjects(_threadCount, hThreads, TRUE, 12000U);
 540         if ((waitResult != WAIT_OBJECT_0) && (waitResult != WAIT_OBJECT_0 + 1U))
 541         {
 542             for (DWORD threadId = 0U; threadId < _threadCount; ++threadId)
 543             {
 544                 if (WaitForSingleObject(hThreads[threadId], 1U) != WAIT_OBJECT_0)
 545                 {
 546                     write_text(hStdErr, L"[tee] Error: Worker thread did not exit cleanly!\n");
 547                     TerminateThread(hThreads[threadId], 1U);
 548                 }
 549             }
 550         }
 551     }
 552
 553     /* Flush the output file */
 554     if ((hMyFile != INVALID_HANDLE_VALUE) && options.flush)
 555     {
 556         FlushFileBuffers(hMyFile);
 557     }
 558
 559     /* Close worker threads */
 560     for (DWORD threadId = 0U; threadId < 2U; ++threadId)
 561     {
 562         CLOSE_HANDLE(hThreads[threadId]);
 563     }
 564
 565     /* Close the output file */
 566     CLOSE_HANDLE(hMyFile);
 567
 568     /* Close events */
 569     for (DWORD threadId = 0U; threadId < 2U; ++threadId)
 570     {
 571         CLOSE_HANDLE(hEventThrdReady[threadId]);
 572         CLOSE_HANDLE(hEventCompleted[threadId]);
 573     }
 574     CLOSE_HANDLE(hEventStop);
 575
 576     /* Exit */
 577     return exitCode;
 578 }
 579
 580 // --------------------------------------------------------------------------
 581 // CRT Startup
 582 // --------------------------------------------------------------------------
 583
 584 #pragma warning(disable: 4702)
 585
 586 int wmainCRTStartup(void)
 587 {
 588     SetErrorMode(SEM_FAILCRITICALERRORS);
 589
 590     int nArgs;
 591     LPWSTR *const szArglist = CommandLineToArgvW(GetCommandLineW(), &nArgs);
 592     if (!szArglist)
 593     {
 594         ExitProcess((UINT)-1);
 595     }
 596
 597     const int retval = wmain(nArgs, szArglist);
 598     LocalFree(szArglist);
 599     ExitProcess((UINT)retval);
 600
 601     return 0;
 602 }