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
  22 #define BUFFSIZE 4096U
  23 #pragma warning(disable: 4706)
  24
  25  // --------------------------------------------------------------------------
  26  // Utilities
  27  // --------------------------------------------------------------------------
  28
  29 static BOOL is_terminal(const HANDLE handle)
  30 {
  31     DWORD mode;
  32     return GetConsoleMode(handle, &mode);
  33 }
  34
  35 #define APPEND_STRING(X) \
  36     do { lstrcpyW(ptr, str##X); ptr += len##X; } while(0)
  37
  38 static wchar_t *concat_3(const wchar_t *const strA, const wchar_t *const strB, const wchar_t *const strC)
  39 {
  40     const size_t lenA = lstrlenW(strA), lenB = lstrlenW(strB), lenC = lstrlenW(strC);
  41     wchar_t *const buffer = (wchar_t*) LocalAlloc(LPTR, sizeof(wchar_t) * (lenA + lenB + lenC + 1U));
  42     if (buffer)
  43     {
  44         wchar_t *ptr = buffer;
  45         APPEND_STRING(A);
  46         APPEND_STRING(B);
  47         APPEND_STRING(C);
  48     }
  49     return buffer;
  50 }
  51
  52 #define SAFE_CLOSE_HANDLE(HANDLE) do \
  53 { \
  54     if (HANDLE) \
  55     { \
  56         CloseHandle((HANDLE)); \
  57     } \
  58 } \
  59 while (0)
  60
  61 // --------------------------------------------------------------------------
  62 // Console CTRL+C handler
  63 // --------------------------------------------------------------------------
  64
  65 static volatile BOOL g_stop = FALSE;
  66
  67 static BOOL WINAPI console_handler(const DWORD ctrlType)
  68 {
  69     switch (ctrlType)
  70     {
  71     case CTRL_C_EVENT:
  72     case CTRL_BREAK_EVENT:
  73     case CTRL_CLOSE_EVENT:
  74         g_stop = TRUE;
  75         return TRUE;
  76     default:
  77         return FALSE;
  78     }
  79 }
  80
  81 // --------------------------------------------------------------------------
  82 // Version
  83 // --------------------------------------------------------------------------
  84
  85 static ULONGLONG get_version(void)
  86 {
  87     const HRSRC hVersion = FindResourceW(NULL, MAKEINTRESOURCE(VS_VERSION_INFO), RT_VERSION);
  88     if (hVersion)
  89     {
  90         const HGLOBAL hResource = LoadResource(NULL, hVersion);
  91         if (hResource)
  92         {
  93             const DWORD sizeOfResource = SizeofResource(NULL, hResource);
  94             if (sizeOfResource >= sizeof(VS_FIXEDFILEINFO))
  95             {
  96                 const PVOID addrResourceBlock = LockResource(hResource);
  97                 if (addrResourceBlock)
  98                 {
  99                     VS_FIXEDFILEINFO *fileInfoData;
 100                     UINT fileInfoSize;
 101                     if (VerQueryValueW(addrResourceBlock, L"\\", &fileInfoData, &fileInfoSize))
 102                     {
 103                         ULARGE_INTEGER fileVersion;
 104                         fileVersion.LowPart  = fileInfoData->dwFileVersionLS;
 105                         fileVersion.HighPart = fileInfoData->dwFileVersionMS;
 106                         return fileVersion.QuadPart;
 107                     }
 108                 }
 109             }
 110         }
 111     }
 112
 113     return 0U;
 114 }
 115
 116 static const wchar_t *get_version_string(void)
 117 {
 118     static wchar_t text[64U] = { '\0' };
 119     lstrcpyW(text, L"tee for Windows v#.#.# [" TEXT(__DATE__) L"]\n");
 120
 121     const ULONGLONG version = get_version();
 122     if (version)
 123     {
 124         text[17U] = L'0' + ((version >> 48) & 0xFFFF);
 125         text[19U] = L'0' + ((version >> 32) & 0xFFFF);
 126         text[21U] = L'0' + ((version >> 16) & 0xFFFF);
 127     }
 128
 129     return text;
 130 }
 131
 132 // --------------------------------------------------------------------------
 133 // Text output
 134 // --------------------------------------------------------------------------
 135
 136 static char *utf16_to_utf8(const wchar_t *const input)
 137 {
 138     const int buff_size = WideCharToMultiByte(CP_UTF8, 0, input, -1, NULL, 0, NULL, NULL);
 139     if (buff_size > 0)
 140     {
 141         char *const buffer = (char*)LocalAlloc(LPTR, buff_size);
 142         if (buffer)
 143         {
 144             const int result = WideCharToMultiByte(CP_UTF8, 0, input, -1, buffer, buff_size, NULL, NULL);
 145             if ((result > 0) && (result <= buff_size))
 146             {
 147                 return buffer;
 148             }
 149             LocalFree(buffer);
 150         }
 151     }
 152     return NULL;
 153 }
 154
 155 static BOOL write_text(const HANDLE handle, const wchar_t *const text)
 156 {
 157     BOOL result = FALSE;
 158     DWORD written;
 159     if (GetConsoleMode(handle, &written))
 160     {
 161         result = WriteConsoleW(handle, text, lstrlenW(text), &written, NULL);
 162     }
 163     else
 164     {
 165         char *const utf8_text = utf16_to_utf8(text);
 166         if (utf8_text)
 167         {
 168             result = WriteFile(handle, utf8_text, lstrlenA(utf8_text), &written, NULL);
 169             LocalFree(utf8_text);
 170         }
 171     }
 172     return result;
 173 }
 174
 175 // --------------------------------------------------------------------------
 176 // Writer thread
 177 // --------------------------------------------------------------------------
 178
 179 static BYTE buffer[2U][BUFFSIZE];
 180 static DWORD bytesTotal[2U] = { 0U, 0U };
 181 static volatile ULONG_PTR index = 0U;
 182
 183 typedef struct
 184 {
 185     HANDLE hOutput,hError;
 186     BOOL flush;
 187     HANDLE hEventReady[2U], hEventCompleted;
 188 }
 189 thread_t;
 190
 191 static DWORD WINAPI writer_thread_start_routine(const LPVOID lpThreadParameter)
 192 {
 193     DWORD bytesWritten;
 194     const thread_t *const param = (thread_t*) lpThreadParameter;
 195
 196     for (;;)
 197     {
 198         switch (WaitForMultipleObjects(2U, param->hEventReady, FALSE, INFINITE))
 199         {
 200         case WAIT_OBJECT_0:
 201             break;
 202         case WAIT_OBJECT_0 + 1U:
 203             SetEvent(param->hEventCompleted);
 204             return 0U;
 205         default:
 206             write_text(param->hError, L"[tee] System error: Failed to wait for event!\n");
 207             return 1U;
 208         }
 209
 210         const ULONG_PTR myIndex = index;
 211
 212         for (DWORD offset = 0U; offset < bytesTotal[myIndex]; offset += bytesWritten)
 213         {
 214             const BOOL result = WriteFile(param->hOutput, buffer[myIndex] + offset, bytesTotal[myIndex] - offset, &bytesWritten, NULL);
 215             if ((!result) || (!bytesWritten))
 216             {
 217                 write_text(param->hError, L"[tee] Error: Not all data could be written!\n");
 218                 break;
 219             }
 220         }
 221
 222         SetEvent(param->hEventCompleted);
 223
 224         if (param->flush)
 225         {
 226             FlushFileBuffers(param->hOutput);
 227         }
 228     }
 229 }
 230
 231 // --------------------------------------------------------------------------
 232 // MAIN
 233 // --------------------------------------------------------------------------
 234
 235 int wmain(const int argc, const wchar_t *const argv[])
 236 {
 237     HANDLE hThreads[2U] = { NULL, NULL };
 238     HANDLE hEventStop = NULL, hEventThrdReady[2U] = { NULL, NULL }, hEventCompleted[2U] = { NULL, NULL };
 239     HANDLE hMyFile = INVALID_HANDLE_VALUE;
 240     int exitCode = 1, argOff = 1;
 241     BOOL append = FALSE, flush = FALSE, ignore = FALSE;
 242     thread_t threadData[2U];
 243
 244     /* Initialize standard streams */
 245     const HANDLE hStdIn = GetStdHandle(STD_INPUT_HANDLE), hStdOut = GetStdHandle(STD_OUTPUT_HANDLE), hStdErr = GetStdHandle(STD_ERROR_HANDLE);
 246     if ((hStdIn == INVALID_HANDLE_VALUE) || (hStdOut == INVALID_HANDLE_VALUE) || (hStdErr == INVALID_HANDLE_VALUE))
 247     {
 248         FatalExit(-1);
 249     }
 250
 251     /* Set up CRTL+C handler */
 252     SetConsoleCtrlHandler(console_handler, TRUE);
 253
 254     /* Print version */
 255     if ((argc > 1) && (lstrcmpiW(argv[1], L"--version") == 0))
 256     {
 257         write_text(hStdErr, get_version_string());
 258         return 0;
 259     }
 260
 261     /* Print manpage */
 262     if ((argc < 2) || (lstrcmpiW(argv[1], L"/?") == 0) || (lstrcmpiW(argv[1], L"--help") == 0))
 263     {
 264         write_text(hStdErr, get_version_string());
 265         write_text(hStdErr, L"\n"
 266             L"Usage:\n"
 267             L"  your_program.exe [...] | tee.exe [options] <output_file>\n\n"
 268             L"Options:\n"
 269             L"  --append   Append to the existing file, instead of truncating\n"
 270             L"  --flush    Flush output file after each write operation\n"
 271             L"  --ignore   Ignore the interrupt signal (SIGINT), e.g. CTRL+C\n\n");
 272         return 1;
 273     }
 274
 275     /* Parse command-line options */
 276     while (argOff < argc)
 277     {
 278         if ((argv[argOff][0U] == L'-') && (argv[argOff][1U] == L'-'))
 279         {
 280             const wchar_t *const option = argv[argOff++] + 2U;
 281             if (*option == L'\0')
 282             {
 283                 break;
 284             }
 285             else if (lstrcmpiW(option, L"append") == 0)
 286             {
 287                 append = TRUE;
 288             }
 289             else if (lstrcmpiW(option, L"flush") == 0)
 290             {
 291                 flush = TRUE;
 292             }
 293             else if (lstrcmpiW(option, L"ignore") == 0)
 294             {
 295                 ignore = TRUE;
 296             }
 297             else
 298             {
 299                 wchar_t *const message = concat_3(L"[tee] Error: Invalid option \"--", option, L"\" encountered!\n");
 300                 if (message)
 301                 {
 302                     write_text(hStdErr, message);
 303                     LocalFree(message);
 304                 }
 305                 return 1;
 306             }
 307         }
 308         else
 309         {
 310             break; /* stop option processing */
 311         }
 312     }
 313
 314     /* Check output file name */
 315     if (argOff >= argc)
 316     {
 317         write_text(hStdErr, L"[tee] Error: Output file name is missing!\n");
 318         return 1;
 319     }
 320
 321     /* Create events */
 322     if (!(hEventStop = CreateEventW(NULL, TRUE, FALSE, NULL)))
 323     {
 324         write_text(hStdErr, L"[tee] System error: Failed to create event!\n\n");
 325         goto cleanup;
 326     }
 327     for (size_t i = 0U; i < 2U; ++i)
 328     {
 329         if (!(hEventThrdReady[i] = CreateEventW(NULL, FALSE, FALSE, NULL)))
 330         {
 331             write_text(hStdErr, L"[tee] System error: Failed to create event!\n\n");
 332             goto cleanup;
 333         }
 334         if (!(hEventCompleted[i] = CreateEventW(NULL, FALSE, FALSE, NULL)))
 335         {
 336             write_text(hStdErr, L"[tee] System error: Failed to create event!\n\n");
 337             goto cleanup;
 338         }
 339     }
 340
 341     /* Open output file */
 342     if ((hMyFile = CreateFileW(argv[argOff], GENERIC_WRITE, FILE_SHARE_READ, NULL, append ? OPEN_ALWAYS : CREATE_ALWAYS, 0U, NULL)) == INVALID_HANDLE_VALUE)
 343     {
 344         write_text(hStdErr, L"[tee] Error: Failed to open the output file for writing!\n");
 345         goto cleanup;
 346     }
 347
 348     /* Seek to the end of the file */
 349     if (append)
 350     {
 351         LARGE_INTEGER offset = { .QuadPart = 0LL };
 352         if (!SetFilePointerEx(hMyFile, offset, NULL, FILE_END))
 353         {
 354             write_text(hStdErr, L"[tee] Error: Failed to move the file pointer to the end of the file!\n");
 355             goto cleanup;
 356         }
 357     }
 358
 359     /* Set up thread data */
 360     for (size_t i = 0; i < 2U; ++i)
 361     {
 362         threadData[i].hOutput = (i > 0U) ? hMyFile : hStdOut;
 363         threadData[i].hError = hStdErr;
 364         threadData[i].flush = flush && (!is_terminal(threadData[i].hOutput));
 365         threadData[i].hEventReady[0U] = hEventThrdReady[i];
 366         threadData[i].hEventReady[1U] = hEventStop;
 367         threadData[i].hEventCompleted = hEventCompleted[i];
 368     }
 369
 370     /* Start threads */
 371     for (size_t i = 0; i < 2U; ++i)
 372     {
 373         if (!(hThreads[i] = CreateThread(NULL, 0U, writer_thread_start_routine, &threadData[i], 0U, NULL)))
 374         {
 375             write_text(hStdErr, L"[tee] System error: Failed to create thread!\n");
 376             goto cleanup;
 377         }
 378     }
 379
 380     /* Are we reading from a pipe? */
 381     const BOOL isPipeInput = (GetFileType(hStdIn) == FILE_TYPE_PIPE);
 382
 383     /* Initialize index */
 384     ULONG_PTR myIndex = 1U - index;
 385
 386     /* Process all input from STDIN stream */
 387     do
 388     {
 389         for (size_t i = 0U; i < 2U; ++i)
 390         {
 391             if (!SetEvent(hEventThrdReady[i]))
 392             {
 393                 write_text(hStdErr, L"[tee] System error: Failed to signal event!\n");
 394                 goto cleanup;
 395             }
 396         }
 397
 398         if (!ReadFile(hStdIn, buffer[myIndex], BUFFSIZE, &bytesTotal[myIndex], NULL))
 399         {
 400             if (GetLastError() != ERROR_BROKEN_PIPE)
 401             {
 402                 write_text(hStdErr, L"[tee] Error: Failed to read input data!\n");
 403                 goto cleanup;
 404             }
 405             break;
 406         }
 407
 408         if ((!bytesTotal[myIndex]) && (!isPipeInput)) /*pipes may return zero bytes, even when more data can become available later!*/
 409         {
 410             break;
 411         }
 412
 413         const DWORD waitResult = WaitForMultipleObjects(2U, hEventCompleted, TRUE, INFINITE);
 414         if ((waitResult != WAIT_OBJECT_0) && (waitResult != WAIT_OBJECT_0 + 1U))
 415         {
 416             write_text(hStdErr, L"[tee] System error: Failed to wait for events!\n");
 417             goto cleanup;
 418         }
 419
 420         myIndex = (ULONG_PTR) InterlockedExchangePointer((PVOID*)&index, (PVOID)myIndex);
 421     }
 422     while ((!g_stop) || ignore);
 423
 424     exitCode = 0;
 425
 426 cleanup:
 427
 428     /* Stop the worker threads */
 429     if (hEventStop)
 430     {
 431         SetEvent(hEventStop);
 432     }
 433
 434     /* Wait for worker threads to exit */
 435     if (hThreads[0U])
 436     {
 437         const DWORD waitResult = WaitForMultipleObjects(hThreads[1U] ? 2U : 1U, hThreads, TRUE, 12000U);
 438         if ((waitResult != WAIT_OBJECT_0) && (waitResult != WAIT_OBJECT_0 + 1U))
 439         {
 440             for (DWORD i = 0U; i < (hThreads[1U] ? 2U : 1U); ++i)
 441             {
 442                 if (WaitForSingleObject(hThreads[i], 125U) != WAIT_OBJECT_0)
 443                 {
 444                     write_text(hStdErr, L"[tee] Error: Worker thread did not exit cleanly!\n");
 445                     TerminateThread(hThreads[i], 1U);
 446                 }
 447             }
 448         }
 449     }
 450
 451     /* Close worker threads */
 452     for (DWORD i = 0U; i < 2U; ++i)
 453     {
 454         SAFE_CLOSE_HANDLE(hThreads[i]);
 455     }
 456
 457     /* Close output file */
 458     if (hMyFile != INVALID_HANDLE_VALUE)
 459     {
 460         if (flush)
 461         {
 462             FlushFileBuffers(hMyFile);
 463         }
 464         CloseHandle(hMyFile);
 465     }
 466
 467     /* Close events */
 468     for (size_t i = 0U; i < 2U; ++i)
 469     {
 470         SAFE_CLOSE_HANDLE(hEventThrdReady[i]);
 471         SAFE_CLOSE_HANDLE(hEventCompleted[i]);
 472     }
 473     SAFE_CLOSE_HANDLE(hEventStop);
 474
 475     /* Exit */
 476     return exitCode;
 477 }
 478
 479 // --------------------------------------------------------------------------
 480 // CRT Startup
 481 // --------------------------------------------------------------------------
 482
 483 #pragma warning(disable: 4702)
 484
 485 int wmainCRTStartup(void)
 486 {
 487     SetErrorMode(SEM_FAILCRITICALERRORS);
 488
 489     int nArgs;
 490     LPWSTR *const szArglist = CommandLineToArgvW(GetCommandLineW(), &nArgs);
 491     if (!szArglist)
 492     {
 493         FatalExit(-1);
 494     }
 495
 496     const int retval = wmain(nArgs, szArglist);
 497     LocalFree(szArglist);
 498     ExitProcess((UINT)retval);
 499
 500     return 0;
 501 }