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