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